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 1dc4cef0bca KAFKA-20493: Transactional RocksDBStore (#22625)
1dc4cef0bca is described below
commit 1dc4cef0bcade1cfc347af4684dece9c998fa8d4
Author: Nick Telford <[email protected]>
AuthorDate: Mon Jun 22 23:52:19 2026 +0100
KAFKA-20493: Transactional RocksDBStore (#22625)
Adds a transactional implementation of `RocksDBStore`, building on the
transaction-buffer framework (KAFKA-20490) and the dual-column-family
accessor work (KAFKA-20492).
`RocksDBTransactionBuffer` implements `TransactionBuffer` for
RocksDB-backed stores: writes are staged in the inherited
`pendingWrites` map and accumulated in a RocksDB `WriteBatch` that is
applied atomically on commit. Reads merge the staged writes over a base
iterator opened against the underlying RocksDB instance, so transactions
observe their own uncommitted writes. The buffer owns the full
range-tombstone lifecycle — `stageDeleteRange`, the `rangeTombstones`
map, and tombstone-aware overrides of `get`, `all`, `range`, `commit`,
`rollback`, and `isEmpty` — since only RocksDB stores use `deleteRange`.
`TransactionalDBAccessor`, a `DBAccessor` decorator, intercepts reads
and writes to route them through the buffer. It is installed in
`openDB()` only when transactional state stores are enabled, so the
non-transactional path is unchanged. The store reports its pending
write-buffer size via `approximateNumUncommittedBytes()`, feeding the
uncommitted-bytes limit added in KAFKA-20491.
For dual-column-family stores, all `put`/`delete`/`deleteRange` and
`all`/`range`/`prefixScan` calls are routed unconditionally through
CF-aware buffer overloads. This fixes two latent invariants under EOS:
old-CF writes were previously flushed outside the shared `WriteBatch`
(so they could commit independently of new-CF writes), and
in-transaction iterators on the old CF saw the pre-delete committed
value because the in-memory tombstone was never placed in
`pendingWrites`. Because `DualColumnFamilyAccessor` always pairs writes
symmetrically, a single shared `pendingWrites` map yields the correct
final state for both CFs and a single `db.write` on commit applies all
column families atomically.
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]>
---
gradle/spotbugs-exclude.xml | 7 +
.../internals/AbstractColumnFamilyAccessor.java | 5 +
.../state/internals/DualColumnFamilyAccessor.java | 5 +
.../streams/state/internals/RocksDBStore.java | 149 +++++++
.../state/internals/RocksDBTransactionBuffer.java | 323 ++++++++++++++++
.../DualColumnFamilyAccessorTransactionalTest.java | 318 +++++++++++++++
.../internals/RocksDBTimestampedStoreTest.java | 58 +++
.../internals/RocksDBTransactionBufferTest.java | 428 +++++++++++++++++++++
8 files changed, 1293 insertions(+)
diff --git a/gradle/spotbugs-exclude.xml b/gradle/spotbugs-exclude.xml
index 6d18691973a..7268ffc6ebe 100644
--- a/gradle/spotbugs-exclude.xml
+++ b/gradle/spotbugs-exclude.xml
@@ -240,6 +240,13 @@ For a detailed description of spotbugs bug categories, see
https://spotbugs.read
<Bug
pattern="AT_NONATOMIC_64BIT_PRIMITIVE,AT_NONATOMIC_OPERATIONS_ON_SHARED_VARIABLE,AT_STALE_THREAD_WRITE_OF_PRIMITIVE"/>
</Match>
+ <Match>
+ <!-- null means "not staged, fall back to base store";
Optional.empty() means staged tombstone -->
+ <Class
name="org.apache.kafka.streams.state.internals.RocksDBTransactionBuffer"/>
+ <Method name="get"/>
+ <Bug pattern="NP_OPTIONAL_RETURN_NULL"/>
+ </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/AbstractColumnFamilyAccessor.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractColumnFamilyAccessor.java
index eb9958b3a7f..c24b689885e 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractColumnFamilyAccessor.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/AbstractColumnFamilyAccessor.java
@@ -121,6 +121,11 @@ abstract class AbstractColumnFamilyAccessor implements
RocksDBStore.ColumnFamily
}
+ @Override
+ public final ColumnFamilyHandle offsetsColumnFamily() {
+ return offsetColumnFamilyHandle;
+ }
+
// Visible for testing
ColumnFamilyHandle offsetColumnFamilyHandle() {
return offsetColumnFamilyHandle;
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java
index 10664c1326b..4b2f7ba5bf4 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessor.java
@@ -262,6 +262,11 @@ class DualColumnFamilyAccessor extends
AbstractColumnFamilyAccessor {
return oldColumnFamily;
}
+ @Override
+ public ColumnFamilyHandle dataColumnFamily() {
+ return newColumnFamily;
+ }
+
// Visible for testing
ColumnFamilyHandle newColumnFamily() {
return newColumnFamily;
diff --git
a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
index 49b38c23465..e2fb110dcd1 100644
---
a/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java
@@ -282,6 +282,12 @@ public class RocksDBStore implements KeyValueStore<Bytes,
byte[]>, BatchWritingS
throw e;
}
+ final boolean transactional = StreamsConfig.InternalConfig.getBoolean(
+ configs, StreamsConfig.TRANSACTIONAL_STATE_STORES_CONFIG, false);
+ if (transactional) {
+ dbAccessor = new TransactionalDBAccessor(dbAccessor, db,
cfAccessor.dataColumnFamily(), cfAccessor.offsetsColumnFamily(), wOptions,
name);
+ }
+
addValueProvidersToMetricsRecorder();
}
@@ -732,6 +738,14 @@ public class RocksDBStore implements KeyValueStore<Bytes,
byte[]>, BatchWritingS
*
* @return an approximate count of key-value mappings in the store.
*/
+ @Override
+ public long approximateNumUncommittedBytes() {
+ if (dbAccessor instanceof TransactionalDBAccessor) {
+ return ((TransactionalDBAccessor)
dbAccessor).buffer.approximateNumUncommittedBytes();
+ }
+ return 0;
+ }
+
@Override
public long approximateNumEntries() {
validateStoreOpen();
@@ -1034,6 +1048,116 @@ public class RocksDBStore implements
KeyValueStore<Bytes, byte[]>, BatchWritingS
}
+ static class TransactionalDBAccessor implements DBAccessor {
+
+ private final DBAccessor underlying;
+ private final RocksDBTransactionBuffer buffer;
+ private final ColumnFamilyHandle offsetsColumnFamily;
+
+ TransactionalDBAccessor(final DBAccessor underlying,
+ final RocksDB db,
+ final ColumnFamilyHandle dataColumnFamily,
+ final ColumnFamilyHandle offsetsColumnFamily,
+ final WriteOptions wOptions,
+ final String storeName) {
+ this.underlying = underlying;
+ this.offsetsColumnFamily = offsetsColumnFamily;
+ this.buffer = new RocksDBTransactionBuffer(db, dataColumnFamily,
wOptions, storeName);
+ }
+
+ @Override
+ public byte[] get(final ColumnFamilyHandle columnFamily, final byte[]
key) throws RocksDBException {
+ if (!columnFamily.equals(offsetsColumnFamily)) {
+ final java.util.Optional<byte[]> staged =
buffer.get(Bytes.wrap(key));
+ if (staged != null) {
+ return staged.orElse(null);
+ }
+ }
+ return underlying.get(columnFamily, key);
+ }
+
+ @Override
+ public byte[] get(final ColumnFamilyHandle columnFamily, final
ReadOptions readOptions, final byte[] key) throws RocksDBException {
+ if (!columnFamily.equals(offsetsColumnFamily)) {
+ final java.util.Optional<byte[]> staged =
buffer.get(Bytes.wrap(key));
+ if (staged != null) {
+ return staged.orElse(null);
+ }
+ }
+ return underlying.get(columnFamily, readOptions, key);
+ }
+
+ @Override
+ public RocksIterator newIterator(final ColumnFamilyHandle
columnFamily) {
+ return underlying.newIterator(columnFamily);
+ }
+
+ @Override
+ public void put(final ColumnFamilyHandle columnFamily, final byte[]
key, final byte[] value) throws RocksDBException {
+ buffer.stage(columnFamily, Bytes.wrap(key), value);
+ }
+
+ @Override
+ public void delete(final ColumnFamilyHandle columnFamily, final byte[]
key) throws RocksDBException {
+ buffer.stage(columnFamily, Bytes.wrap(key), null);
+ }
+
+ @Override
+ public void deleteRange(final ColumnFamilyHandle columnFamily, final
byte[] from, final byte[] to) throws RocksDBException {
+ buffer.stageDeleteRange(columnFamily, Bytes.wrap(from),
Bytes.wrap(to));
+ }
+
+ @Override
+ public long approximateNumEntries(final ColumnFamilyHandle
columnFamily) throws RocksDBException {
+ return underlying.approximateNumEntries(columnFamily);
+ }
+
+ @Override
+ public void flush(final ColumnFamilyHandle... columnFamilies) throws
RocksDBException {
+ underlying.flush(columnFamilies);
+ }
+
+ @Override
+ public void reset() {
+ underlying.reset();
+ }
+
+ @Override
+ public void close() {
+ buffer.close();
+ underlying.close();
+ }
+
+ @Override
+ public ManagedKeyValueIterator<Bytes, byte[]> all(final
ColumnFamilyHandle cf, final String storeName, final boolean forward) {
+ return buffer.all(cf, forward);
+ }
+
+ @Override
+ public ManagedKeyValueIterator<Bytes, byte[]> range(final
ColumnFamilyHandle cf, final String storeName,
+ final Bytes
from, final Bytes to,
+ final boolean
forward, final boolean toInclusive) {
+ return buffer.range(cf, from, to, forward, toInclusive);
+ }
+
+ @Override
+ public ManagedKeyValueIterator<Bytes, byte[]> prefixScan(final
ColumnFamilyHandle cf, final String storeName,
+ final
Bytes prefix, final Bytes to) {
+ return buffer.range(cf, prefix, to, true, false);
+ }
+
+ @Override
+ public void commitStagedWrites() {
+ buffer.commit();
+ }
+
+ @Override
+ public void rollbackStagedWrites() {
+ buffer.rollback();
+ }
+
+ }
+
interface ColumnFamilyAccessor {
void put(final DBAccessor accessor, final byte[] key, final byte[]
value);
@@ -1086,6 +1210,26 @@ public class RocksDBStore implements
KeyValueStore<Bytes, byte[]>, BatchWritingS
Position open(final RocksDBStore.DBAccessor accessor, final boolean
ignoreInvalidState) throws RocksDBException, StreamsException;
Long getCommittedOffset(final RocksDBStore.DBAccessor accessor, final
TopicPartition partition) throws RocksDBException;
+
+ /**
+ * Returns the primary data column family handle.
+ *
+ * <p>This is the CF that all live puts target. It is passed to the
transaction buffer
+ * as its default CF for WriteBatch staging.
+ *
+ * <p>For dual-CF (upgrade-mode) stores this is the new-format CF,
since all puts
+ * land there and reads check it first.
+ */
+ ColumnFamilyHandle dataColumnFamily();
+
+ /**
+ * Returns the column family handle used to persist offset metadata.
+ *
+ * <p>Reads from this CF must bypass the staged-write buffer so they
always reflect
+ * committed state, guarding against the case where a data key
coincidentally matches
+ * an offset key in the buffer.
+ */
+ ColumnFamilyHandle offsetsColumnFamily();
}
class SingleColumnFamilyAccessor extends AbstractColumnFamilyAccessor {
@@ -1195,6 +1339,11 @@ public class RocksDBStore implements
KeyValueStore<Bytes, byte[]>, BatchWritingS
}
}
+ @Override
+ public ColumnFamilyHandle dataColumnFamily() {
+ return columnFamily;
+ }
+
// Visible for testing
ColumnFamilyHandle columnFamily() {
return columnFamily;
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
new file mode 100644
index 00000000000..ec7d4cd41bf
--- /dev/null
+++
b/streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBuffer.java
@@ -0,0 +1,323 @@
+/*
+ * 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.errors.ProcessorStateException;
+
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+import org.rocksdb.RocksIterator;
+import org.rocksdb.WriteBatch;
+import org.rocksdb.WriteOptions;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+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 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
+ * {@link AbstractTransactionBuffer}, so a {@code WriteBatchWithIndex} is not
needed.
+ * <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.
+ */
+class RocksDBTransactionBuffer extends AbstractTransactionBuffer<Bytes> {
+
+ private final RocksDB db;
+ private final ColumnFamilyHandle cfHandle;
+ private final WriteOptions wOptions;
+ private final String storeName;
+ private WriteBatch writeBatch;
+ private volatile NavigableMap<Bytes, List<Bytes>> rangeTombstones =
Collections.emptyNavigableMap();
+
+ RocksDBTransactionBuffer(final RocksDB db,
+ final ColumnFamilyHandle cfHandle,
+ final WriteOptions wOptions,
+ final String storeName) {
+ this.db = db;
+ this.cfHandle = cfHandle;
+ this.wOptions = wOptions;
+ this.storeName = storeName;
+ this.writeBatch = new WriteBatch();
+ }
+
+ @Override
+ int estimateKeySize(final Bytes key) {
+ return key.get().length;
+ }
+
+ @Override
+ void stageToBackend(final Bytes key, final byte[] value) {
+ stage(cfHandle, key, value);
+ }
+
+ /**
+ * Stages a write for an explicit column family. Updates the shared read
buffer
+ * ({@code pendingWrites}) and appends the write to the shared {@link
WriteBatch}
+ * 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);
+ 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);
+ }
+ }
+
+ /**
+ * Stages a range deletion for an explicit column family. Updates the
shared
+ * {@code pendingWrites} and {@code rangeTombstones} so iterators opened
before
+ * commit hide the deleted range, and appends the range delete to the
shared
+ * {@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);
+ 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);
+ }
+ }
+
+ @Override
+ public Optional<byte[]> get(final Bytes key) {
+ final Optional<byte[]> staged = pendingWrites.get(key);
+ if (staged != null) {
+ return staged;
+ }
+ if (isCoveredByRangeTombstone(key, rangeTombstones)) {
+ return Optional.empty();
+ }
+ return null;
+ }
+
+ @Override
+ public boolean isEmpty() {
+ return super.isEmpty() && rangeTombstones.isEmpty();
+ }
+
+ ManagedKeyValueIterator<Bytes, byte[]> all(final ColumnFamilyHandle cf,
final boolean forward) {
+ if (Thread.currentThread() == ownerThread) {
+ final ManagedKeyValueIterator<Bytes, byte[]> baseIter =
newBaseIterator(cf, null, null, forward, true);
+ return new StagedMergeIterator<>(pendingWrites, baseIter, forward);
+ }
+ return snapshotScan(cf, null, null, forward, true);
+ }
+
+ ManagedKeyValueIterator<Bytes, byte[]> range(final ColumnFamilyHandle cf,
+ 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 ManagedKeyValueIterator<Bytes, byte[]> baseIter =
newBaseIterator(cf, from, to, forward, toInclusive);
+ return new StagedMergeIterator<>(stagingView, baseIter, forward);
+ }
+ return snapshotScan(cf, from, to, forward, toInclusive);
+ }
+
+ private ManagedKeyValueIterator<Bytes, byte[]> snapshotScan(final
ColumnFamilyHandle cf,
+ final Bytes
from, final Bytes to,
+ final boolean
forward, final boolean toInclusive) {
+ snapshotLock.readLock().lock();
+ try {
+ 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<>(stagingSnapshot, baseIter,
forward);
+ } finally {
+ snapshotLock.readLock().unlock();
+ }
+ }
+
+ @Override
+ ManagedKeyValueIterator<Bytes, byte[]> newBaseIterator(final Bytes from,
final Bytes to) {
+ return newBaseIterator(cfHandle, from, to, true, true);
+ }
+
+ @Override
+ ManagedKeyValueIterator<Bytes, byte[]> newBaseIterator(final Bytes from,
final Bytes to,
+ final boolean
forward, final boolean toInclusive) {
+ return newBaseIterator(cfHandle, from, to, forward, toInclusive);
+ }
+
+ private ManagedKeyValueIterator<Bytes, byte[]> newBaseIterator(final
ColumnFamilyHandle cf,
+ final Bytes
from, final Bytes to,
+ final
boolean forward, final boolean toInclusive) {
+ final RocksIterator rocksIterator = db.newIterator(cf);
+ final ManagedKeyValueIterator<Bytes, byte[]> iter;
+ if (from != null && to != null) {
+ iter = new RocksDBRangeIterator(storeName, rocksIterator, from,
to, forward, toInclusive);
+ } else if (from != null && forward) {
+ rocksIterator.seek(from.get());
+ iter = new RocksDbIterator(storeName, rocksIterator, true);
+ } else if (!forward) {
+ if (to != null) {
+ rocksIterator.seekForPrev(to.get());
+ } else {
+ rocksIterator.seekToLast();
+ }
+ iter = new RocksDbIterator(storeName, rocksIterator, false);
+ } else {
+ rocksIterator.seekToFirst();
+ iter = new RocksDbIterator(storeName, rocksIterator, true);
+ }
+ // RocksDbIterator requires onClose to be set before close() is called.
+ // Since this iterator is used internally by StagedMergeIterator (not
+ // tracked by RocksDBStore's open-iterator set), use a no-op callback.
+ iter.onClose(() -> { });
+ if (rangeTombstones.isEmpty()) {
+ return iter;
+ }
+ return new RangeTombstoneFilterIterator(iter, rangeTombstones);
+ }
+
+ @Override
+ void flushToBase() {
+ try {
+ db.write(wOptions, writeBatch);
+ } catch (final RocksDBException e) {
+ throw new ProcessorStateException("Error committing transaction
buffer for store " + storeName, e);
+ }
+ writeBatch.close();
+ writeBatch = new WriteBatch();
+ rangeTombstones = Collections.emptyNavigableMap();
+ }
+
+ @Override
+ void discardPendingBatch() {
+ writeBatch.clear();
+ rangeTombstones = Collections.emptyNavigableMap();
+ }
+
+ @Override
+ public long approximateNumUncommittedBytes() {
+ return super.approximateNumUncommittedBytes() +
writeBatch.getDataSize();
+ }
+
+ @Override
+ public void close() {
+ writeBatch.close();
+ }
+
+ static boolean isCoveredByRangeTombstone(final Bytes key,
+ final NavigableMap<Bytes,
List<Bytes>> tombstones) {
+ if (tombstones.isEmpty()) {
+ return false;
+ }
+ for (final Map.Entry<Bytes, List<Bytes>> entry :
tombstones.headMap(key, true).entrySet()) {
+ for (final Bytes to : entry.getValue()) {
+ if (key.compareTo(to) < 0) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+
+ private static class RangeTombstoneFilterIterator implements
ManagedKeyValueIterator<Bytes, byte[]> {
+
+ private final ManagedKeyValueIterator<Bytes, byte[]> wrapped;
+ private final NavigableMap<Bytes, List<Bytes>> tombstones;
+ private KeyValue<Bytes, byte[]> prefetched;
+ private boolean closed = false;
+ private Runnable closeCallback;
+
+ RangeTombstoneFilterIterator(final ManagedKeyValueIterator<Bytes,
byte[]> wrapped,
+ final NavigableMap<Bytes, List<Bytes>>
tombstones) {
+ this.wrapped = wrapped;
+ this.tombstones = tombstones;
+ }
+
+ @Override
+ public void onClose(final Runnable closeCallback) {
+ this.closeCallback = closeCallback;
+ }
+
+ @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<Bytes, byte[]> next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ final KeyValue<Bytes, byte[]> result = prefetched;
+ prefetched = null;
+ return result;
+ }
+
+ @Override
+ public Bytes peekNextKey() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ return prefetched.key;
+ }
+
+ @Override
+ public void close() {
+ closed = true;
+ try {
+ wrapped.close();
+ } finally {
+ if (closeCallback != null) {
+ closeCallback.run();
+ }
+ }
+ }
+
+ private KeyValue<Bytes, byte[]> computeNext() {
+ while (wrapped.hasNext()) {
+ final KeyValue<Bytes, byte[]> entry = wrapped.next();
+ if (!isCoveredByRangeTombstone(entry.key, tombstones)) {
+ return entry;
+ }
+ }
+ return null;
+ }
+ }
+}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessorTransactionalTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessorTransactionalTest.java
new file mode 100644
index 00000000000..d61fc2c6fb7
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/DualColumnFamilyAccessorTransactionalTest.java
@@ -0,0 +1,318 @@
+/*
+ * 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.processor.StateStoreContext;
+import org.apache.kafka.streams.query.Position;
+import org.apache.kafka.test.TestUtils;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.rocksdb.ColumnFamilyDescriptor;
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.ColumnFamilyOptions;
+import org.rocksdb.DBOptions;
+import org.rocksdb.FlushOptions;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+import org.rocksdb.WriteOptions;
+
+import java.io.File;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+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;
+import static org.mockito.Mockito.lenient;
+import static org.mockito.Mockito.mock;
+
+/**
+ * Regression tests ensuring that {@link DualColumnFamilyAccessor} scans
honour writes staged in a
+ * {@link RocksDBStore.TransactionalDBAccessor} transaction buffer before they
are committed to
+ * RocksDB.
+ *
+ * Prior to the fix on KIP-892/txn-rocksdb-accessor, DualColumnFamilyAccessor
built its inner
+ * iterators via {@code accessor.newIterator()}, which TransactionalDBAccessor
passes straight
+ * through to RocksDB without merging the buffer. This caused
range/all/prefixScan to miss staged
+ * writes and deletions, breaking read-your-writes under EOS.
+ */
+public class DualColumnFamilyAccessorTransactionalTest {
+
+ static {
+ RocksDB.loadLibrary();
+ }
+
+ private static final String STORE_NAME = "test-store";
+
+ private File dbDir;
+ private RocksDB db;
+ private ColumnFamilyHandle offsetsCFHandle;
+ private ColumnFamilyHandle oldCFHandle;
+ private ColumnFamilyHandle newCFHandle;
+ private WriteOptions wOptions;
+ private FlushOptions flushOptions;
+
+ private RocksDBStore.DBAccessor txnAccessor;
+ private DualColumnFamilyAccessor accessor;
+
+ 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() throws RocksDBException {
+ dbDir = TestUtils.tempDirectory();
+
+ final DBOptions dbOptions = new DBOptions();
+ dbOptions.setCreateIfMissing(true);
+ dbOptions.setCreateMissingColumnFamilies(true);
+
+ final ColumnFamilyOptions cfOptions = new ColumnFamilyOptions();
+ final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
+ new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY,
cfOptions),
+ new
ColumnFamilyDescriptor("old-cf".getBytes(StandardCharsets.UTF_8), cfOptions),
+ new
ColumnFamilyDescriptor("new-cf".getBytes(StandardCharsets.UTF_8), cfOptions)
+ );
+ final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
+ db = RocksDB.open(dbOptions, dbDir.getAbsolutePath(), cfDescriptors,
cfHandles);
+ offsetsCFHandle = cfHandles.get(0);
+ oldCFHandle = cfHandles.get(1);
+ newCFHandle = cfHandles.get(2);
+
+ wOptions = new WriteOptions();
+ wOptions.setDisableWAL(true);
+ flushOptions = new FlushOptions();
+
+ final RocksDBStore.DBAccessor directAccessor =
+ new RocksDBStore.DirectDBAccessor(db, flushOptions, wOptions);
+ txnAccessor = new RocksDBStore.TransactionalDBAccessor(
+ directAccessor, db, newCFHandle, offsetsCFHandle, wOptions,
STORE_NAME);
+
+ final RocksDBStore store = mock(RocksDBStore.class);
+ store.position = Position.emptyPosition();
+ store.context = mock(StateStoreContext.class);
+ lenient().when(store.name()).thenReturn(STORE_NAME);
+
+ accessor = new DualColumnFamilyAccessor(
+ offsetsCFHandle, oldCFHandle, newCFHandle,
+ v -> ("converted:" + new String(v, StandardCharsets.UTF_8))
+ .getBytes(StandardCharsets.UTF_8),
+ store,
+ new AtomicBoolean(true));
+ }
+
+ @AfterEach
+ public void tearDown() throws RocksDBException {
+ if (txnAccessor != null) txnAccessor.close();
+ if (offsetsCFHandle != null) offsetsCFHandle.close();
+ if (oldCFHandle != null) oldCFHandle.close();
+ if (newCFHandle != null) newCFHandle.close();
+ if (db != null) db.close();
+ if (wOptions != null) wOptions.close();
+ if (flushOptions != null) flushOptions.close();
+ }
+
+ private List<KeyValue<Bytes, byte[]>> drain(final
ManagedKeyValueIterator<Bytes, byte[]> it) {
+ it.onClose(() -> { });
+ final List<KeyValue<Bytes, byte[]>> results = new ArrayList<>();
+ while (it.hasNext()) {
+ results.add(it.next());
+ }
+ it.close();
+ return results;
+ }
+
+ @Test
+ public void shouldSeeStagedPutInAll() throws RocksDBException {
+ // Stage a write on the new CF via the dual-CF accessor (no commit).
+ accessor.put(txnAccessor, key("b").get(), val("new-b"));
+
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertEquals(1, results.size());
+ assertArrayEquals(key("b").get(), results.get(0).key.get());
+ assertArrayEquals(val("new-b"), results.get(0).value);
+ }
+
+ @Test
+ public void shouldSeeStagedPutInRange() throws RocksDBException {
+ accessor.put(txnAccessor, key("b").get(), val("new-b"));
+ accessor.put(txnAccessor, key("d").get(), val("new-d"));
+
+ final List<KeyValue<Bytes, byte[]>> results =
+ drain(accessor.range(txnAccessor, key("a"), key("c"), true));
+
+ assertEquals(1, results.size());
+ assertArrayEquals(key("b").get(), results.get(0).key.get());
+ }
+
+ @Test
+ public void shouldSeeStagedPutInPrefixScan() throws RocksDBException {
+ accessor.put(txnAccessor, key("foo:1").get(), val("new-foo1"));
+ accessor.put(txnAccessor, key("bar:1").get(), val("new-bar"));
+
+ final List<KeyValue<Bytes, byte[]>> results =
+ drain(accessor.prefixScan(txnAccessor,
Bytes.wrap("foo:".getBytes(StandardCharsets.UTF_8))));
+
+ assertEquals(1, results.size());
+ assertArrayEquals(key("foo:1").get(), results.get(0).key.get());
+ }
+
+ @Test
+ public void shouldSeeStagedDeleteInRange() throws RocksDBException {
+ // Commit a base value to new CF via a direct write, then stage a
delete.
+ db.put(newCFHandle, wOptions, key("b").get(), val("committed-b"));
+
+ accessor.put(txnAccessor, key("b").get(), null);
+
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertTrue(results.isEmpty(), "staged delete should suppress the
committed value");
+ }
+
+ @Test
+ public void shouldSeeStagedDeleteRangeInRange() throws RocksDBException {
+ // Commit base values to new CF, then stage a deleteRange.
+ db.put(newCFHandle, wOptions, key("a").get(), val("a"));
+ db.put(newCFHandle, wOptions, key("b").get(), val("b"));
+ db.put(newCFHandle, wOptions, key("c").get(), val("c"));
+
+ accessor.deleteRange(txnAccessor, key("a").get(), key("c").get());
+
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertEquals(1, results.size(), "keys a and b should be hidden by the
staged range tombstone");
+ assertArrayEquals(key("c").get(), results.get(0).key.get());
+ }
+
+ @Test
+ public void shouldPreferStagedNewValueOverCommittedOldValue() throws
RocksDBException {
+ // Populate old CF only with a committed entry.
+ db.put(oldCFHandle, wOptions, key("x").get(), val("old-x"));
+
+ // Migrate: put() stages delete on old CF and put on new CF.
+ accessor.put(txnAccessor, key("x").get(), val("new-x"));
+
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertEquals(1, results.size(), "same key should appear exactly once");
+ assertArrayEquals(key("x").get(), results.get(0).key.get());
+ assertArrayEquals(val("new-x"), results.get(0).value,
+ "staged new-format value should win over committed old-format
value");
+ }
+
+ @Test
+ public void shouldHideStagedDeleteOfUnmigratedKey() throws
RocksDBException {
+ // Key lives only in the old CF (unmigrated). Stage put(key, null),
which issues
+ // delete(oldCF, key) + delete(newCF, key). Both go into the shared
read buffer, so
+ // scans on either CF see the staged tombstone and suppress the
pre-delete base value.
+ db.put(oldCFHandle, wOptions, key("x").get(), val("old-x"));
+
+ accessor.put(txnAccessor, key("x").get(), null);
+
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertTrue(results.isEmpty(),
+ "staged old-CF tombstone should hide the unmigrated key from
scans");
+ }
+
+ @Test
+ public void shouldSeeCommittedValuesAfterCommit() throws RocksDBException {
+ accessor.put(txnAccessor, key("a").get(), val("a"));
+ accessor.put(txnAccessor, key("b").get(), val("b"));
+ txnAccessor.commitStagedWrites();
+
+ // After commit the values are in RocksDB; a fresh scan should return
both.
+ final List<KeyValue<Bytes, byte[]>> results =
drain(accessor.all(txnAccessor, true));
+
+ assertEquals(2, results.size());
+ assertArrayEquals(key("a").get(), results.get(0).key.get());
+ assertArrayEquals(key("b").get(), results.get(1).key.get());
+ assertFalse(results.isEmpty());
+ }
+
+ // --- Point get() read-your-writes tests ---
+ //
+ // These tests specifically exercise the get() code path through
+ // TransactionalDBAccessor.get(), which checks the staged-write buffer
+ // for any CF except the offsets CF. The buffer is keyed by key only
+ // (no CF), so it must be consulted for both oldCF and newCF reads.
+
+ @Test
+ public void shouldSeeStagedPutViaGet() throws RocksDBException {
+ accessor.put(txnAccessor, key("a").get(), val("new-a"));
+
+ assertArrayEquals(val("new-a"), accessor.get(txnAccessor,
key("a").get()),
+ "staged put should be visible through get() before commit");
+ }
+
+ @Test
+ public void shouldSeeStagedDeleteViaGet() throws RocksDBException {
+ // Commit a base value to the new CF, then stage a delete.
+ db.put(newCFHandle, wOptions, key("a").get(), val("committed-a"));
+
+ accessor.put(txnAccessor, key("a").get(), null);
+
+ assertNull(accessor.get(txnAccessor, key("a").get()),
+ "staged delete should be visible through get() before commit");
+ }
+
+ @Test
+ public void shouldPreferStagedNewValueOverCommittedOldValueViaGet() throws
RocksDBException {
+ // Key exists only in old CF (unmigrated), then we stage a put to
migrate it.
+ db.put(oldCFHandle, wOptions, key("x").get(), val("old-x"));
+
+ accessor.put(txnAccessor, key("x").get(), val("new-x"));
+
+ assertArrayEquals(val("new-x"), accessor.get(txnAccessor,
key("x").get()),
+ "staged new-format value should win over committed old-format
value via get()");
+ }
+
+ @Test
+ public void shouldHideStagedDeleteOfUnmigratedKeyViaGet() throws
RocksDBException {
+ // Key lives only in old CF. Stage put(key, null) which tombstones
both CFs.
+ db.put(oldCFHandle, wOptions, key("x").get(), val("old-x"));
+
+ accessor.put(txnAccessor, key("x").get(), null);
+
+ assertNull(accessor.get(txnAccessor, key("x").get()),
+ "staged old-CF tombstone should hide the unmigrated key via
get()");
+ }
+
+ @Test
+ public void shouldReturnNullForAbsentKeyViaGet() throws RocksDBException {
+ assertNull(accessor.get(txnAccessor, key("absent").get()),
+ "get() for a key with no staged write and no committed value
should return null");
+ }
+}
diff --git
a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreTest.java
b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreTest.java
index 3c8f87bda49..e73cecc96cb 100644
---
a/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreTest.java
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStoreTest.java
@@ -16,13 +16,18 @@
*/
package org.apache.kafka.streams.state.internals;
+import org.apache.kafka.common.TopicPartition;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.apache.kafka.common.utils.Bytes;
import org.apache.kafka.common.utils.LogCaptureAppender;
import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.StreamsConfig;
import org.apache.kafka.streams.errors.ProcessorStateException;
import org.apache.kafka.streams.state.KeyValueIterator;
+import org.apache.kafka.test.InternalMockProcessorContext;
+import org.apache.kafka.test.StreamsTestUtils;
+import org.apache.kafka.test.TestUtils;
import org.hamcrest.core.IsNull;
import org.junit.jupiter.api.Test;
@@ -36,6 +41,8 @@ import java.io.File;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;
+import java.util.Map;
+import java.util.Properties;
import static java.util.Arrays.asList;
import static
org.apache.kafka.streams.state.internals.RocksDBStore.OFFSETS_COLUMN_FAMILY_NAME;
@@ -45,6 +52,7 @@ import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
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.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -501,6 +509,56 @@ public class RocksDBTimestampedStoreTest extends
RocksDBStoreTest {
}
}
+ /**
+ * Regression test for the null {@code dataColumnFamily} bug under EOS.
+ *
+ * <p>When {@code TRANSACTIONAL_STATE_STORES_CONFIG} is enabled, {@link
RocksDBStore#openDB}
+ * wraps the underlying accessor in a {@link
RocksDBStore.TransactionalDBAccessor}. The
+ * accessor's {@code get()} consults the staged-write buffer only when the
request targets
+ * the distinguished data CF ({@code columnFamily.equals(cfHandle)}).
Before the fix,
+ * {@link RocksDBTimestampedStore} never set {@code dataColumnFamily}, so
{@code cfHandle}
+ * was {@code null} and the check was always false — staged writes were
invisible to
+ * {@code get()}, breaking read-your-writes.
+ */
+ @Test
+ public void
shouldReadYourWritesViaGetWhenTransactionalTimestampedStoreOpenedUnderEOS() {
+ final Properties props = StreamsTestUtils.getStreamsConfig();
+ props.setProperty(StreamsConfig.PROCESSING_GUARANTEE_CONFIG,
StreamsConfig.EXACTLY_ONCE_V2);
+ props.setProperty(StreamsConfig.TRANSACTIONAL_STATE_STORES_CONFIG,
"true");
+ final File stateDir = TestUtils.tempDirectory();
+ final InternalMockProcessorContext<?, ?> txnContext =
+ new InternalMockProcessorContext<>(stateDir, new
StreamsConfig(props));
+
+ final RocksDBStore txnStore = new RocksDBTimestampedStore(DB_NAME,
METRICS_SCOPE);
+ try {
+ txnStore.init(txnContext, txnStore);
+
+ final Bytes key = new Bytes("k1".getBytes());
+ final byte[] value = "v1".getBytes();
+
+ // Put a value — this is staged in the transaction buffer but not
yet committed to RocksDB.
+ txnStore.put(key, value);
+
+ // get() must see the staged write (read-your-writes).
+ assertArrayEquals(value, txnStore.get(key),
+ "transactional TimestampedStore should return staged value
via get() before commit");
+
+ // Commit and verify the value persists.
+ final TopicPartition tp = new TopicPartition("changelog", 0);
+ txnStore.commit(Map.of(tp, 1L));
+
+ assertArrayEquals(value, txnStore.get(key),
+ "transactional TimestampedStore should return value via
get() after commit");
+
+ // Delete should also be read-your-writes.
+ txnStore.put(key, null);
+ assertNull(txnStore.get(key),
+ "transactional TimestampedStore should return null for a
staged delete via get() before commit");
+ } finally {
+ txnStore.close();
+ }
+ }
+
private void prepareOldStore() {
final RocksDBStore keyValueStore = new RocksDBStore(DB_NAME,
METRICS_SCOPE);
try {
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
new file mode 100644
index 00000000000..0fab3b44e91
--- /dev/null
+++
b/streams/src/test/java/org/apache/kafka/streams/state/internals/RocksDBTransactionBufferTest.java
@@ -0,0 +1,428 @@
+/*
+ * 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.apache.kafka.test.TestUtils;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
+import org.rocksdb.ColumnFamilyDescriptor;
+import org.rocksdb.ColumnFamilyHandle;
+import org.rocksdb.ColumnFamilyOptions;
+import org.rocksdb.DBOptions;
+import org.rocksdb.RocksDB;
+import org.rocksdb.RocksDBException;
+import org.rocksdb.WriteOptions;
+
+import java.io.File;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import java.util.Optional;
+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 RocksDBTransactionBufferTest {
+
+ static {
+ RocksDB.loadLibrary();
+ }
+
+ private File dbDir;
+ private RocksDB db;
+ private ColumnFamilyHandle cfHandle;
+ private ColumnFamilyHandle otherCfHandle;
+ private WriteOptions wOptions;
+ private RocksDBTransactionBuffer 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() throws RocksDBException {
+ dbDir = TestUtils.tempDirectory();
+
+ final DBOptions dbOptions = new DBOptions();
+ dbOptions.setCreateIfMissing(true);
+ dbOptions.setCreateMissingColumnFamilies(true);
+
+ final ColumnFamilyOptions cfOptions = new ColumnFamilyOptions();
+ final List<ColumnFamilyDescriptor> cfDescriptors = Arrays.asList(
+ new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY,
cfOptions),
+ new
ColumnFamilyDescriptor("other-cf".getBytes(StandardCharsets.UTF_8), cfOptions)
+ );
+ final List<ColumnFamilyHandle> cfHandles = new ArrayList<>();
+ db = RocksDB.open(dbOptions, dbDir.getAbsolutePath(), cfDescriptors,
cfHandles);
+ cfHandle = cfHandles.get(0);
+ otherCfHandle = cfHandles.get(1);
+
+ wOptions = new WriteOptions();
+ wOptions.setDisableWAL(true);
+
+ // Pre-populate base data
+ db.put(cfHandle, wOptions, key("a").get(), val("base-a"));
+ db.put(cfHandle, wOptions, key("b").get(), val("base-b"));
+ db.put(cfHandle, wOptions, key("c").get(), val("base-c"));
+
+ buffer = new RocksDBTransactionBuffer(db, cfHandle, wOptions,
"test-store");
+ }
+
+ @AfterEach
+ public void tearDown() {
+ if (buffer != null) buffer.close();
+ if (cfHandle != null) cfHandle.close();
+ if (otherCfHandle != null) otherCfHandle.close();
+ if (db != null) db.close();
+ if (wOptions != null) wOptions.close();
+ }
+
+ @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 shouldCommitStagedWritesToRocksDB() throws RocksDBException {
+ 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(db.get(cfHandle, key("a").get())));
+ assertNull(db.get(cfHandle, key("b").get()));
+ assertEquals("base-c", str(db.get(cfHandle, key("c").get())));
+ assertEquals("new-d", str(db.get(cfHandle, key("d").get())));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldRollbackWithoutAffectingBase() throws RocksDBException {
+ buffer.stage(key("a"), val("new-a"));
+ buffer.stage(key("d"), val("new-d"));
+
+ buffer.rollback();
+
+ assertEquals("base-a", str(db.get(cfHandle, key("a").get())));
+ assertNull(db.get(cfHandle, key("d").get()));
+ 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() throws
RocksDBException {
+ buffer.stage(key("x"), val("staged-x"));
+ buffer.rollback();
+
+ assertNull(db.get(cfHandle, key("x").get()));
+ assertNull(buffer.get(key("x")));
+ }
+
+ @Test
+ public void shouldHideDeletedRangeFromPointReads() {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("c"));
+
+ assertEquals(Optional.empty(), buffer.get(key("a")));
+ assertEquals(Optional.empty(), buffer.get(key("b")));
+ assertNull(buffer.get(key("c"))); // exclusive upper bound
+ }
+
+ @Test
+ public void shouldHideDeletedRangeFromScans() {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("c"));
+
+ 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("c"), keys);
+ }
+ }
+
+ @Test
+ public void shouldAllowStageAfterDeleteRangeToOverride() {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("d"));
+ buffer.stage(key("b"), val("new-b"));
+
+ final Optional<byte[]> result = buffer.get(key("b"));
+ assertTrue(result.isPresent());
+ assertArrayEquals(val("new-b"), result.get());
+ assertEquals(Optional.empty(), buffer.get(key("a")));
+ }
+
+ @Test
+ public void shouldCommitRangeDeleteToRocksDB() throws RocksDBException {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("c"));
+ buffer.commit();
+
+ assertNull(db.get(cfHandle, key("a").get()));
+ assertNull(db.get(cfHandle, key("b").get()));
+ assertEquals("base-c", str(db.get(cfHandle, key("c").get())));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldCommitRangeDeleteWithOverridingPut() throws
RocksDBException {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("d"));
+ buffer.stage(key("b"), val("new-b"));
+ buffer.commit();
+
+ assertNull(db.get(cfHandle, key("a").get()));
+ assertEquals("new-b", str(db.get(cfHandle, key("b").get())));
+ assertNull(db.get(cfHandle, key("c").get()));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldRollbackRangeDelete() throws RocksDBException {
+ buffer.stageDeleteRange(cfHandle, key("a"), key("c"));
+ buffer.rollback();
+
+ assertEquals("base-a", str(db.get(cfHandle, key("a").get())));
+ assertEquals("base-b", str(db.get(cfHandle, key("b").get())));
+ assertEquals("base-c", str(db.get(cfHandle, key("c").get())));
+ assertTrue(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldReportNotEmptyWithRangeTombstones() {
+ assertTrue(buffer.isEmpty());
+ buffer.stageDeleteRange(cfHandle, key("a"), key("c"));
+ assertFalse(buffer.isEmpty());
+ }
+
+ @Test
+ public void shouldSupportMultipleCommitRollbackCycles() throws
RocksDBException {
+ buffer.stage(key("x"), val("v1"));
+ buffer.commit();
+ assertEquals("v1", str(db.get(cfHandle, key("x").get())));
+
+ buffer.stage(key("x"), val("v2"));
+ buffer.rollback();
+ // After rollback, base still has v1
+ assertEquals("v1", str(db.get(cfHandle, key("x").get())));
+ // get() returns null (not staged), caller should check base
+ assertNull(buffer.get(key("x")));
+
+ buffer.stage(key("x"), val("v3"));
+ buffer.commit();
+ assertEquals("v3", str(db.get(cfHandle, key("x").get())));
+ }
+
+ // --- CF-aware overload tests ---
+
+ @Test
+ public void shouldStageWriteToOtherCF() throws RocksDBException {
+ buffer.stage(otherCfHandle, key("x"), val("other-x"));
+
+ // Shared read buffer sees the staged value
+ assertTrue(buffer.get(key("x")).isPresent());
+ assertArrayEquals(val("other-x"), buffer.get(key("x")).get());
+ // Not yet flushed to RocksDB
+ assertNull(db.get(otherCfHandle, key("x").get()));
+
+ buffer.commit();
+ assertEquals("other-x", str(db.get(otherCfHandle, key("x").get())));
+ assertNull(buffer.get(key("x")));
+ }
+
+ @Test
+ public void shouldStageDeleteToOtherCF() throws RocksDBException {
+ db.put(otherCfHandle, wOptions, key("x").get(), val("other-x"));
+
+ buffer.stage(otherCfHandle, key("x"), null);
+
+ // Shared read buffer shows tombstone
+ assertEquals(Optional.empty(), buffer.get(key("x")));
+ // Not yet flushed
+ assertArrayEquals(val("other-x"), db.get(otherCfHandle,
key("x").get()));
+
+ buffer.commit();
+ assertNull(db.get(otherCfHandle, key("x").get()));
+ }
+
+ @Test
+ public void shouldStageDeleteRangeToOtherCF() throws RocksDBException {
+ db.put(otherCfHandle, wOptions, key("a").get(), val("other-a"));
+ db.put(otherCfHandle, wOptions, key("b").get(), val("other-b"));
+
+ buffer.stageDeleteRange(otherCfHandle, key("a"), key("c"));
+
+ // Shared tombstone hides keys in range from buffer reads
+ assertEquals(Optional.empty(), buffer.get(key("a")));
+ assertEquals(Optional.empty(), buffer.get(key("b")));
+ assertNull(buffer.get(key("c"))); // exclusive upper bound: not covered
+
+ buffer.commit();
+ assertNull(db.get(otherCfHandle, key("a").get()));
+ assertNull(db.get(otherCfHandle, key("b").get()));
+ }
+
+ @Test
+ public void shouldCommitWritesAcrossCFsAtomically() throws
RocksDBException {
+ buffer.stage(cfHandle, key("d"), val("data-d"));
+ buffer.stage(otherCfHandle, key("o"), val("other-o"));
+
+ // Neither visible in RocksDB before commit
+ assertNull(db.get(cfHandle, key("d").get()));
+ assertNull(db.get(otherCfHandle, key("o").get()));
+
+ buffer.commit();
+
+ assertEquals("data-d", str(db.get(cfHandle, key("d").get())));
+ assertEquals("other-o", str(db.get(otherCfHandle, key("o").get())));
+ assertTrue(buffer.isEmpty());
+ }
+}
