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 18cba70a660 KAFKA-20490: Add transaction buffer framework (#22323)
18cba70a660 is described below
commit 18cba70a66051c9f660f72aa1efee8305bc20f41
Author: Nick Telford <[email protected]>
AuthorDate: Tue Jun 16 20:59:15 2026 +0100
KAFKA-20490: Add transaction buffer framework (#22323)
Introduces the core building blocks for transactional state stores:
TransactionBuffer, AbstractTransactionBuffer, and StagedMergeIterator.
TransactionBuffer defines the staging interface for buffering writes
before a transaction commits. AbstractTransactionBuffer provides a
ConcurrentSkipListMap-backed implementation. Reads on the owner thread
(the stream thread driving the task) take a lock-free fast path directly
against the staged map; reads on non-owner threads, such as interactive
query, snapshot the staged entries at iteration start for a consistent
view without blocking writers.
StagedMergeIterator is a ManagedKeyValueIterator that merges the staged
overlay with a base store iterator, respecting point-delete tombstones
for entries deleted within the transaction. It registers with the
store's open-iterator tracking so RocksDB can defer compactions while it
is live.
The abstract base is tombstone-agnostic: point deletes are handled
uniformly by all implementations. All types are generic over the key
type so that window and session stores can reuse the same framework with
composite keys.
Reviewers: Bill Bejeck <[email protected]>
---
gradle/spotbugs-exclude.xml | 15 +
.../state/internals/AbstractTransactionBuffer.java | 175 ++++++++
.../state/internals/StagedMergeIterator.java | 172 ++++++++
.../streams/state/internals/TransactionBuffer.java | 94 +++++
.../internals/AbstractTransactionBufferTest.java | 444 +++++++++++++++++++++
.../state/internals/StagedMergeIteratorTest.java | 362 +++++++++++++++++
6 files changed, 1262 insertions(+)
diff --git a/gradle/spotbugs-exclude.xml b/gradle/spotbugs-exclude.xml
index f02f9e04bb3..6d18691973a 100644
--- a/gradle/spotbugs-exclude.xml
+++ b/gradle/spotbugs-exclude.xml
@@ -225,6 +225,21 @@ For a detailed description of spotbugs bug categories, see
https://spotbugs.read
<Bug pattern="NP_OPTIONAL_RETURN_NULL"/>
</Match>
+ <Match>
+ <!-- null means "not staged, fall back to base store";
Optional.empty() means staged tombstone -->
+ <Class
name="org.apache.kafka.streams.state.internals.AbstractTransactionBuffer"/>
+ <Method name="get"/>
+ <Bug pattern="NP_OPTIONAL_RETURN_NULL"/>
+ </Match>
+
+ <Match>
+ <!-- pendingWritesBytes is only accessed by the owner thread
(stage/commit/rollback all enforce this);
+ non-volatile and non-atomic access is intentional to avoid
synchronization overhead. -->
+ <Class
name="org.apache.kafka.streams.state.internals.AbstractTransactionBuffer"/>
+ <Field name="pendingWritesBytes"/>
+ <Bug
pattern="AT_NONATOMIC_64BIT_PRIMITIVE,AT_NONATOMIC_OPERATIONS_ON_SHARED_VARIABLE,AT_STALE_THREAD_WRITE_OF_PRIMITIVE"/>
+ </Match>
+
<Match>
<!-- Suppress warnings about ignoring the return value of await.
This is done intentionally because we use other clues to determine
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
new file mode 100644
index 00000000000..21e71745c33
--- /dev/null
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractTransactionBuffer.java
@@ -0,0 +1,175 @@
+/*
+ * 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);
+ }
+
+ /**
+ * @throws IllegalArgumentException if {@code from > to} and {@code
forward == true}; or {@code from < to} and {@code forward == false}.
+ */
+ @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 ManagedKeyValueIterator<K, byte[]> baseIter =
newBaseIterator(from, to, forward, toInclusive);
+ return new StagedMergeIterator<>(stagingView, baseIter, forward);
+ }
+ return snapshotScan(from, to, forward, toInclusive);
+ }
+
+ @Override
+ public void commit() {
+ snapshotLock.writeLock().lock();
+ try {
+ flushToBase();
+ pendingWrites.clear();
+ pendingWritesBytes = 0;
+ } finally {
+ snapshotLock.writeLock().unlock();
+ }
+ }
+
+ @Override
+ public void rollback() {
+ pendingWrites.clear();
+ pendingWritesBytes = 0;
+ discardPendingBatch();
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return pendingWrites.isEmpty();
+ }
+
+ @Override
+ public long approximateNumUncommittedBytes() {
+ return pendingWritesBytes;
+ }
+
+ @Override
+ public void close() {
+ // Subclasses can override to release resources
+ }
+
+ // -- Internal helpers (package-private for subclass use) --
+
+ /**
+ * Constructs an iteratator over a *snapshot* of the current transaction
buffer state.
+ * This ensures snapshot-isolation for interactive queries, and prevents
iterators suddenly becoming invalid when
+ * the stream thread clears the transaction buffer state on
commit/rollback.
+ */
+ ManagedKeyValueIterator<K, byte[]> snapshotScan(final K from, final K to,
+ final boolean forward,
final boolean toInclusive) {
+ snapshotLock.readLock().lock();
+ try {
+ 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<>(stagingSnapshot, baseIter,
forward);
+ } finally {
+ snapshotLock.readLock().unlock();
+ }
+ }
+
+ /**
+ * Constructs a view over the transaction buffer's read buffer consisting
of only the keys in the given range.
+ * If {@code from} or {@code to} are {@code null}, the range will be
open-ended; if they are both {@code null}, the
+ * range will include all keys (it will just return the original read
buffer directly).
+ */
+ NavigableMap<K, Optional<byte[]>> boundStaging(final K from, final K to,
final boolean toInclusive) {
+ if (from != null && to != null) {
+ return pendingWrites.subMap(from, true, to, toInclusive);
+ } else if (from != null) {
+ return pendingWrites.tailMap(from, true);
+ } else if (to != null) {
+ return pendingWrites.headMap(to, toInclusive);
+ } else {
+ return pendingWrites;
+ }
+ }
+}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.java
new file mode 100644
index 00000000000..76bd3abc5b2
--- /dev/null
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/StagedMergeIterator.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 org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.state.KeyValueIterator;
+
+import java.util.Iterator;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+
+/**
+ * A merge iterator that overlays staged writes (from a {@link NavigableMap})
on top of
+ * a base store iterator. Staging entries take precedence over base entries
for the same key.
+ * Tombstones ({@link Optional#empty()}) in the staging map cause the
corresponding key
+ * to be skipped in the merged output.
+ * <p>
+ * Supports both forward and reverse iteration. When iterating in reverse, the
staging
+ * map is traversed in descending order and the merge comparison is inverted
so that
+ * the largest key is emitted first.
+ *
+ * @param <K> the key type, must be {@link Comparable}
+ * @param <V> the value type
+ */
+class StagedMergeIterator<K extends Comparable<K>, V> implements
ManagedKeyValueIterator<K, V> {
+
+ private final Iterator<Map.Entry<K, Optional<V>>> stagingIterator;
+ private final KeyValueIterator<K, V> baseIterator;
+ private final boolean forward;
+
+ private Map.Entry<K, Optional<V>> nextStaging;
+ private KeyValue<K, V> nextBase;
+ private KeyValue<K, V> prefetched;
+ private boolean closed = false;
+ private Runnable closeCallback;
+
+ StagedMergeIterator(final NavigableMap<K, Optional<V>> staging,
+ final KeyValueIterator<K, V> baseIterator) {
+ this(staging, baseIterator, true);
+ }
+
+ StagedMergeIterator(final NavigableMap<K, Optional<V>> staging,
+ final KeyValueIterator<K, V> baseIterator,
+ final boolean forward) {
+ this.forward = forward;
+ final NavigableMap<K, Optional<V>> orderedStaging = forward ? staging
: staging.descendingMap();
+ this.stagingIterator = orderedStaging.entrySet().iterator();
+ this.baseIterator = baseIterator;
+ advanceStaging();
+ advanceBase();
+ }
+
+ private void advanceStaging() {
+ nextStaging = stagingIterator.hasNext() ? stagingIterator.next() :
null;
+ }
+
+ private void advanceBase() {
+ nextBase = baseIterator.hasNext() ? baseIterator.next() : null;
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (closed) {
+ throw new IllegalStateException("Iterator has already been
closed.");
+ }
+ if (prefetched != null) {
+ return true;
+ }
+ prefetched = computeNext();
+ return prefetched != null;
+ }
+
+ @Override
+ public KeyValue<K, V> next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ final KeyValue<K, V> result = prefetched;
+ prefetched = null;
+ return result;
+ }
+
+ @Override
+ public K peekNextKey() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return prefetched.key;
+ }
+
+ @Override
+ public void onClose(final Runnable closeCallback) {
+ this.closeCallback = closeCallback;
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ try {
+ baseIterator.close();
+ } finally {
+ if (closeCallback != null) {
+ closeCallback.run();
+ }
+ }
+ }
+
+ private KeyValue<K, V> computeNext() {
+ while (nextStaging != null || nextBase != null) {
+ if (nextStaging == null) {
+ // staging exhausted — emit base
+ final KeyValue<K, V> result = nextBase;
+ advanceBase();
+ return result;
+ }
+ if (nextBase == null) {
+ // base exhausted — skip tombstones, emit non-tombstone
staging entries
+ final Map.Entry<K, Optional<V>> entry = nextStaging;
+ advanceStaging();
+ if (entry.getValue().isPresent()) {
+ return new KeyValue<>(entry.getKey(),
entry.getValue().get());
+ }
+ // tombstone — continue loop
+ continue;
+ }
+
+ // Compare keys; in reverse mode, invert so the largest key is
"first"
+ final int rawCmp = nextStaging.getKey().compareTo(nextBase.key);
+ final int cmp = forward ? rawCmp : Integer.compare(0, rawCmp);
+ if (cmp < 0) {
+ // staging key comes first — emit it or skip if tombstone
+ final Map.Entry<K, Optional<V>> entry = nextStaging;
+ advanceStaging();
+ if (entry.getValue().isPresent()) {
+ return new KeyValue<>(entry.getKey(),
entry.getValue().get());
+ }
+ // tombstone — continue loop
+ } else if (cmp > 0) {
+ // base key comes first
+ final KeyValue<K, V> result = nextBase;
+ advanceBase();
+ return result;
+ } else {
+ // same key — staging takes precedence, skip base
+ advanceBase();
+ final Map.Entry<K, Optional<V>> entry = nextStaging;
+ advanceStaging();
+ if (entry.getValue().isPresent()) {
+ return new KeyValue<>(entry.getKey(),
entry.getValue().get());
+ }
+ // tombstone — both skipped, continue loop
+ }
+ }
+ return null;
+ }
+}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/TransactionBuffer.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/TransactionBuffer.java
new file mode 100644
index 00000000000..09ec173d21c
--- /dev/null
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/TransactionBuffer.java
@@ -0,0 +1,94 @@
+/*
+ * 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.io.Closeable;
+import java.util.Optional;
+
+/**
+ * A transaction buffer that accumulates writes and flushes them atomically.
+ * <p>
+ * Staged writes are visible to all readers via {@link #get(Comparable)} and
scan methods,
+ * but are only applied to the underlying store on {@link #commit()}.
+ *
+ * @param <K> the key type, must be {@link Comparable}
+ */
+interface TransactionBuffer<K extends Comparable<K>> extends Closeable {
+
+ /**
+ * Stage a put or delete. A null value represents a tombstone (pending
delete).
+ * Must be called from the owner thread only.
+ */
+ void stage(K key, byte[] value);
+
+ /**
+ * Look up a key in the staging buffer.
+ * <p>
+ * Returns {@code null} if the key has no staged entry (caller should fall
back to the
+ * base store). Returns {@link Optional#empty()} if the key has a staged
tombstone
+ * (pending delete). Returns {@code Optional.of(value)} if the key has a
staged value.
+ * <p>
+ * Can be called from any thread without locking.
+ */
+ Optional<byte[]> get(K key);
+
+ /**
+ * Return a scan iterator over all entries in the given direction (staged
overlay
+ * merged with base store). Safe to call from any thread. Owner-thread
calls use a
+ * lock-free fast path; other threads acquire a read lock to snapshot
staged writes.
+ */
+ ManagedKeyValueIterator<K, byte[]> all(boolean forward);
+
+ /**
+ * Return a range scan iterator with configurable direction and upper
bound inclusiveness
+ * (staged overlay merged with base store). Safe to call from any thread.
Owner-thread
+ * calls use a lock-free fast path; other threads acquire a read lock to
snapshot staged
+ * writes.
+ */
+ ManagedKeyValueIterator<K, byte[]> range(K from, K to, boolean forward,
boolean toInclusive);
+
+ /**
+ * Atomically apply all staged writes to the underlying store and clear
the staging area.
+ * Must be called from the owner thread only.
+ */
+ void commit();
+
+ /**
+ * Discard all staged writes without applying them.
+ * Must be called from the owner thread only.
+ */
+ void rollback();
+
+ /**
+ * Returns true if there are no staged writes.
+ */
+ boolean isEmpty();
+
+ /**
+ * Returns an approximation of the number of uncommitted bytes currently
staged in this buffer.
+ * This is not exact — it does not account for overwrites of the same key,
and may not include
+ * all overhead (e.g. object headers in the staging map). It is intended
for use in triggering
+ * early commits when the buffer grows too large.
+ */
+ long approximateNumUncommittedBytes();
+
+ /**
+ * Release any resources held by this buffer.
+ */
+ @Override
+ void close();
+}
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
new file mode 100644
index 00000000000..9527b4b24af
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/AbstractTransactionBufferTest.java
@@ -0,0 +1,444 @@
+/*
+ * 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 org.apache.kafka.streams.KeyValue;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+
+import java.util.ArrayList;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.NavigableMap;
+import java.util.Optional;
+import java.util.TreeMap;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class AbstractTransactionBufferTest {
+
+ private TestBuffer buffer;
+
+ @BeforeEach
+ void setUp() {
+ buffer = new TestBuffer();
+ }
+
+ // -- stage / get --
+
+ @Test
+ void getStagedValueReturnsOptionalOfValue() {
+ final byte[] val = {1, 2, 3};
+ buffer.stage(1, val);
+ assertArrayEquals(val, buffer.get(1).get());
+ }
+
+ @Test
+ void getStagedTombstoneReturnsEmptyOptional() {
+ buffer.stage(1, null);
+ assertEquals(Optional.empty(), buffer.get(1));
+ }
+
+ @Test
+ void getUnstaggedKeyReturnsNull() {
+ assertNull(buffer.get(99));
+ }
+
+ // -- isEmpty / approximateNumUncommittedBytes --
+
+ @Test
+ void isEmptyAndZeroBytesOnFreshBuffer() {
+ assertTrue(buffer.isEmpty());
+ assertEquals(0L, buffer.approximateNumUncommittedBytes());
+ }
+
+ @Test
+ void isNotEmptyAfterStage() {
+ buffer.stage(1, new byte[]{1});
+ assertFalse(buffer.isEmpty());
+ }
+
+ @Test
+ void byteCountReflectsStagedValue() {
+ final byte[] val = {1, 2, 3};
+ buffer.stage(1, val);
+ assertEquals(Integer.BYTES + val.length,
buffer.approximateNumUncommittedBytes());
+ }
+
+ @Test
+ void byteCountReflectsStagedTombstone() {
+ buffer.stage(1, null);
+ assertEquals(Integer.BYTES, buffer.approximateNumUncommittedBytes());
+ }
+
+ @Test
+ void byteCountAccumulatesAcrossMultipleStages() {
+ buffer.stage(1, new byte[4]);
+ buffer.stage(2, new byte[8]);
+ assertEquals(2 * Integer.BYTES + 4 + 8,
buffer.approximateNumUncommittedBytes());
+ }
+
+ // -- commit --
+
+ @Test
+ void commitClearsPendingWritesAndBytes() {
+ buffer.stage(1, new byte[]{42});
+ buffer.commit();
+ assertTrue(buffer.isEmpty());
+ assertEquals(0L, buffer.approximateNumUncommittedBytes());
+ }
+
+ @Test
+ void commitFlushesWritesToBase() {
+ final byte[] val = {7, 8, 9};
+ buffer.stage(5, val);
+ buffer.commit();
+ assertArrayEquals(val, buffer.base.get(5));
+ }
+
+ @Test
+ void commitFlushesTombstoneToBase() {
+ buffer.base.put(3, new byte[]{1});
+ buffer.stage(3, null);
+ buffer.commit();
+ assertFalse(buffer.base.containsKey(3));
+ }
+
+ @Test
+ void subsequentBatchAfterCommitIsIndependent() {
+ buffer.stage(1, new byte[]{1});
+ buffer.commit();
+ buffer.stage(2, new byte[]{2});
+ assertEquals(Integer.BYTES + 1,
buffer.approximateNumUncommittedBytes());
+ buffer.commit();
+ assertTrue(buffer.base.containsKey(1));
+ assertTrue(buffer.base.containsKey(2));
+ }
+
+ // -- rollback --
+
+ @Test
+ void rollbackClearsPendingWritesAndBytes() {
+ buffer.stage(1, new byte[]{42});
+ buffer.rollback();
+ assertTrue(buffer.isEmpty());
+ assertEquals(0L, buffer.approximateNumUncommittedBytes());
+ }
+
+ @Test
+ void rollbackLeavesBaseUnchanged() {
+ buffer.base.put(3, new byte[]{99});
+ buffer.stage(3, new byte[]{1});
+ buffer.rollback();
+ assertArrayEquals(new byte[]{99}, buffer.base.get(3));
+ }
+
+ @Test
+ void rollbackDiscardsNewEntries() {
+ buffer.stage(5, new byte[]{1});
+ buffer.rollback();
+ assertFalse(buffer.base.containsKey(5));
+ }
+
+ // -- all (owner thread) --
+
+ @Test
+ void allForwardMergesStagedAndBase() {
+ buffer.base.put(2, new byte[]{2});
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1, 2, 3), drainKeys(buffer.all(true)));
+ }
+
+ @Test
+ void allReverseReturnsKeysDescending() {
+ buffer.base.put(2, new byte[]{2});
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(3, 2, 1), drainKeys(buffer.all(false)));
+ }
+
+ @Test
+ void allStagedTombstoneHidesBaseEntry() {
+ buffer.base.put(2, new byte[]{2});
+ buffer.stage(2, null);
+ assertFalse(drainKeys(buffer.all(true)).contains(2));
+ }
+
+ // -- range (owner thread) --
+
+ @Test
+ void rangeInclusiveUpperBound() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1, 2), drainKeys(buffer.range(1, 2, true, true)));
+ }
+
+ @Test
+ void rangeExclusiveUpperBound() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1), drainKeys(buffer.range(1, 2, true, false)));
+ }
+
+ @Test
+ void rangeNullFromOpensLowerBound() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1, 2), drainKeys(buffer.range(null, 2, true,
true)));
+ }
+
+ @Test
+ void rangeNullToOpensUpperBound() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(2, 3), drainKeys(buffer.range(2, null, true,
true)));
+ }
+
+ @Test
+ void rangeNullBothBoundsWithReverseReturnsSameAsAllReverse() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(3, 2, 1), drainKeys(buffer.range(null, null,
false, true)));
+ }
+
+ // -- boundStaging --
+
+ @Test
+ void boundStagingBothNullReturnsFullMap() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ assertEquals(buffer.pendingWrites, buffer.boundStaging(null, null,
true));
+ }
+
+ @Test
+ void boundStagingOnlyFromReturnsTailMap() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(2, 3), new ArrayList<>(buffer.boundStaging(2,
null, true).keySet()));
+ }
+
+ @Test
+ void boundStagingOnlyToInclusiveReturnsHeadMapInclusive() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1, 2), new ArrayList<>(buffer.boundStaging(null,
2, true).keySet()));
+ }
+
+ @Test
+ void boundStagingOnlyToExclusiveReturnsHeadMapExclusive() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ assertEquals(List.of(1), new ArrayList<>(buffer.boundStaging(null, 2,
false).keySet()));
+ }
+
+ @Test
+ void boundStagingBothBoundsInclusiveReturnsSubMap() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ buffer.stage(4, new byte[]{4});
+ assertEquals(List.of(2, 3), new ArrayList<>(buffer.boundStaging(2, 3,
true).keySet()));
+ }
+
+ @Test
+ void boundStagingBothBoundsExclusiveUpperReturnsSubMap() {
+ buffer.stage(1, new byte[]{1});
+ buffer.stage(2, new byte[]{2});
+ buffer.stage(3, new byte[]{3});
+ buffer.stage(4, new byte[]{4});
+ assertEquals(List.of(2), new ArrayList<>(buffer.boundStaging(2, 3,
false).keySet()));
+ }
+
+ // -- non-owner thread: point lookup --
+
+ @Test
+ void nonOwnerThreadCanReadStagedEntry() throws Exception {
+ buffer.stage(1, new byte[]{42});
+ final ExecutorService exec = Executors.newSingleThreadExecutor();
+ try {
+ final Future<Optional<byte[]>> future = exec.submit(() ->
buffer.get(1));
+ final Optional<byte[]> result = future.get();
+ assertNotNull(result);
+ assertArrayEquals(new byte[]{42}, result.get());
+ } finally {
+ exec.shutdown();
+ }
+ }
+
+ // -- non-owner thread: snapshot isolation --
+
+ @Test
+ void nonOwnerAllSnapshotsAtCreationTime() throws Exception {
+ buffer.stage(1, new byte[]{1});
+
+ final ExecutorService exec = Executors.newSingleThreadExecutor();
+ try {
+ // Non-owner captures the iterator — snapshot of staging is taken
here
+ final Future<ManagedKeyValueIterator<Integer, byte[]>> futureIter =
+ exec.submit(() -> buffer.all(true));
+ final ManagedKeyValueIterator<Integer, byte[]> iter =
futureIter.get();
+
+ // Owner stages a new entry after the snapshot was taken
+ buffer.stage(2, new byte[]{2});
+
+ // Iterator must not reflect the post-snapshot stage
+ assertEquals(List.of(1), drainKeys(iter));
+ } finally {
+ exec.shutdown();
+ }
+ }
+
+ @Test
+ void nonOwnerRangeSnapshotsAtCreationTime() throws Exception {
+ buffer.stage(1, new byte[]{1});
+
+ final ExecutorService exec = Executors.newSingleThreadExecutor();
+ try {
+ final Future<ManagedKeyValueIterator<Integer, byte[]>> futureIter =
+ exec.submit(() -> buffer.range(null, null, true, true));
+ final ManagedKeyValueIterator<Integer, byte[]> iter =
futureIter.get();
+
+ buffer.stage(2, new byte[]{2});
+
+ assertEquals(List.of(1), drainKeys(iter));
+ } finally {
+ exec.shutdown();
+ }
+ }
+
+ // -- Helpers --
+
+ private static List<Integer> drainKeys(final
ManagedKeyValueIterator<Integer, byte[]> iter) {
+ final List<Integer> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key);
+ }
+ return keys;
+ }
+
+ // -- Test double --
+
+ static class TestBuffer extends AbstractTransactionBuffer<Integer> {
+ final TreeMap<Integer, byte[]> base = new TreeMap<>();
+ private final TreeMap<Integer, Optional<byte[]>> pendingBatch = new
TreeMap<>();
+
+ @Override
+ void stageToBackend(final Integer key, final byte[] value) {
+ pendingBatch.put(key, Optional.ofNullable(value));
+ }
+
+ @Override
+ ManagedKeyValueIterator<Integer, byte[]> newBaseIterator(final Integer
from, final Integer to) {
+ return newBaseIterator(from, to, true, true);
+ }
+
+ @Override
+ ManagedKeyValueIterator<Integer, byte[]> newBaseIterator(
+ final Integer from, final Integer to, final boolean forward, final
boolean toInclusive) {
+ final NavigableMap<Integer, byte[]> view = boundedBaseView(from,
to, toInclusive);
+ return new TestIterator(forward ? view : view.descendingMap());
+ }
+
+ private NavigableMap<Integer, byte[]> boundedBaseView(
+ final Integer from, final Integer to, final boolean toInclusive) {
+ if (from != null && to != null) {
+ return base.subMap(from, true, to, toInclusive);
+ } else if (from != null) {
+ return base.tailMap(from, true);
+ } else if (to != null) {
+ return base.headMap(to, toInclusive);
+ }
+ return base;
+ }
+
+ @Override
+ void flushToBase() {
+ for (final Map.Entry<Integer, Optional<byte[]>> entry :
pendingBatch.entrySet()) {
+ if (entry.getValue().isPresent()) {
+ base.put(entry.getKey(), entry.getValue().get());
+ } else {
+ base.remove(entry.getKey());
+ }
+ }
+ pendingBatch.clear();
+ }
+
+ @Override
+ void discardPendingBatch() {
+ pendingBatch.clear();
+ }
+
+ @Override
+ int estimateKeySize(final Integer key) {
+ return Integer.BYTES;
+ }
+ }
+
+ private static class TestIterator implements
ManagedKeyValueIterator<Integer, byte[]> {
+ private final Iterator<Map.Entry<Integer, byte[]>> iter;
+ private Runnable closeCallback;
+
+ TestIterator(final NavigableMap<Integer, byte[]> view) {
+ this.iter = view.entrySet().iterator();
+ }
+
+ @Override
+ public void onClose(final Runnable closeCallback) {
+ this.closeCallback = closeCallback;
+ }
+
+ @Override
+ public boolean hasNext() {
+ return iter.hasNext();
+ }
+
+ @Override
+ public KeyValue<Integer, byte[]> next() {
+ final Map.Entry<Integer, byte[]> entry = iter.next();
+ return new KeyValue<>(entry.getKey(), entry.getValue());
+ }
+
+ @Override
+ public Integer peekNextKey() {
+ throw new UnsupportedOperationException();
+ }
+
+ @Override
+ public void close() {
+ if (closeCallback != null) closeCallback.run();
+ }
+ }
+}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java
new file mode 100644
index 00000000000..0d99e8cebce
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/StagedMergeIteratorTest.java
@@ -0,0 +1,362 @@
+/*
+ * 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 org.apache.kafka.common.utils.Bytes;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.state.KeyValueIterator;
+
+import org.junit.jupiter.api.Test;
+
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.NavigableMap;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.TreeMap;
+
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class StagedMergeIteratorTest {
+
+ private static Bytes key(final String k) {
+ return Bytes.wrap(k.getBytes(StandardCharsets.UTF_8));
+ }
+
+ private static byte[] val(final String v) {
+ return v.getBytes(StandardCharsets.UTF_8);
+ }
+
+ private static String str(final byte[] bytes) {
+ return new String(bytes, StandardCharsets.UTF_8);
+ }
+
+ @Test
+ public void shouldMergeDisjointStagingAndBase() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.of(val("staged-a")));
+ staging.put(key("c"), Optional.of(val("staged-c")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("b"), val("base-b")),
+ new KeyValue<>(key("d"), val("base-d"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "b", "c", "d"), keys);
+ }
+ }
+
+ @Test
+ public void shouldPreferStagingOverBaseForSameKey() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.of(val("staged-a")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ assertTrue(iter.hasNext());
+ final KeyValue<Bytes, byte[]> entry = iter.next();
+ assertEquals("a", entry.key.toString());
+ assertEquals("staged-a", str(entry.value));
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldSkipTombstones() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.empty()); // tombstone
+ staging.put(key("c"), Optional.of(val("staged-c")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("b"), val("base-b"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("b", "c"), keys);
+ }
+ }
+
+ @Test
+ public void shouldSkipBaseKeyWhenStagingHasTombstoneForSameKey() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("b"), Optional.empty()); // tombstone for key in base
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a")),
+ new KeyValue<>(key("b"), val("base-b")),
+ new KeyValue<>(key("c"), val("base-c"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "c"), keys);
+ }
+ }
+
+ @Test
+ public void shouldHandleEmptyStaging() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a")),
+ new KeyValue<>(key("b"), val("base-b"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "b"), keys);
+ }
+ }
+
+ @Test
+ public void shouldHandleEmptyBase() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.of(val("staged-a")));
+ staging.put(key("b"), Optional.of(val("staged-b")));
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(List.of()))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "b"), keys);
+ }
+ }
+
+ @Test
+ public void shouldHandleBothEmpty() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(List.of()))) {
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldHandleAllTombstones() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.empty());
+ staging.put(key("b"), Optional.empty());
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a")),
+ new KeyValue<>(key("b"), val("base-b"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldSupportPeekNextKey() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("b"), Optional.of(val("staged-b")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ assertEquals(key("a"), iter.peekNextKey());
+ iter.next(); // consume a
+ assertEquals(key("b"), iter.peekNextKey());
+ iter.next(); // consume b
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldThrowOnNextWhenExhausted() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(List.of()))) {
+ assertThrows(NoSuchElementException.class, iter::next);
+ }
+ }
+
+ @Test
+ public void shouldHandleConsecutiveTombstonesInStaging() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.empty());
+ staging.put(key("b"), Optional.empty());
+ staging.put(key("c"), Optional.empty());
+ staging.put(key("d"), Optional.of(val("staged-d")));
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(List.of()))) {
+ assertTrue(iter.hasNext());
+ assertEquals("d", iter.next().key.toString());
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldMaintainSortOrderWithInterleavedEntries() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("b"), Optional.of(val("staged-b")));
+ staging.put(key("d"), Optional.of(val("staged-d")));
+ staging.put(key("f"), Optional.of(val("staged-f")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a")),
+ new KeyValue<>(key("c"), val("base-c")),
+ new KeyValue<>(key("e"), val("base-e")),
+ new KeyValue<>(key("g"), val("base-g"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries))) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "b", "c", "d", "e", "f", "g"), keys);
+ }
+ }
+
+ @Test
+ public void shouldReverseMergeDisjointStagingAndBase() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.of(val("staged-a")));
+ staging.put(key("c"), Optional.of(val("staged-c")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("d"), val("base-d")),
+ new KeyValue<>(key("b"), val("base-b"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries), false)) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("d", "c", "b", "a"), keys);
+ }
+ }
+
+ @Test
+ public void shouldReversePreferStagingOverBaseForSameKey() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("a"), Optional.of(val("staged-a")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("a"), val("base-a"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries), false)) {
+ assertTrue(iter.hasNext());
+ final KeyValue<Bytes, byte[]> entry = iter.next();
+ assertEquals("a", entry.key.toString());
+ assertEquals("staged-a", str(entry.value));
+ assertFalse(iter.hasNext());
+ }
+ }
+
+ @Test
+ public void shouldReverseSkipTombstones() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("c"), Optional.empty()); // tombstone
+ staging.put(key("a"), Optional.of(val("staged-a")));
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("b"), val("base-b"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries), false)) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("b", "a"), keys);
+ }
+ }
+
+ @Test
+ public void shouldReverseSkipBaseKeyWhenStagingHasTombstoneForSameKey() {
+ final NavigableMap<Bytes, Optional<byte[]>> staging = new TreeMap<>();
+ staging.put(key("b"), Optional.empty()); // tombstone for key in base
+
+ final List<KeyValue<Bytes, byte[]>> baseEntries = List.of(
+ new KeyValue<>(key("c"), val("base-c")),
+ new KeyValue<>(key("b"), val("base-b")),
+ new KeyValue<>(key("a"), val("base-a"))
+ );
+
+ try (KeyValueIterator<Bytes, byte[]> iter = new
StagedMergeIterator(staging, new ListIterator(baseEntries), false)) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("c", "a"), keys);
+ }
+ }
+
+ /**
+ * Simple list-backed KeyValueIterator for testing.
+ */
+ private static class ListIterator implements KeyValueIterator<Bytes,
byte[]> {
+ private final List<KeyValue<Bytes, byte[]>> entries;
+ private int index = 0;
+
+ ListIterator(final List<KeyValue<Bytes, byte[]>> entries) {
+ this.entries = new ArrayList<>(entries);
+ }
+
+ @Override
+ public boolean hasNext() {
+ return index < entries.size();
+ }
+
+ @Override
+ public KeyValue<Bytes, byte[]> next() {
+ if (!hasNext()) throw new NoSuchElementException();
+ return entries.get(index++);
+ }
+
+ @Override
+ public Bytes peekNextKey() {
+ if (!hasNext()) throw new NoSuchElementException();
+ return entries.get(index).key;
+ }
+
+ @Override
+ public void close() {
+ // no-op
+ }
+ }
+}
