KAFKA-14491: [19/N] Combine versioned store RocksDB instances into one

streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegment.java

@@ -16,6 +16,8 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import static org.apache.kafka.streams.state.internals.RocksDBStore.incrementWithoutOverflow;
+
 import java.nio.ByteBuffer;
 import java.util.Arrays;
 import java.util.Collections;
@@ -76,6 +78,12 @@ public int compareTo(final LogicalKeyValueSegment segment) {
 
     @Override
     public synchronized void destroy() {
+        if (id < 0) {
+            throw new IllegalStateException("Negative segment ID indicates a reserved segment, "
+                + "which should not be destroyed. Reserved segments are cleaned up only when "
+                + "an entire store is closed, via the close() method rather than destroy().");
+        }
+
         final Bytes keyPrefix = prefixKeyFormatter.getPrefix();
 
         // this deleteRange() call deletes all entries with the given prefix, because the
@@ -176,7 +184,7 @@ public synchronized KeyValueIterator<Bytes, byte[]> range(final Bytes from, fina
         // with empty bytes from the returned iterator. this filtering is accomplished by
         // passing the prefix filter into StrippedPrefixKeyValueIteratorAdapter().
         final Bytes toBound = to == null
-            ? Bytes.increment(prefixKeyFormatter.getPrefix())
+            ? incrementWithoutOverflow(prefixKeyFormatter.getPrefix())
             : prefixKeyFormatter.addPrefix(to);
         final KeyValueIterator<Bytes, byte[]> iteratorWithKeyPrefixes = physicalStore.range(
             fromBound,

streams/src/main/java/org/apache/kafka/streams/state/internals/LogicalKeyValueSegments.java

@@ -16,15 +16,33 @@
  */
 package org.apache.kafka.streams.state.internals;
 
+import java.util.HashMap;
+import java.util.Map;
 import org.apache.kafka.streams.processor.ProcessorContext;
 import org.apache.kafka.streams.processor.internals.ProcessorContextUtils;
 import org.apache.kafka.streams.state.internals.metrics.RocksDBMetricsRecorder;
 
+/**
+ * A {@link Segments} implementation which uses a single underlying RocksDB instance.
+ * Regular segments with {@code segmentId >= 0} expire according to the specified
+ * retention period. "Reserved" segments with {@code segmentId < 0} do not expire
+ * and are completely separate from regular segments in that methods such as
+ * {@link #getSegmentForTimestamp(long)}, {@link #getOrCreateSegment(long, ProcessorContext)},
+ * {@link #getOrCreateSegmentIfLive(long, ProcessorContext, long)},
+ * {@link #segments(long, long, boolean)}, and {@link #allSegments(boolean)}
+ * only return regular segments and not reserved segments. The methods {@link #flush()}
+ * and {@link #close()} flush and close both regular and reserved segments, due to
+ * the fact that both types of segments share the same physical RocksDB instance.
+ * To create a reserved segment, use {@link #createReservedSegment(long, String)} instead.
+ */
 public class LogicalKeyValueSegments extends AbstractSegments<LogicalKeyValueSegment> {
 
     private final RocksDBMetricsRecorder metricsRecorder;
     private final RocksDBStore physicalStore;
 
+    // reserved segments do not expire, and are tracked here separately from regular segments
+    private final Map<Long, LogicalKeyValueSegment> reservedSegments = new HashMap<>();
+
     LogicalKeyValueSegments(final String name,
                             final String parentDir,
                             final long retentionPeriod,
@@ -41,6 +59,12 @@ public LogicalKeyValueSegment getOrCreateSegment(final long segmentId,
         if (segments.containsKey(segmentId)) {
             return segments.get(segmentId);
         } else {
+            if (segmentId < 0) {
+                throw new IllegalArgumentException(
+                    "Negative segment IDs are reserved for reserved segments, "
+                        + "and should be created through createReservedSegment() instead");
+            }
+
             final LogicalKeyValueSegment newSegment = new LogicalKeyValueSegment(segmentId, segmentName(segmentId), physicalStore);
 
             if (segments.put(segmentId, newSegment) != null) {
@@ -51,6 +75,26 @@ public LogicalKeyValueSegment getOrCreateSegment(final long segmentId,
         }
     }
 
+    LogicalKeyValueSegment createReservedSegment(final long segmentId,
+                                                 final String segmentName) {
+        if (segmentId >= 0) {
+            throw new IllegalArgumentException("segmentId for a reserved segment must be negative");
+        }
+
+        final LogicalKeyValueSegment newSegment = new LogicalKeyValueSegment(segmentId, segmentName, physicalStore);
+
+        if (reservedSegments.put(segmentId, newSegment) != null) {
+            throw new IllegalStateException("LogicalKeyValueSegment already exists.");
+        }
+
+        return newSegment;
+    }
+
+    // VisibleForTesting
+    LogicalKeyValueSegment getReservedSegment(final long segmentId) {
+        return reservedSegments.get(segmentId);
+    }
+
     @Override
     public void openExisting(final ProcessorContext context, final long streamTime) {
         metricsRecorder.init(ProcessorContextUtils.getMetricsImpl(context), context.taskId());
@@ -71,6 +115,24 @@ public void flush() {
     public void close() {
         // close the logical segments first to close any open iterators
         super.close();
+
+        // same for reserved segments
+        for (final LogicalKeyValueSegment segment : reservedSegments.values()) {
+            segment.close();
+        }
+        reservedSegments.clear();
+
         physicalStore.close();
     }
+
+    @Override
+    public String segmentName(final long segmentId) {
+        if (segmentId < 0) {
+            throw new IllegalArgumentException(
+                "Negative segment IDs are reserved for reserved segments, "
+                    + "which have custom names that should not be accessed from this method");
+        }
+
+        return super.segmentName(segmentId);
+    }
 }

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBStore.java

@@ -90,7 +90,7 @@ public class RocksDBStore implements KeyValueStore<Bytes, byte[]>, BatchWritingS
     private static final long BLOCK_CACHE_SIZE = 50 * 1024 * 1024L;
     private static final long BLOCK_SIZE = 4096L;
     private static final int MAX_WRITE_BUFFERS = 3;
-    private static final String DB_FILE_DIR = "rocksdb";
+    static final String DB_FILE_DIR = "rocksdb";
 
     final String name;
     private final String parentDir;
@@ -124,7 +124,6 @@ public class RocksDBStore implements KeyValueStore<Bytes, byte[]>, BatchWritingS
     protected Position position;
     private OffsetCheckpoint positionCheckpoint;
 
-    // VisibleForTesting
     public RocksDBStore(final String name,
                         final String metricsScope) {
         this(name, DB_FILE_DIR, new RocksDBMetricsRecorder(metricsScope, name));
@@ -424,7 +423,10 @@ void deleteRange(final Bytes keyFrom, final Bytes keyTo) {
 
         validateStoreOpen();
 
-        // End of key is exclusive, so we increment it by 1 byte to make keyTo inclusive
+        // End of key is exclusive, so we increment it by 1 byte to make keyTo inclusive.
+        // RocksDB's deleteRange() does not support a null upper bound so in the event
+        // of overflow from increment(), the operation cannot be performed and an
+        // IndexOutOfBoundsException will be thrown.
         dbAccessor.deleteRange(keyFrom.get(), Bytes.increment(keyTo).get());
     }
 
@@ -756,7 +758,7 @@ public ManagedKeyValueIterator<Bytes, byte[]> all(final boolean forward) {
 
         @Override
         public ManagedKeyValueIterator<Bytes, byte[]> prefixScan(final Bytes prefix) {
-            final Bytes to = Bytes.increment(prefix);
+            final Bytes to = incrementWithoutOverflow(prefix);
             return new RocksDBRangeIterator(
                     name,
                     db.newIterator(columnFamily),
@@ -821,4 +823,20 @@ public Options getOptions() {
     public Position getPosition() {
         return position;
     }
+
+    /**
+     * Same as {@link Bytes#increment(Bytes)} but {@code null} is returned instead of throwing
+     * {@code IndexOutOfBoundsException} in the event of overflow.
+     *
+     * @param input bytes to increment
+     * @return A new copy of the incremented byte array, or {@code null} if incrementing would
+     *         result in overflow.
+     */
+    static Bytes incrementWithoutOverflow(final Bytes input) {
+        try {
+            return Bytes.increment(input);
+        } catch (final IndexOutOfBoundsException e) {
+            return null;
+        }
+    }
 }

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBTimestampedStore.java

@@ -236,7 +236,7 @@ public ManagedKeyValueIterator<Bytes, byte[]> all(final boolean forward) {
 
         @Override
         public ManagedKeyValueIterator<Bytes, byte[]> prefixScan(final Bytes prefix) {
-            final Bytes to = Bytes.increment(prefix);
+            final Bytes to = incrementWithoutOverflow(prefix);
             return new RocksDBDualCFRangeIterator(
                 name,
                 db.newIterator(newColumnFamily),

streams/src/main/java/org/apache/kafka/streams/state/internals/RocksDBVersionedStore.java

@@ -17,6 +17,7 @@
 package org.apache.kafka.streams.state.internals;
 
 import static org.apache.kafka.streams.StreamsConfig.InternalConfig.IQ_CONSISTENCY_OFFSET_VECTOR_ENABLED;
+import static org.apache.kafka.streams.state.internals.RocksDBStore.DB_FILE_DIR;
 
 import java.io.File;
 import java.nio.ByteBuffer;
@@ -84,7 +85,7 @@ public class RocksDBVersionedStore implements VersionedKeyValueStore<Bytes, byte
     private final long gracePeriod;
     private final RocksDBMetricsRecorder metricsRecorder;
 
-    private final RocksDBStore latestValueStore;
+    private final LogicalKeyValueSegment latestValueStore; // implemented as a "reserved segment" of the segments store
     private final LogicalKeyValueSegments segmentStores;
     private final RocksDBVersionedStoreClient versionedStoreClient;
     private final RocksDBVersionedStoreRestoreWriteBuffer restoreWriteBuffer;
@@ -106,8 +107,9 @@ public class RocksDBVersionedStore implements VersionedKeyValueStore<Bytes, byte
         // history retention >= grace period, for sound semantics.
         this.gracePeriod = historyRetention;
         this.metricsRecorder = new RocksDBMetricsRecorder(metricsScope, name);
-        this.latestValueStore = new RocksDBStore(latestValueStoreName(name), name, metricsRecorder);
-        this.segmentStores = new LogicalKeyValueSegments(segmentsStoreName(name), name, historyRetention, segmentInterval, metricsRecorder);
+        // pass store name as segments name so state dir subdirectory uses the store name
+        this.segmentStores = new LogicalKeyValueSegments(name, DB_FILE_DIR, historyRetention, segmentInterval, metricsRecorder);
+        this.latestValueStore = this.segmentStores.createReservedSegment(-1L, latestValueStoreName(name));
         this.versionedStoreClient = new RocksDBVersionedStoreClient();
         this.restoreWriteBuffer = new RocksDBVersionedStoreRestoreWriteBuffer(versionedStoreClient);
     }
@@ -246,19 +248,18 @@ public String name() {
 
     @Override
     public void flush() {
-        // order shouldn't matter since failure to flush is a fatal exception
         segmentStores.flush();
-        latestValueStore.flush();
+        // flushing segments store includes flushing latest value store, since they share the
+        // same physical RocksDB instance
     }
 
     @Override
     public void close() {
         open = false;
 
-        // close latest value store first so that calls to get() immediately begin to fail with
-        // store not open, as all calls to get() first get() from latest value store
-        latestValueStore.close();
         segmentStores.close();
+        // closing segments store includes closing latest value store, since they share the
+        // same physical RocksDB instance
     }
 
     @Override
@@ -293,7 +294,6 @@ public void init(final ProcessorContext context, final StateStore root) {
 
         metricsRecorder.init(ProcessorContextUtils.getMetricsImpl(context), context.taskId());
 
-        latestValueStore.openDB(context.appConfigs(), context.stateDir());
         segmentStores.openExisting(context, observedStreamTime);
 
         final File positionCheckpointFile = new File(context.stateDir(), name() + ".position");
@@ -928,8 +928,4 @@ static byte[] from(final byte[] rawValue, final long timestamp) {
     private static String latestValueStoreName(final String storeName) {
         return storeName + ".latestValues";
     }
-
-    private static String segmentsStoreName(final String storeName) {
-        return storeName + ".segments";
-    }
 }