bbejeck commented on code in PR #22323: URL: https://github.com/apache/kafka/pull/22323#discussion_r3422582356
########## streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java: ########## @@ -0,0 +1,172 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.kafka.streams.state.internals; + +import java.util.NavigableMap; +import java.util.Optional; +import java.util.TreeMap; +import java.util.concurrent.ConcurrentSkipListMap; +import java.util.concurrent.locks.ReentrantReadWriteLock; + +/** + * Base class for {@link TransactionBuffer} implementations. Provides the shared two-layer + * staging design: a thread-safe {@link ConcurrentSkipListMap} for reads (any thread) and + * backend-specific write accumulation for atomic commit. + * <p> + * Point lookups ({@link #get(Comparable)}) are lock-free. Scan methods automatically detect the + * owner thread and use a lock-free fast path; non-owner threads acquire a read lock to + * snapshot the staging map atomically with base iterator creation. + * + * @param <K> the key type, must be {@link Comparable} + */ +abstract class AbstractTransactionBuffer<K extends Comparable<K>> implements TransactionBuffer<K> { + + final ConcurrentSkipListMap<K, Optional<byte[]>> pendingWrites = new ConcurrentSkipListMap<>(); + final ReentrantReadWriteLock snapshotLock = new ReentrantReadWriteLock(); + final Thread ownerThread; + long pendingWritesBytes; + + AbstractTransactionBuffer() { + this.ownerThread = Thread.currentThread(); + } + + // -- Abstract methods to be implemented by backend-specific subclasses -- + + /** Append the write to the backend-specific batch (e.g. WriteBatch for RocksDB). */ + abstract void stageToBackend(K key, byte[] value); + + /** Create a base store iterator for the given range. Upper bound is inclusive. Forward direction. */ + abstract ManagedKeyValueIterator<K, byte[]> newBaseIterator(K from, K to); + + /** Create a base store iterator with configurable direction and upper bound inclusiveness. */ + ManagedKeyValueIterator<K, byte[]> newBaseIterator(final K from, final K to, + final boolean forward, final boolean toInclusive) { + return newBaseIterator(from, to); + } + + /** Atomically apply the accumulated writes to the base store. */ + abstract void flushToBase(); + + /** Discard the backend-specific pending batch without applying it. */ + abstract void discardPendingBatch(); + + /** Estimate the byte size of a key, for uncommitted bytes tracking. */ + abstract int estimateKeySize(K key); + + // -- TransactionBuffer implementation -- + + @Override + public void stage(final K key, final byte[] value) { + pendingWrites.put(key, Optional.ofNullable(value)); + pendingWritesBytes += estimateKeySize(key) + (value != null ? value.length : 0); + stageToBackend(key, value); + } + + @Override + public Optional<byte[]> get(final K key) { + return pendingWrites.get(key); + } + + @Override + public ManagedKeyValueIterator<K, byte[]> all(final boolean forward) { + if (Thread.currentThread() == ownerThread) { + final ManagedKeyValueIterator<K, byte[]> baseIter = newBaseIterator(null, null, forward, true); + return new StagedMergeIterator<>(pendingWrites, baseIter, forward); + } + return snapshotScan(null, null, forward, true); + } + + @Override + public ManagedKeyValueIterator<K, byte[]> range(final K from, final K to, final boolean forward, final boolean toInclusive) { + if (Thread.currentThread() == ownerThread) { + final ManagedKeyValueIterator<K, byte[]> baseIter = newBaseIterator(from, to, forward, toInclusive); + final NavigableMap<K, Optional<byte[]>> stagingView = boundStaging(from, to, toInclusive); Review Comment: nit: If `forward == true` and `from` > `to` will throw an `IllegalArguemntException` and a new RocksDB iterator created above at line 96 could turn into a leak. Let's swap the order of creation here and add a Javadoc comment about forward means from < to. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
