HDDS-12742. Make RDBStoreAbstractIterator Return Reference-Counted KeyValues

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, UncheckedAutoCloseableSupplier<RawKeyValue<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,45 @@ final RAW getPrefix() {
 
   @Override
   public final void forEachRemaining(
-      Consumer<? super Table.KeyValue<RAW, RAW>> action) {
+      Consumer<? super UncheckedAutoCloseableSupplier<RawKeyValue<RAW>>> action) {
     while (hasNext()) {
-      action.accept(next());
+      UncheckedAutoCloseableSupplier<RawKeyValue<RAW>> entry = next();
+      action.accept(entry);
     }
   }
 
   private void setCurrentEntry() {
-    if (rocksDBIterator.get().isValid()) {
+    if (currentEntry != null) {
+      currentEntry.release();
+    }
+
+    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()));
+    return hasNext;
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> next() {
-    setCurrentEntry();
-    if (currentEntry != null) {
+  public final UncheckedAutoCloseableSupplier<RawKeyValue<RAW>> next() {
+    if (hasNext()) {
+      UncheckedAutoCloseableSupplier<RawKeyValue<RAW>> entry = currentEntry.retainAndReleaseOnClose();
+      this.previousKeyValue = entry.get();
       rocksDBIterator.get().next();
-      return currentEntry;
+      setCurrentEntry();
+      return entry;
     }
     throw new NoSuchElementException("RocksDB Store has no more elements");
   }
@@ -129,19 +148,23 @@ public final void seekToLast() {
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> seek(RAW key) {
+  public final UncheckedAutoCloseableSupplier<RawKeyValue<RAW>> seek(RAW key) {
     seek0(key);
     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 currentEntry.retainAndReleaseOnClose();
+    }
+    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 +173,7 @@ public final void removeFromDB() throws IOException {
   @Override
   public void close() {
     rocksDBIterator.close();
+    closed = true;
+    setCurrentEntry();
   }
 }

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,100 @@
 package org.apache.hadoop.hdds.utils.db;
 
 import java.io.IOException;
+import java.util.Deque;
+import java.util.HashSet;
 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.ratis.util.Preconditions;
+import org.apache.ratis.util.ReferenceCountedObject;
 
 /**
  * Implement {@link RDBStoreAbstractIterator} using {@link CodecBuffer}.
+ * Any Key or Value returned will be only valid within the lifecycle of this iterator.
  */
 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 2 buffers one for setting next value and one for sending the current value.
+    maxNumberOfBuffersInMemory = Math.max(2, 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 = new HashSet<>();
+    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 +139,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;
+      });
     }
   }
 

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

@@ -28,6 +28,7 @@
 import org.apache.hadoop.hdds.utils.MetadataKeyFilters;
 import org.apache.hadoop.hdds.utils.db.RocksDatabase.ColumnFamily;
 import org.apache.hadoop.util.Time;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -218,14 +219,63 @@ public TableIterator<byte[], KeyValue<byte[], byte[]>> iterator()
   @Override
   public TableIterator<byte[], KeyValue<byte[], byte[]>> iterator(byte[] prefix)
       throws IOException {
-    return new RDBStoreByteArrayIterator(db.newIterator(family, false), this,
-        prefix);
+    RDBStoreByteArrayIterator itr  = new RDBStoreByteArrayIterator(db.newIterator(family, false), this, prefix);
+    return new TableIterator<byte[], KeyValue<byte[], byte[]>>() {
+      @Override
+      public void seekToFirst() {
+        itr.seekToFirst();
+      }
+
+      @Override
+      public void seekToLast() {
+        itr.seekToLast();
+      }
+
+      @Override
+      public KeyValue<byte[], byte[]> seek(byte[] bytes) {
+        try (UncheckedAutoCloseableSupplier<RawKeyValue<byte[]>> kv = itr.seek(bytes)) {
+          return kv == null ? null : kv.get();
+        }
+      }
+
+      @Override
+      public void removeFromDB() throws IOException {
+        itr.removeFromDB();
+      }
+
+      @Override
+      public void close() {
+        itr.close();
+      }
+
+      @Override
+      public boolean hasNext() {
+        return itr.hasNext();
+      }
+
+      @Override
+      public KeyValue<byte[], byte[]> next() {
+        try (UncheckedAutoCloseableSupplier<RawKeyValue<byte[]>> kv = itr.next()) {
+          return kv.get();
+        }
+      }
+    };
   }
 
-  TableIterator<CodecBuffer, KeyValue<CodecBuffer, CodecBuffer>> iterator(
+  public TableIterator<byte[], UncheckedAutoCloseableSupplier<RawKeyValue<byte[]>>> closeableSupplierIterator(
+      byte[] prefix) throws IOException {
+    return new RDBStoreByteArrayIterator(db.newIterator(family, false), this, prefix);
+  }
+
+  TableIterator<CodecBuffer, UncheckedAutoCloseableSupplier<RawKeyValue<CodecBuffer>>> iterator(
       CodecBuffer prefix) throws IOException {
+    return iterator(prefix, 2);
+  }
+
+  TableIterator<CodecBuffer, UncheckedAutoCloseableSupplier<RawKeyValue<CodecBuffer>>> iterator(
+      CodecBuffer prefix, int maxNumberOfBuffers) throws IOException {
     return new RDBStoreCodecBufferIterator(db.newIterator(family, false),
-        this, prefix);
+        this, prefix, maxNumberOfBuffers);
   }
 
   @Override