HDDS-12742. Make RDBStoreAbstractIterator Return Reference-Counted KeyValues

hadoop-hdds/container-service/src/main/java/org/apache/hadoop/ozone/container/metadata/DatanodeTable.java

@@ -89,6 +89,21 @@ public void deleteWithBatch(BatchOperation batch, KEY key)
         "version.");
   }
 
+  @Override
+  public KeyValueSpliterator<KEY, VALUE> spliterator(int maxParallelism, boolean closeOnException) {
+    throw new UnsupportedOperationException("Iterating tables directly is not" +
+        " supported for datanode containers due to differing schema " +
+        "version.");
+  }
+
+  @Override
+  public KeyValueSpliterator<KEY, VALUE> spliterator(KEY startKey, KEY prefix, int maxParallelism,
+      boolean closeOnException) {
+    throw new UnsupportedOperationException("Iterating tables directly is not" +
+        " supported for datanode containers due to differing schema " +
+        "version.");
+  }
+
   @Override
   public String getName() {
     return table.getName();

hadoop-hdds/framework/src/main/java/org/apache/hadoop/hdds/utils/db/RDBStoreAbstractIterator.java

@@ -21,6 +21,8 @@
 import java.util.NoSuchElementException;
 import java.util.function.Consumer;
 import org.apache.hadoop.hdds.utils.db.managed.ManagedRocksIterator;
+import org.apache.ratis.util.ReferenceCountedObject;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -30,30 +32,34 @@
  * @param <RAW> the raw type.
  */
 abstract class RDBStoreAbstractIterator<RAW>
-    implements TableIterator<RAW, Table.KeyValue<RAW, RAW>> {
+    implements TableIterator<RAW, RDBStoreAbstractIterator.AutoCloseableRawKeyValue<RAW>> {
 
   private static final Logger LOG =
       LoggerFactory.getLogger(RDBStoreAbstractIterator.class);
 
   private final ManagedRocksIterator rocksDBIterator;
   private final RDBTable rocksDBTable;
-  private Table.KeyValue<RAW, RAW> currentEntry;
+  private ReferenceCountedObject<RawKeyValue<RAW>> currentEntry;
+  private RawKeyValue<RAW> previousKeyValue;
   // This is for schemas that use a fixed-length
   // prefix for each key.
   private final RAW prefix;
+  private Boolean hasNext;
+  private boolean closed;
 
   RDBStoreAbstractIterator(ManagedRocksIterator iterator, RDBTable table,
       RAW prefix) {
     this.rocksDBIterator = iterator;
     this.rocksDBTable = table;
     this.prefix = prefix;
+    this.currentEntry = null;
+    this.hasNext = false;
+    this.closed = false;
+    this.previousKeyValue = null;
   }
 
-  /** @return the key for the current entry. */
-  abstract RAW key();
-
   /** @return the {@link Table.KeyValue} for the current entry. */
-  abstract Table.KeyValue<RAW, RAW> getKeyValue();
+  abstract ReferenceCountedObject<RawKeyValue<RAW>> getKeyValue();
 
   /** Seek to the given key. */
   abstract void seek0(RAW key);
@@ -78,32 +84,52 @@ final RAW getPrefix() {
 
   @Override
   public final void forEachRemaining(
-      Consumer<? super Table.KeyValue<RAW, RAW>> action) {
+      Consumer<? super AutoCloseableRawKeyValue<RAW>> action) {
     while (hasNext()) {
-      action.accept(next());
+      AutoCloseableRawKeyValue<RAW> entry = next();
+      action.accept(entry);
+    }
+  }
+
+  private void releaseEntry() {
+    if (currentEntry != null) {
+      currentEntry.release();
     }
+    currentEntry = null;
+    hasNext = null;
   }
 
   private void setCurrentEntry() {
-    if (rocksDBIterator.get().isValid()) {
+    boolean isValid = !closed && rocksDBIterator.get().isValid();
+    if (isValid) {
       currentEntry = getKeyValue();
+      currentEntry.retain();
     } else {
       currentEntry = null;
     }
+    setHasNext(isValid, currentEntry);
+  }
+
+  public void setHasNext(boolean isValid, ReferenceCountedObject<RawKeyValue<RAW>> entry) {
+    this.hasNext = isValid && (prefix == null || startsWithPrefix(entry.get().getKey()));
   }
 
   @Override
   public final boolean hasNext() {
-    return rocksDBIterator.get().isValid() &&
-        (prefix == null || startsWithPrefix(key()));
+    if (hasNext == null) {
+      setCurrentEntry();
+    }
+    return hasNext;
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> next() {
-    setCurrentEntry();
-    if (currentEntry != null) {
+  public final AutoCloseableRawKeyValue<RAW> next() {
+    if (hasNext()) {
+      AutoCloseableRawKeyValue<RAW> entry = new AutoCloseableRawKeyValue<>(currentEntry);
+      this.previousKeyValue = currentEntry.get();
       rocksDBIterator.get().next();
-      return currentEntry;
+      releaseEntry();
+      return entry;
     }
     throw new NoSuchElementException("RocksDB Store has no more elements");
   }
@@ -115,7 +141,7 @@ public final void seekToFirst() {
     } else {
       seek0(prefix);
     }
-    setCurrentEntry();
+    releaseEntry();
   }
 
   @Override
@@ -125,23 +151,28 @@ public final void seekToLast() {
     } else {
       throw new UnsupportedOperationException("seekToLast: prefix != null");
     }
-    setCurrentEntry();
+    releaseEntry();
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> seek(RAW key) {
+  public final AutoCloseableRawKeyValue<RAW> seek(RAW key) {
     seek0(key);
+    releaseEntry();
     setCurrentEntry();
-    return currentEntry;
+    // Current entry should be only closed when the next() and thus closing the returned entry should be a noop.
+    if (hasNext()) {
+      return new AutoCloseableRawKeyValue<>(currentEntry);
+    }
+    return null;
   }
 
   @Override
   public final void removeFromDB() throws IOException {
     if (rocksDBTable == null) {
       throw new UnsupportedOperationException("remove");
     }
-    if (currentEntry != null) {
-      delete(currentEntry.getKey());
+    if (previousKeyValue != null) {
+      delete(previousKeyValue.getKey());
     } else {
       LOG.info("Failed to removeFromDB: currentEntry == null");
     }
@@ -150,5 +181,21 @@ public final void removeFromDB() throws IOException {
   @Override
   public void close() {
     rocksDBIterator.close();
+    closed = true;
+    releaseEntry();
+  }
+
+  public static final class AutoCloseableRawKeyValue<RAW> extends RawKeyValue<RAW> implements AutoCloseable {
+    private final UncheckedAutoCloseableSupplier<RawKeyValue<RAW>> keyValue;
+
+    private AutoCloseableRawKeyValue(ReferenceCountedObject<RawKeyValue<RAW>> kv) {
+      super(kv.get().getKey(), kv.get().getValue());
+      this.keyValue = kv.retainAndReleaseOnClose();
+    }
+
+    @Override
+    public void close() {
+      keyValue.close();
+    }
   }
 }

hadoop-hdds/framework/src/main/java/org/apache/hadoop/hdds/utils/db/RDBStoreByteArrayIterator.java

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.Arrays;
 import org.apache.hadoop.hdds.utils.db.managed.ManagedRocksIterator;
+import org.apache.ratis.util.ReferenceCountedObject;
 
 /**
  * RocksDB store iterator using the byte[] API.
@@ -33,14 +34,10 @@ class RDBStoreByteArrayIterator extends RDBStoreAbstractIterator<byte[]> {
   }
 
   @Override
-  byte[] key() {
-    return getRocksDBIterator().get().key();
-  }
-
-  @Override
-  Table.KeyValue<byte[], byte[]> getKeyValue() {
+  ReferenceCountedObject<RawKeyValue<byte[]>> getKeyValue() {
     final ManagedRocksIterator i = getRocksDBIterator();
-    return RawKeyValue.create(i.get().key(), i.get().value());
+    RawKeyValue<byte[]> rawKV = RawKeyValue.create(i.get().key(), i.get().value());
+    return ReferenceCountedObject.wrap(rawKV);
   }
 
   @Override

hadoop-hdds/framework/src/main/java/org/apache/hadoop/hdds/utils/db/RDBStoreCodecBufferIterator.java

@@ -18,48 +18,105 @@
 package org.apache.hadoop.hdds.utils.db;
 
 import java.io.IOException;
+import java.util.Deque;
 import java.util.Optional;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.Set;
+import java.util.concurrent.BlockingDeque;
+import java.util.concurrent.LinkedBlockingDeque;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
+import org.apache.commons.lang3.exception.UncheckedInterruptedException;
 import org.apache.hadoop.hdds.utils.db.managed.ManagedRocksIterator;
+import org.apache.hadoop.util.Sets;
 import org.apache.ratis.util.Preconditions;
+import org.apache.ratis.util.ReferenceCountedObject;
 
 /**
- * Implement {@link RDBStoreAbstractIterator} using {@link CodecBuffer}.
- */
+ * An implementation of {@link RDBStoreAbstractIterator} that uses {@link CodecBuffer}
+ * for efficient memory management when iterating over RocksDB entries.
+ * This iterator employs a buffer pooling strategy to minimize memory allocations
+ * during iteration. Key and value buffers are pre-allocated and reused through a
+ * reference-counting mechanism, which significantly reduces GC pressure when
+ * processing large datasets.
+ * Key features:
+ **/
 class RDBStoreCodecBufferIterator extends RDBStoreAbstractIterator<CodecBuffer> {
 
-  private final Buffer keyBuffer;
-  private final Buffer valueBuffer;
-  private final AtomicBoolean closed = new AtomicBoolean();
+  private final BlockingDeque<RawKeyValue<Buffer>> availableBufferStack;
+  private final Set<RawKeyValue<Buffer>> inUseBuffers;
+  private final AtomicReference<Boolean> closed = new AtomicReference<>(false);
 
   RDBStoreCodecBufferIterator(ManagedRocksIterator iterator, RDBTable table,
-      CodecBuffer prefix) {
+      CodecBuffer prefix, int maxNumberOfBuffersInMemory) {
     super(iterator, table, prefix);
-
+    // We need atleast 1 buffers one for setting next value and one for sending the current value.
+    maxNumberOfBuffersInMemory = Math.max(1, maxNumberOfBuffersInMemory);
     final String name = table != null ? table.getName() : null;
-    this.keyBuffer = new Buffer(
-        new CodecBuffer.Capacity(name + "-iterator-key", 1 << 10),
-        buffer -> getRocksDBIterator().get().key(buffer));
-    this.valueBuffer = new Buffer(
-        new CodecBuffer.Capacity(name + "-iterator-value", 4 << 10),
-        buffer -> getRocksDBIterator().get().value(buffer));
+    this.availableBufferStack = new LinkedBlockingDeque<>(maxNumberOfBuffersInMemory);
+    this.inUseBuffers = Sets.newConcurrentHashSet();
+    for (int i = 0; i < maxNumberOfBuffersInMemory; i++) {
+      Buffer keyBuffer = new Buffer(
+          new CodecBuffer.Capacity(name + "-iterator-key-" + i, 1 << 10),
+          buffer -> getRocksDBIterator().get().key(buffer));
+      Buffer valueBuffer = new Buffer(
+          new CodecBuffer.Capacity(name + "-iterator-value-" + i, 4 << 10),
+          buffer -> getRocksDBIterator().get().value(buffer));
+      availableBufferStack.add(new RawKeyValue<>(keyBuffer, valueBuffer));
+    }
+
     seekToFirst();
   }
 
   void assertOpen() {
     Preconditions.assertTrue(!closed.get(), "Already closed");
   }
 
-  @Override
-  CodecBuffer key() {
+  private <V> V getFromDeque(BlockingDeque<V> deque, Set<V> inUseSet) {
+    V popped;
+    do {
+      try {
+        popped = deque.poll(1000, TimeUnit.MILLISECONDS);
+      } catch (InterruptedException e) {
+        throw new UncheckedInterruptedException(e);
+      }
+    } while (popped == null);
     assertOpen();
-    return keyBuffer.getFromDb();
+    inUseSet.add(popped);
+    return popped;
+  }
+
+  private ReferenceCountedObject<RawKeyValue<CodecBuffer>> getReferenceCountedBuffer(
+      RawKeyValue<Buffer> key, Deque<RawKeyValue<Buffer>> stack, Set<RawKeyValue<Buffer>> inUseSet,
+      Function<RawKeyValue<Buffer>, RawKeyValue<CodecBuffer>> transformer) {
+    RawKeyValue<CodecBuffer> value = transformer.apply(key);
+    return ReferenceCountedObject.wrap(value, () -> {
+    }, completelyReleased -> {
+      if (!completelyReleased) {
+        return;
+      }
+      closed.updateAndGet((prev) -> {
+        // If already closed the data structure should not be manipulated since the buffer would have already been
+        // closed.
+        if (!prev) {
+          //Entire block done inside this code block to avoid race condition with close() method.
+          //Remove from the set before adding it back to the stack. Otherwise there could be a race condition with
+          // #getFromDeque function.
+          inUseSet.remove(key);
+          stack.push(key);
+        }
+        return prev;
+      });
+    });
   }
 
   @Override
-  Table.KeyValue<CodecBuffer, CodecBuffer> getKeyValue() {
+  ReferenceCountedObject<RawKeyValue<CodecBuffer>> getKeyValue() {
     assertOpen();
-    return Table.newKeyValue(key(), valueBuffer.getFromDb());
+    RawKeyValue<Buffer> kvBuffer = getFromDeque(availableBufferStack, inUseBuffers);
+    Function<RawKeyValue<Buffer>, RawKeyValue<CodecBuffer>> transformer =
+        kv -> new RawKeyValue<>(kv.getKey().getFromDb(), kv.getValue().getFromDb());
+    return getReferenceCountedBuffer(kvBuffer, availableBufferStack, inUseBuffers, transformer);
   }
 
   @Override
@@ -87,13 +144,26 @@ boolean startsWithPrefix(CodecBuffer key) {
     return key.startsWith(prefix);
   }
 
+  private <V> void release(Deque<V> valueStack, Set<V> inUseSet, Function<V, Void> releaser) {
+    while (!valueStack.isEmpty()) {
+      V popped = valueStack.pop();
+      releaser.apply(popped);
+    }
+    for (V inUseValue : inUseSet) {
+      releaser.apply(inUseValue);
+    }
+  }
+
   @Override
   public void close() {
     if (closed.compareAndSet(false, true)) {
       super.close();
       Optional.ofNullable(getPrefix()).ifPresent(CodecBuffer::release);
-      keyBuffer.release();
-      valueBuffer.release();
+      release(availableBufferStack, inUseBuffers, kv -> {
+        kv.getKey().release();
+        kv.getValue().release();
+        return null;
+      });
     }
   }