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 d2cbb26e2b4 KAFKA-20494: Transactional InMemoryKeyValueStore (#22626)
d2cbb26e2b4 is described below
commit d2cbb26e2b4dec3fc96c5fb666f74af5104793ca
Author: Nick Telford <[email protected]>
AuthorDate: Mon Jun 22 21:29:36 2026 +0100
KAFKA-20494: Transactional InMemoryKeyValueStore (#22626)
Adds a transactional implementation of `InMemoryKeyValueStore`, building
on the transaction-buffer framework (KAFKA-20490).
`InMemoryTransactionBuffer` implements `TransactionBuffer` for in-memory
stores backed by a `ConcurrentNavigableMap`: writes are staged in the
inherited `pendingWrites` map and flushed to the backing map on commit,
while base iterators are created directly from the backing map. Reads
merge staged writes over the base iterator so a transaction observes its
own uncommitted writes.
`InMemoryKeyValueStore` routes reads and writes through the buffer only
when transactional state stores are enabled, leaving the
non-transactional path unchanged. Restore bypasses the buffer to write
directly to the backing map, and the store reports its pending
write-buffer size via `approximateNumUncommittedBytes()`, feeding the
uncommitted-bytes limit added in KAFKA-20491.
Unlike RocksDB stores, in-memory stores never use `deleteRange`, so the
buffer carries no range-tombstone handling.
This is part of
[KIP-892](https://cwiki.apache.org/confluence/display/KAFKA/KIP-892%3A+Transactional+Semantics+for+State+Stores).
Reviewers: Bill Bejeck <[email protected]>
---
.../state/internals/AbstractTransactionBuffer.java | 25 ++-
.../state/internals/InMemoryKeyValueStore.java | 61 +++++-
.../state/internals/InMemoryTransactionBuffer.java | 174 ++++++++++++++++
.../internals/AbstractTransactionBufferTest.java | 7 +
.../internals/InMemoryTransactionBufferTest.java | 226 +++++++++++++++++++++
5 files changed, 487 insertions(+), 6 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 21e71745c33..b67f0e7b191 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
@@ -58,6 +58,18 @@ abstract class AbstractTransactionBuffer<K extends
Comparable<K>> implements Tra
return newBaseIterator(from, to);
}
+ /**
+ * Create a base iterator that is isolated from concurrent base-store
mutation.
+ * <p>
+ * Called only from the non-owner ({@link #snapshotScan}) path while the
{@link #snapshotLock}
+ * read-lock is held. Subclasses must provide mutation isolation
appropriate to their backend
+ * (e.g. an eager range copy for in-memory, a native snapshot for RocksDB).
+ * <p>
+ * The owner fast path uses the live {@link #newBaseIterator} and must not
use this method.
+ */
+ abstract ManagedKeyValueIterator<K, byte[]> newBaseSnapshotIterator(K
from, K to,
+
boolean forward, boolean toInclusive);
+
/** Atomically apply the accumulated writes to the base store. */
abstract void flushToBase();
@@ -117,9 +129,14 @@ abstract class AbstractTransactionBuffer<K extends
Comparable<K>> implements Tra
@Override
public void rollback() {
- pendingWrites.clear();
- pendingWritesBytes = 0;
- discardPendingBatch();
+ snapshotLock.writeLock().lock();
+ try {
+ pendingWrites.clear();
+ pendingWritesBytes = 0;
+ discardPendingBatch();
+ } finally {
+ snapshotLock.writeLock().unlock();
+ }
}
@Override
@@ -149,7 +166,7 @@ abstract class AbstractTransactionBuffer<K extends
Comparable<K>> implements Tra
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);
+ final ManagedKeyValueIterator<K, byte[]> baseIter =
newBaseSnapshotIterator(from, to, forward, toInclusive);
return new StagedMergeIterator<>(stagingSnapshot, baseIter,
forward);
} finally {
snapshotLock.readLock().unlock();
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
index 4b5336169bf..ae020cf0c08 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryKeyValueStore.java
@@ -58,6 +58,7 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
private final Position position = Position.emptyPosition();
private volatile boolean open = false;
private StateStoreContext context;
+ private InMemoryTransactionBuffer transactionBuffer;
public InMemoryKeyValueStore(final String name) {
this.name = name;
@@ -85,7 +86,8 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
(RecordBatchingStateRestoreCallback) records -> {
synchronized (position) {
for (final ConsumerRecord<byte[], byte[]> record :
records) {
- put(Bytes.wrap(record.key()), record.value());
+ final Bytes key = Bytes.wrap(record.key());
+ putInternal(key, record.value());
ChangelogRecordDeserializationHelper.applyChecksAndUpdatePosition(
record,
consistencyEnabled,
@@ -99,6 +101,13 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
open = true;
this.context = stateStoreContext;
+ final boolean transactional = StreamsConfig.InternalConfig.getBoolean(
+ stateStoreContext.appConfigs(),
+ StreamsConfig.TRANSACTIONAL_STATE_STORES_CONFIG,
+ false);
+ if (transactional) {
+ this.transactionBuffer = new InMemoryTransactionBuffer(map);
+ }
}
@Override
@@ -133,11 +142,21 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
@Override
public synchronized byte[] get(final Bytes key) {
+ if (transactionBuffer != null) {
+ final java.util.Optional<byte[]> staged =
transactionBuffer.get(key);
+ if (staged != null) {
+ return staged.orElse(null);
+ }
+ }
return map.get(key);
}
@Override
public synchronized void put(final Bytes key, final byte[] value) {
+ if (transactionBuffer != null) {
+ transactionBuffer.stage(key, value);
+ return;
+ }
putInternal(key, value);
}
@@ -165,6 +184,12 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
@Override
public synchronized void putAll(final List<KeyValue<Bytes, byte[]>>
entries) {
+ if (transactionBuffer != null) {
+ for (final KeyValue<Bytes, byte[]> entry : entries) {
+ transactionBuffer.stage(entry.key, entry.value);
+ }
+ return;
+ }
for (final KeyValue<Bytes, byte[]> entry : entries) {
putInternal(entry.key, entry.value);
}
@@ -176,11 +201,19 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
final Bytes from = Bytes.wrap(prefixKeySerializer.serialize(null,
prefix));
final Bytes to = ByteUtils.increment(from);
+ if (transactionBuffer != null) {
+ return transactionBuffer.range(from, to, true, false);
+ }
return new InMemoryKeyValueIterator(map.subMap(from, true, to,
false).keySet(), true);
}
@Override
public synchronized byte[] delete(final Bytes key) {
+ if (transactionBuffer != null) {
+ final byte[] oldValue = get(key);
+ transactionBuffer.stage(key, null);
+ return oldValue;
+ }
return map.remove(key);
}
@@ -195,6 +228,9 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
}
private KeyValueIterator<Bytes, byte[]> range(final Bytes from, final
Bytes to, final boolean forward) {
+ if (transactionBuffer != null) {
+ return transactionBuffer.range(from, to, forward, true);
+ }
if (from == null && to == null) {
return getKeyValueIterator(map.keySet(), forward);
} else if (from == null) {
@@ -231,13 +267,34 @@ public class InMemoryKeyValueStore implements
KeyValueStore<Bytes, byte[]> {
return map.size();
}
+ @Override
+ public long approximateNumUncommittedBytes() {
+ if (transactionBuffer != null) {
+ return transactionBuffer.approximateNumUncommittedBytes();
+ }
+ return 0;
+ }
+
@Override
public void commit(final Map<TopicPartition, Long> changelogOffsets) {
- // do-nothing since it is in-memory
+ commitStagedWrites();
+ }
+
+ void commitStagedWrites() {
+ if (transactionBuffer != null) {
+ transactionBuffer.commit();
+ }
+ }
+
+ void rollbackStagedWrites() {
+ if (transactionBuffer != null) {
+ transactionBuffer.rollback();
+ }
}
@Override
public void close() {
+ rollbackStagedWrites();
map.clear();
open = false;
}
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBuffer.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBuffer.java
new file mode 100644
index 00000000000..b1b2f319f5c
--- /dev/null
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBuffer.java
@@ -0,0 +1,174 @@
+/*
+ * 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 java.util.Map;
+import java.util.NavigableMap;
+import java.util.NoSuchElementException;
+import java.util.Optional;
+import java.util.TreeMap;
+
+/**
+ * A {@link TransactionBuffer} implementation for {@link
InMemoryKeyValueStore}.
+ * <p>
+ * The owner thread (StreamThread) scans use a lock-free fast path: a live
iterator over
+ * the base {@code TreeMap} (owner is the sole writer; no concurrent mutation
during its own
+ * scan) merged live with the staged {@link
java.util.concurrent.ConcurrentSkipListMap}.
+ * <p>
+ * Non-owner (Interactive Query) scans use {@link #newBaseSnapshotIterator},
which eagerly
+ * copies the bounded range into a fresh {@link java.util.TreeMap} while the
+ * {@link AbstractTransactionBuffer#snapshotLock} read-lock is held. Because
commit acquires
+ * the write-lock before mutating the base map, the copy is atomic with
respect to commit,
+ * giving IQ threads true point-in-time snapshot isolation. After construction
the IQ iterator
+ * never touches the live base map.
+ */
+class InMemoryTransactionBuffer extends AbstractTransactionBuffer<Bytes> {
+
+ private final NavigableMap<Bytes, byte[]> baseMap;
+
+ InMemoryTransactionBuffer(final NavigableMap<Bytes, byte[]> baseMap) {
+ this.baseMap = baseMap;
+ }
+
+ @Override
+ int estimateKeySize(final Bytes key) {
+ return key.get().length;
+ }
+
+ @Override
+ void stageToBackend(final Bytes key, final byte[] value) {
+ // no-op — staging map is sufficient; no write-batch concept for
in-memory
+ }
+
+ @Override
+ ManagedKeyValueIterator<Bytes, byte[]> newBaseIterator(final Bytes from,
final Bytes to) {
+ return newBaseIterator(from, to, true, true);
+ }
+
+ @Override
+ ManagedKeyValueIterator<Bytes, byte[]> newBaseIterator(final Bytes from,
final Bytes to,
+ final boolean
forward, final boolean toInclusive) {
+ final NavigableMap<Bytes, byte[]> view = boundView(from, to,
toInclusive);
+ return new BaseMapIterator(forward ? view : view.descendingMap());
+ }
+
+ /**
+ * Non-owner (IQ) path: eagerly copies the bounded range from the base map
into a fresh
+ * {@link java.util.TreeMap} while the caller holds the snapshot
read-lock, providing true
+ * point-in-time snapshot isolation. The returned iterator never touches
the live base map.
+ */
+ @Override
+ ManagedKeyValueIterator<Bytes, byte[]> newBaseSnapshotIterator(final Bytes
from, final Bytes to,
+ final
boolean forward, final boolean toInclusive) {
+ final NavigableMap<Bytes, byte[]> copy = new TreeMap<>(boundView(from,
to, toInclusive));
+ return new BaseMapIterator(forward ? copy : copy.descendingMap());
+ }
+
+ private NavigableMap<Bytes, byte[]> boundView(final Bytes from, final
Bytes to, final boolean toInclusive) {
+ if (from != null && to != null) {
+ return baseMap.subMap(from, true, to, toInclusive);
+ } else if (from != null) {
+ return baseMap.tailMap(from, true);
+ } else if (to != null) {
+ return baseMap.headMap(to, toInclusive);
+ } else {
+ return baseMap;
+ }
+ }
+
+ @Override
+ void flushToBase() {
+ for (final Map.Entry<Bytes, Optional<byte[]>> entry :
pendingWrites.entrySet()) {
+ if (entry.getValue().isPresent()) {
+ baseMap.put(entry.getKey(), entry.getValue().get());
+ } else {
+ baseMap.remove(entry.getKey());
+ }
+ }
+ }
+
+ @Override
+ void discardPendingBatch() {
+ // no-op — no backend batch to discard
+ }
+
+ /**
+ * An iterator over the entries of a supplied {@link NavigableMap} view of
the base store.
+ * When fed a private copy (the IQ snapshot path) it provides isolation
from concurrent
+ * base-map mutation. When fed a live sub-view (the owner fast path) it is
not isolated,
+ * but that is safe because the owner thread is the sole writer of the
base map.
+ */
+ static class BaseMapIterator implements ManagedKeyValueIterator<Bytes,
byte[]> {
+ private final java.util.Iterator<Map.Entry<Bytes, byte[]>> delegate;
+ private Map.Entry<Bytes, byte[]> next;
+ private boolean closed = false;
+ private Runnable closeCallback;
+
+ BaseMapIterator(final NavigableMap<Bytes, byte[]> map) {
+ this.delegate = map.entrySet().iterator();
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (closed) {
+ throw new IllegalStateException("Iterator has already been
closed.");
+ }
+ if (next != null) {
+ return true;
+ }
+ if (delegate.hasNext()) {
+ next = delegate.next();
+ return true;
+ }
+ return false;
+ }
+
+ @Override
+ public KeyValue<Bytes, byte[]> next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ final Map.Entry<Bytes, byte[]> entry = next;
+ next = null;
+ return new KeyValue<>(entry.getKey(), entry.getValue());
+ }
+
+ @Override
+ public Bytes peekNextKey() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return next.getKey();
+ }
+
+ @Override
+ public void onClose(final Runnable closeCallback) {
+ this.closeCallback = closeCallback;
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ if (closeCallback != null) {
+ closeCallback.run();
+ }
+ }
+ }
+}
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 9527b4b24af..86d871f6201 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
@@ -384,6 +384,13 @@ public class AbstractTransactionBufferTest {
return base;
}
+ @Override
+ ManagedKeyValueIterator<Integer, byte[]> newBaseSnapshotIterator(
+ final Integer from, final Integer to, final boolean forward, final
boolean toInclusive) {
+ final NavigableMap<Integer, byte[]> copy = new
TreeMap<>(boundedBaseView(from, to, toInclusive));
+ return new TestIterator(forward ? copy : copy.descendingMap());
+ }
+
@Override
void flushToBase() {
for (final Map.Entry<Integer, Optional<byte[]>> entry :
pendingBatch.entrySet()) {
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBufferTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBufferTest.java
new file mode 100644
index 00000000000..24d81c1da01
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/InMemoryTransactionBufferTest.java
@@ -0,0 +1,226 @@
+/*
+ * 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.BeforeEach;
+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.Optional;
+import java.util.TreeMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicReference;
+
+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.assertNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+public class InMemoryTransactionBufferTest {
+
+ private NavigableMap<Bytes, byte[]> baseMap;
+ private InMemoryTransactionBuffer buffer;
+
+ 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 bytes == null ? null : new String(bytes,
StandardCharsets.UTF_8);
+ }
+
+ @BeforeEach
+ public void setUp() {
+ baseMap = new TreeMap<>();
+ baseMap.put(key("a"), val("base-a"));
+ baseMap.put(key("b"), val("base-b"));
+ baseMap.put(key("c"), val("base-c"));
+ buffer = new InMemoryTransactionBuffer(baseMap);
+ }
+
+ @Test
+ public void shouldReturnNullForUnstagedKey() {
+ assertNull(buffer.get(key("a")));
+ assertNull(buffer.get(key("z")));
+ }
+
+ @Test
+ public void shouldReturnStagedValue() {
+ buffer.stage(key("a"), val("staged-a"));
+ final Optional<byte[]> staged = buffer.get(key("a"));
+ assertTrue(staged.isPresent());
+ assertArrayEquals(val("staged-a"), staged.get());
+ }
+
+ @Test
+ public void shouldReturnEmptyOptionalForStagedTombstone() {
+ buffer.stage(key("a"), null);
+ final Optional<byte[]> staged = buffer.get(key("a"));
+ assertEquals(Optional.empty(), staged);
+ }
+
+ @Test
+ public void shouldReturnStagedValueForNewKey() {
+ buffer.stage(key("z"), val("staged-z"));
+ final Optional<byte[]> staged = buffer.get(key("z"));
+ assertTrue(staged.isPresent());
+ assertArrayEquals(val("staged-z"), staged.get());
+ }
+
+ @Test
+ public void shouldReportIsEmpty() {
+ assertTrue(buffer.isEmpty());
+ buffer.stage(key("x"), val("v"));
+ assertFalse(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldCommitStagedWritesToBase() {
+ buffer.stage(key("a"), val("new-a"));
+ buffer.stage(key("d"), val("new-d"));
+ buffer.stage(key("b"), null); // delete b
+
+ buffer.commit();
+
+ assertEquals("new-a", str(baseMap.get(key("a"))));
+ assertNull(baseMap.get(key("b")));
+ assertEquals("base-c", str(baseMap.get(key("c"))));
+ assertEquals("new-d", str(baseMap.get(key("d"))));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldRollbackWithoutAffectingBase() {
+ buffer.stage(key("a"), val("new-a"));
+ buffer.stage(key("d"), val("new-d"));
+
+ buffer.rollback();
+
+ assertEquals("base-a", str(baseMap.get(key("a"))));
+ assertNull(baseMap.get(key("d")));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldMergeStagedWritesInAllScan() {
+ buffer.stage(key("a"), val("staged-a"));
+ buffer.stage(key("d"), val("staged-d"));
+ buffer.stage(key("b"), null); // tombstone
+
+ try (KeyValueIterator<Bytes, byte[]> iter = buffer.all(true)) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ assertEquals(List.of("a", "c", "d"), keys);
+ }
+ }
+
+ @Test
+ public void shouldMergeStagedWritesInOwnerAllScan() {
+ buffer.stage(key("a"), val("staged-a"));
+ buffer.stage(key("d"), val("staged-d"));
+
+ try (KeyValueIterator<Bytes, byte[]> iter = buffer.all(true)) {
+ 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 shouldMergeStagedWritesInRangeScan() {
+ buffer.stage(key("b"), val("staged-b"));
+
+ try (KeyValueIterator<Bytes, byte[]> iter = buffer.range(key("a"),
key("c"), true, true)) {
+ final List<String> keys = new ArrayList<>();
+ final List<String> values = new ArrayList<>();
+ while (iter.hasNext()) {
+ final KeyValue<Bytes, byte[]> entry = iter.next();
+ keys.add(entry.key.toString());
+ values.add(str(entry.value));
+ }
+ assertEquals(List.of("a", "b", "c"), keys);
+ assertEquals(List.of("base-a", "staged-b", "base-c"), values);
+ }
+ }
+
+ @Test
+ public void shouldSeeStagedWritesFromAnotherThread() throws Exception {
+ buffer.stage(key("x"), val("staged-x"));
+
+ final AtomicReference<Optional<byte[]>> result = new
AtomicReference<>();
+ final CountDownLatch latch = new CountDownLatch(1);
+
+ final Thread reader = new Thread(() -> {
+ result.set(buffer.get(key("x")));
+ latch.countDown();
+ });
+ reader.start();
+ latch.await();
+
+ assertTrue(result.get().isPresent());
+ assertArrayEquals(val("staged-x"), result.get().get());
+ }
+
+ @Test
+ public void shouldScanFromAnotherThread() throws Exception {
+ buffer.stage(key("d"), val("staged-d"));
+
+ final AtomicReference<List<String>> result = new AtomicReference<>();
+ final CountDownLatch latch = new CountDownLatch(1);
+
+ final Thread reader = new Thread(() -> {
+ try (KeyValueIterator<Bytes, byte[]> iter = buffer.all(true)) {
+ final List<String> keys = new ArrayList<>();
+ while (iter.hasNext()) {
+ keys.add(iter.next().key.toString());
+ }
+ result.set(keys);
+ }
+ latch.countDown();
+ });
+ reader.start();
+ latch.await();
+
+ assertEquals(List.of("a", "b", "c", "d"), result.get());
+ }
+
+ @Test
+ public void shouldNotShowStagedWritesInBaseAfterRollback() {
+ buffer.stage(key("x"), val("staged-x"));
+ buffer.rollback();
+
+ assertNull(baseMap.get(key("x")));
+ assertNull(buffer.get(key("x")));
+ }
+
+}
