bbejeck commented on code in PR #22625: URL: https://github.com/apache/kafka/pull/22625#discussion_r3455180557
########## 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 Review Comment: Is this comment still accurate re: used by `StagedMergeIterator` ? If it needs updating just tack it on to one of the following PRs no need to block this one -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
