This is an automated email from the ASF dual-hosted git repository.
bbejeck pushed a commit to branch trunk
in repository https://gitbox.apache.org/repos/asf/kafka.git
The following commit(s) were added to refs/heads/trunk by this push:
new db8f66ae36c KAFKA-20490: Replace txn buffer skip-list with TreeMap
(#22648)
db8f66ae36c is described below
commit db8f66ae36c9964f096d8f7fca37e09d4c6da1f4
Author: Nick Telford <[email protected]>
AuthorDate: Wed Jun 24 23:28:27 2026 +0100
KAFKA-20490: Replace txn buffer skip-list with TreeMap (#22648)
The transaction buffer's staging map (`pendingWrites`) is written
exclusively by the owning `StreamThread`, but it must also serve
point-in-time, snapshot-isolated reads from interactive-query (IQ)
threads. It was a `ConcurrentSkipListMap` so that IQ readers could
traverse it safely while the owner mutated it concurrently. That
concurrency came at a cost: every owner `get`/`put` on the hot path paid
skip-list overhead even though the owner is the sole mutator, regressing
`StreamThread` throughput.
This change replaces the `ConcurrentSkipListMap` with a plain `TreeMap`
guarded by the buffer's existing `snapshotLock`. Owner writes take the
write lock; non-owner (IQ) reads take the read lock; owner reads and
scan-copies run lock-free, because the owner is the only thread that
ever mutates the map. Because an owner scan iterates the same `TreeMap`
it may concurrently mutate, the owner eagerly copies the bounded staging
range up front to avoid a same-thread `ConcurrentModificationException`.
The same locking discipline is applied across
`RocksDBTransactionBuffer`'s `stage`, `stageDeleteRange`, `get`, and
`scan` paths.
It also fixes a pre-existing copy-on-write bug in the `rangeTombstones`
handling. The shallow `TreeMap` copy shared its mutable `List<Bytes>`
values with the original, so range-deleting the same `from` key twice
mutated a list still being traversed by an in-flight non-owner iterator,
corrupting its view. The affected value lists are now deep-copied on
write so iterators retain a stable snapshot.
New concurrency and `ConcurrentModificationException` regression tests
in `AbstractTransactionBufferTest` and `RocksDBTransactionBufferTest`
cover these paths.
Reviewers: Bill Bejeck <[email protected]>
---
.../state/internals/AbstractTransactionBuffer.java | 60 ++++++++++---
.../state/internals/RocksDBTransactionBuffer.java | 84 +++++++++++++-----
.../internals/AbstractTransactionBufferTest.java | 99 ++++++++++++++++++++++
.../internals/RocksDBTransactionBufferTest.java | 53 ++++++++++++
4 files changed, 261 insertions(+), 35 deletions(-)
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java
index b67f0e7b191..865fe9fb07e 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java
@@ -19,23 +19,36 @@ 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.
+ * staging design: a {@link TreeMap} staging buffer for uncommitted writes
plus 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.
+ * Exactly one thread — the {@link #ownerThread} that constructs the buffer
(the StreamThread) —
+ * ever mutates {@code pendingWrites}; Interactive Query (IQ) threads are
read-only. The staging
+ * map is therefore a plain (non-thread-safe) {@link TreeMap} guarded by
{@link #snapshotLock}:
+ * <ul>
+ * <li>Owner structural mutations ({@code stage}/{@code commit}/{@code
rollback}) hold the
+ * <b>write</b> lock.</li>
+ * <li>The owner reads ({@code get}) and scan-copies <b>lock-free</b>:
because it is the sole
+ * mutator and is single-threaded, a read is never concurrent with a
structural change, and
+ * IQ threads never mutate — so the tree is quiescent for the read.</li>
+ * <li>Non-owner (IQ) reads/snapshots hold the <b>read</b> lock; owner
writes become visible to
+ * them via the write-unlock → read-lock happens-before edge.</li>
+ * <li>Owner scans eagerly copy the bounded staging range into a private
{@link TreeMap} so a
+ * subsequent owner {@code stage} cannot invalidate an open iterator
with a
+ * {@link java.util.ConcurrentModificationException}.</li>
+ * </ul>
+ * {@code isEmpty()} and {@code approximateNumUncommittedBytes()} read
owner-thread-only state and
+ * are not synchronized.
*
* @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 NavigableMap<K, Optional<byte[]>> pendingWrites = new TreeMap<>();
final ReentrantReadWriteLock snapshotLock = new ReentrantReadWriteLock();
final Thread ownerThread;
long pendingWritesBytes;
@@ -83,21 +96,40 @@ abstract class AbstractTransactionBuffer<K extends
Comparable<K>> implements Tra
@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);
+ snapshotLock.writeLock().lock();
+ try {
+ pendingWrites.put(key, Optional.ofNullable(value));
+ pendingWritesBytes += estimateKeySize(key) + (value != null ?
value.length : 0);
+ stageToBackend(key, value);
+ } finally {
+ snapshotLock.writeLock().unlock();
+ }
}
@Override
public Optional<byte[]> get(final K key) {
- return pendingWrites.get(key);
+ // Owner reads lock-free: it is the sole mutator and single-threaded,
so this read is never
+ // concurrent with a structural change. Non-owner (IQ) reads take the
read lock to exclude a
+ // concurrent owner write on the non-thread-safe TreeMap.
+ if (Thread.currentThread() == ownerThread) {
+ return pendingWrites.get(key);
+ }
+ snapshotLock.readLock().lock();
+ try {
+ return pendingWrites.get(key);
+ } finally {
+ snapshotLock.readLock().unlock();
+ }
}
@Override
public ManagedKeyValueIterator<K, byte[]> all(final boolean forward) {
if (Thread.currentThread() == ownerThread) {
+ // Eager copy of the staging layer (lock-free for the sole
mutator) so a later owner
+ // stage cannot invalidate this iterator; the base store is
iterated live.
+ final NavigableMap<K, Optional<byte[]>> stagingSnapshot = new
TreeMap<>(pendingWrites);
final ManagedKeyValueIterator<K, byte[]> baseIter =
newBaseIterator(null, null, forward, true);
- return new StagedMergeIterator<>(pendingWrites, baseIter, forward);
+ return new StagedMergeIterator<>(stagingSnapshot, baseIter,
forward);
}
return snapshotScan(null, null, forward, true);
}
@@ -108,9 +140,9 @@ abstract class AbstractTransactionBuffer<K extends
Comparable<K>> implements Tra
@Override
public ManagedKeyValueIterator<K, byte[]> range(final K from, final K to,
final boolean forward, final boolean toInclusive) {
if (Thread.currentThread() == ownerThread) {
- final NavigableMap<K, Optional<byte[]>> stagingView =
boundStaging(from, to, toInclusive);
+ final NavigableMap<K, Optional<byte[]>> stagingSnapshot = new
TreeMap<>(boundStaging(from, to, toInclusive));
final ManagedKeyValueIterator<K, byte[]> baseIter =
newBaseIterator(from, to, forward, toInclusive);
- return new StagedMergeIterator<>(stagingView, baseIter, forward);
+ return new StagedMergeIterator<>(stagingSnapshot, baseIter,
forward);
}
return snapshotScan(from, to, forward, toInclusive);
}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBuffer.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBuffer.java
index 6b33217667a..5191b7f721b 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBuffer.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBuffer.java
@@ -39,9 +39,13 @@ import java.util.TreeMap;
/**
* A {@link TransactionBuffer} implementation for RocksDB-backed stores.
* Uses a {@link WriteBatch} (without index) to accumulate writes for atomic
commit.
- * Reads are handled entirely by the shared {@code ConcurrentSkipListMap} in
+ * Reads are handled entirely by the shared staging map in
* {@link AbstractTransactionBuffer}, so a {@code WriteBatchWithIndex} is not
needed.
* <p>
+ * Follows the base class locking discipline: structural mutations of the
staging map
+ * ({@link #stage}, {@link #stageDeleteRange}) hold the {@code snapshotLock}
write lock; owner
+ * reads/scan-copies are lock-free; non-owner (IQ) reads hold the read lock.
+ * <p>
* Range deletions ({@link #stageDeleteRange}) are only supported by
RocksDB-backed stores
* and are owned entirely by this class; {@link AbstractTransactionBuffer}
carries no
* tombstone state.
@@ -82,16 +86,21 @@ class RocksDBTransactionBuffer extends
AbstractTransactionBuffer<Bytes> {
* under {@code cf}, so every staged CF is committed atomically on {@link
#commit()}.
*/
void stage(final ColumnFamilyHandle cf, final Bytes key, final byte[]
value) {
- pendingWrites.put(key, Optional.ofNullable(value));
- pendingWritesBytes += estimateKeySize(key) + (value != null ?
value.length : 0);
+ snapshotLock.writeLock().lock();
try {
- if (value != null) {
- writeBatch.put(cf, key.get(), value);
- } else {
- writeBatch.delete(cf, key.get());
+ pendingWrites.put(key, Optional.ofNullable(value));
+ pendingWritesBytes += estimateKeySize(key) + (value != null ?
value.length : 0);
+ try {
+ if (value != null) {
+ writeBatch.put(cf, key.get(), value);
+ } else {
+ writeBatch.delete(cf, key.get());
+ }
+ } catch (final RocksDBException e) {
+ throw new ProcessorStateException("Error staging write in
transaction buffer for store " + storeName, e);
}
- } catch (final RocksDBException e) {
- throw new ProcessorStateException("Error staging write in
transaction buffer for store " + storeName, e);
+ } finally {
+ snapshotLock.writeLock().unlock();
}
}
@@ -102,21 +111,49 @@ class RocksDBTransactionBuffer extends
AbstractTransactionBuffer<Bytes> {
* {@link WriteBatch} under {@code cf}.
*/
void stageDeleteRange(final ColumnFamilyHandle cf, final Bytes from, final
Bytes to) {
- pendingWrites.subMap(from, true, to, false).clear();
- final TreeMap<Bytes, List<Bytes>> copy = new
TreeMap<>(rangeTombstones);
- copy.computeIfAbsent(from, k -> new ArrayList<>()).add(to);
- rangeTombstones = copy;
- pendingWritesBytes += estimateKeySize(from) + estimateKeySize(to);
+ snapshotLock.writeLock().lock();
try {
- writeBatch.deleteRange(cf, from.get(), to.get());
- } catch (final RocksDBException e) {
- throw new ProcessorStateException("Error staging delete range in
transaction buffer for store " + storeName, e);
+ // Multi-node subMap().clear() on the (non-thread-safe) TreeMap
staging buffer must hold
+ // the write lock; doing it together with the rangeTombstones swap
also makes the two
+ // updates atomic with respect to a non-owner snapshot scan.
+ pendingWrites.subMap(from, true, to, false).clear();
+ // Copy-on-write: clone the map AND the affected key's list, so a
non-owner iterator
+ // holding the previous rangeTombstones map (and its lists) is
never mutated underneath
+ // it. A shallow map copy alone would share the List<Bytes> values
and corrupt in-flight
+ // iterators when the same `from` is deleted twice.
+ final TreeMap<Bytes, List<Bytes>> copy = new
TreeMap<>(rangeTombstones);
+ final List<Bytes> existing = copy.get(from);
+ final List<Bytes> updated = existing == null ? new ArrayList<>() :
new ArrayList<>(existing);
+ updated.add(to);
+ copy.put(from, updated);
+ rangeTombstones = copy;
+ pendingWritesBytes += estimateKeySize(from) + estimateKeySize(to);
+ try {
+ writeBatch.deleteRange(cf, from.get(), to.get());
+ } catch (final RocksDBException e) {
+ throw new ProcessorStateException("Error staging delete range
in transaction buffer for store " + storeName, e);
+ }
+ } finally {
+ snapshotLock.writeLock().unlock();
}
}
@Override
public Optional<byte[]> get(final Bytes key) {
- final Optional<byte[]> staged = pendingWrites.get(key);
+ // Owner reads lock-free (sole, single-threaded mutator); non-owner
(IQ) reads take the read
+ // lock to exclude a concurrent owner write on the TreeMap. The
volatile rangeTombstones
+ // reference is read lock-free either way.
+ final Optional<byte[]> staged;
+ if (Thread.currentThread() == ownerThread) {
+ staged = pendingWrites.get(key);
+ } else {
+ snapshotLock.readLock().lock();
+ try {
+ staged = pendingWrites.get(key);
+ } finally {
+ snapshotLock.readLock().unlock();
+ }
+ }
if (staged != null) {
return staged;
}
@@ -133,8 +170,11 @@ class RocksDBTransactionBuffer extends
AbstractTransactionBuffer<Bytes> {
ManagedKeyValueIterator<Bytes, byte[]> all(final ColumnFamilyHandle cf,
final boolean forward) {
if (Thread.currentThread() == ownerThread) {
+ // Eager copy of the staging layer (lock-free for the sole
mutator) so a later owner
+ // stage cannot invalidate this iterator; the base store is
iterated live.
+ final NavigableMap<Bytes, Optional<byte[]>> stagingSnapshot = new
TreeMap<>(pendingWrites);
final ManagedKeyValueIterator<Bytes, byte[]> baseIter =
newBaseIterator(cf, null, null, forward, true);
- return new StagedMergeIterator<>(pendingWrites, baseIter, forward);
+ return new StagedMergeIterator<>(stagingSnapshot, baseIter,
forward);
}
return snapshotScan(cf, null, null, forward, true);
}
@@ -143,9 +183,9 @@ class RocksDBTransactionBuffer extends
AbstractTransactionBuffer<Bytes> {
final Bytes from, final Bytes
to,
final boolean forward, final
boolean toInclusive) {
if (Thread.currentThread() == ownerThread) {
- final NavigableMap<Bytes, Optional<byte[]>> stagingView =
boundStaging(from, to, toInclusive);
+ final NavigableMap<Bytes, Optional<byte[]>> stagingSnapshot = new
TreeMap<>(boundStaging(from, to, toInclusive));
final ManagedKeyValueIterator<Bytes, byte[]> baseIter =
newBaseIterator(cf, from, to, forward, toInclusive);
- return new StagedMergeIterator<>(stagingView, baseIter, forward);
+ return new StagedMergeIterator<>(stagingSnapshot, baseIter,
forward);
}
return snapshotScan(cf, from, to, forward, toInclusive);
}
@@ -182,6 +222,8 @@ class RocksDBTransactionBuffer extends
AbstractTransactionBuffer<Bytes> {
final boolean forward,
final boolean toInclusive
) {
+ // A RocksDB iterator already exposes a point-in-time consistent view
of the base store
+ // as of its creation, so no extra snapshotting is required beyond the
live base iterator.
return newBaseIterator(cfHandle, from, to, forward, toInclusive);
}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractTransactionBufferTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractTransactionBufferTest.java
index 86d871f6201..571b3c84dd0 100644
---
a/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractTransactionBufferTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractTransactionBufferTest.java
@@ -31,6 +31,7 @@ import java.util.TreeMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicBoolean;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -339,6 +340,104 @@ public class AbstractTransactionBufferTest {
}
}
+ // -- owner mutate-during-iteration (no CME on the TreeMap staging buffer)
--
+
+ @Test
+ void ownerStageDuringOwnAllIterationDoesNotThrow() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(3, new byte[]{3});
+
+ final ManagedKeyValueIterator<Integer, byte[]> iter = buffer.all(true);
+ assertEquals(1, iter.next().key);
+
+ // Owner structurally mutates the staging buffer while its own
iterator is open.
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(4, new byte[]{4});
+
+ // No ConcurrentModificationException, and the iterator reflects the
snapshot at creation
+ // (the keys staged after the iterator was opened are absent).
+ assertEquals(List.of(3), drainKeys(iter));
+ }
+
+ @Test
+ void ownerStageDuringOwnRangeIterationDoesNotThrow() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(5, new byte[]{5});
+
+ final ManagedKeyValueIterator<Integer, byte[]> iter = buffer.range(1,
9, true, true);
+ assertEquals(1, iter.next().key);
+
+ buffer.stage(3, new byte[]{3});
+
+ assertEquals(List.of(5), drainKeys(iter));
+ }
+
+ // -- non-owner snapshot survives a concurrent owner commit --
+
+ @Test
+ void nonOwnerAllSurvivesConcurrentOwnerCommit() throws Exception {
+ buffer.base.put(2, new byte[]{2});
+ buffer.stage(1, new byte[]{1});
+
+ final ExecutorService exec = Executors.newSingleThreadExecutor();
+ try {
+ final ManagedKeyValueIterator<Integer, byte[]> iter =
+ exec.submit(() -> buffer.all(true)).get();
+
+ // Owner commits (flushes staging to base and clears
pendingWrites) while the non-owner
+ // snapshot iterator is open.
+ buffer.commit();
+
+ // Snapshot captured at creation time: staged key 1 plus base key
2.
+ assertEquals(List.of(1, 2), drainKeys(iter));
+ } finally {
+ exec.shutdown();
+ }
+ }
+
+ // -- concurrency stress: owner writer + non-owner readers on the TreeMap
staging buffer --
+
+ @Test
+ void concurrentOwnerWritesAndNonOwnerReadsAreConsistent() throws Exception
{
+ final int readerCount = 4;
+ final ExecutorService exec = Executors.newFixedThreadPool(readerCount);
+ final AtomicBoolean stop = new AtomicBoolean(false);
+ final List<Future<?>> readers = new ArrayList<>();
+ try {
+ for (int r = 0; r < readerCount; r++) {
+ readers.add(exec.submit(() -> {
+ while (!stop.get()) {
+ buffer.get(7);
+ // Non-owner scan → snapshotScan: must never throw and
must be sorted.
+ final List<Integer> keys = drainKeys(buffer.all(true));
+ for (int i = 1; i < keys.size(); i++) {
+ if (keys.get(i - 1) >= keys.get(i)) {
+ throw new AssertionError("non-owner scan
returned unsorted keys: " + keys);
+ }
+ }
+ }
+ return null;
+ }));
+ }
+
+ for (int i = 0; i < 10_000; i++) {
+ buffer.stage(i % 64, new byte[]{(byte) i});
+ if (i % 200 == 0) {
+ buffer.commit();
+ }
+ if (i % 500 == 0) {
+ buffer.rollback();
+ }
+ }
+ stop.set(true);
+ for (final Future<?> reader : readers) {
+ reader.get(); // surfaces any AssertionError /
ConcurrentModificationException
+ }
+ } finally {
+ exec.shutdownNow();
+ }
+ }
+
// -- Helpers --
private static List<Integer> drainKeys(final
ManagedKeyValueIterator<Integer, byte[]> iter) {
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBufferTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBufferTest.java
index 0fab3b44e91..b0f7f11e37b 100644
---
a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBufferTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBufferTest.java
@@ -39,6 +39,7 @@ import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicReference;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
@@ -260,6 +261,58 @@ public class RocksDBTransactionBufferTest {
assertEquals(List.of("a", "b", "c", "d"), result.get());
}
+ @Test
+ public void shouldHandleConcurrentStageDeleteRangeAndNonOwnerReads()
throws Exception {
+ final int readerCount = 4;
+ final AtomicBoolean stop = new AtomicBoolean(false);
+ final AtomicReference<Throwable> failure = new AtomicReference<>();
+ final List<Thread> readers = new ArrayList<>();
+ for (int r = 0; r < readerCount; r++) {
+ final Thread reader = new Thread(() -> {
+ try {
+ while (!stop.get()) {
+ buffer.get(key("k50"));
+ // Non-owner scan → snapshotScan: must never throw and
must be sorted, even
+ // while the owner concurrently runs stageDeleteRange
(subMap().clear()).
+ try (KeyValueIterator<Bytes, byte[]> iter =
buffer.all(true)) {
+ Bytes prev = null;
+ while (iter.hasNext()) {
+ final Bytes k = iter.next().key;
+ if (prev != null && prev.compareTo(k) >= 0) {
+ throw new AssertionError("non-owner scan
returned unsorted keys: " + prev + " >= " + k);
+ }
+ prev = k;
+ }
+ }
+ }
+ } catch (final Throwable t) {
+ failure.compareAndSet(null, t);
+ }
+ });
+ readers.add(reader);
+ reader.start();
+ }
+ try {
+ for (int i = 0; i < 5_000; i++) {
+ buffer.stage(key(String.format("k%02d", i % 80)), val("v" +
i));
+ if (i % 50 == 0) {
+ buffer.stageDeleteRange(cfHandle, key("k10"), key("k30"));
+ }
+ if (i % 200 == 0) {
+ buffer.commit();
+ }
+ }
+ } finally {
+ stop.set(true);
+ for (final Thread reader : readers) {
+ reader.join();
+ }
+ }
+ if (failure.get() != null) {
+ throw new AssertionError("concurrent non-owner reader failed",
failure.get());
+ }
+ }
+
@Test
public void shouldNotShowStagedWritesInBaseAfterRollback() throws
RocksDBException {
buffer.stage(key("x"), val("staged-x"));
