HDDS-12742. Make RDBStoreAbstractIterator Return Reference-Counted KeyValues

hadoop-hdds/common/src/main/java/org/apache/hadoop/hdds/utils/db/PutToByteBuffer.java

@@ -32,6 +32,6 @@
  * @param <E> The exception type this function may throw.
  */
 @FunctionalInterface
-interface PutToByteBuffer<E extends Exception>
+public interface PutToByteBuffer<E extends Exception>
     extends CheckedFunction<ByteBuffer, Integer, E> {
 }

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/Buffer.java

@@ -0,0 +1,81 @@
+/*
+ * 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.hadoop.hdds.utils.db;
+
+import org.apache.ratis.util.Preconditions;
+
+/**
+ * A utility class for managing an underlying {@link CodecBuffer} with dynamic capacity adjustment
+ * based on the requirements of a data source. This class encapsulates operations to allocate,
+ * prepare, and release the buffer, as well as retrieve data from a source.
+ */
+public class Buffer {
+  private final CodecBuffer.Capacity initialCapacity;
+  private final PutToByteBuffer<RuntimeException> source;
+  private CodecBuffer buffer;
+
+  public Buffer(CodecBuffer.Capacity initialCapacity,
+      PutToByteBuffer<RuntimeException> source) {
+    this.initialCapacity = initialCapacity;
+    this.source = source;
+  }
+
+  public void release() {
+    if (buffer != null) {
+      buffer.release();
+    }
+  }
+
+  public void prepare() {
+    if (buffer == null) {
+      allocate();
+    } else {
+      buffer.clear();
+    }
+  }
+
+  public void allocate() {
+    if (buffer != null) {
+      buffer.release();
+    }
+    buffer = CodecBuffer.allocateDirect(-initialCapacity.get());
+  }
+
+  public CodecBuffer getFromDb() {
+    for (prepare(); ; allocate()) {
+      final Integer required = buffer.putFromSource(source);
+      if (required == null) {
+        return null; // the source is unavailable
+      } else if (required == buffer.readableBytes()) {
+        return buffer; // buffer size is big enough
+      }
+      // buffer size too small, try increasing the capacity.
+      if (buffer.setCapacity(required)) {
+        buffer.clear();
+        // retry with the new capacity
+        final int retried = buffer.putFromSource(source);
+        Preconditions.assertSame(required.intValue(), retried, "required");
+        return buffer;
+      }
+
+      // failed to increase the capacity
+      // increase initial capacity and reallocate it
+      initialCapacity.increase(required);
+    }
+  }
+}

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

@@ -226,7 +226,7 @@ public RDBStore build() throws IOException {
       return new RDBStore(dbFile, rocksDBOption, statistics, writeOptions, tableConfigs,
           openReadOnly, dbJmxBeanNameName, enableCompactionDag,
           maxDbUpdatesSizeThreshold, createCheckpointDirs, configuration,
-          enableRocksDbMetrics);
+          enableRocksDbMetrics, rocksDBConfiguration.isThreadSafeIteratorEnabled());
     } finally {
       tableConfigs.forEach(TableConfig::close);
     }

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

@@ -73,6 +73,7 @@ public class RDBStore implements DBStore {
   private final long maxDbUpdatesSizeThreshold;
   private final ManagedDBOptions dbOptions;
   private final ManagedStatistics statistics;
+  private final boolean initializeReferenceCountedIterator;
 
   @SuppressWarnings("parameternumber")
   RDBStore(File dbFile, ManagedDBOptions dbOptions, ManagedStatistics statistics,
@@ -82,13 +83,14 @@ public class RDBStore implements DBStore {
                   long maxDbUpdatesSizeThreshold,
                   boolean createCheckpointDirs,
                   ConfigurationSource configuration,
-                  boolean enableRocksDBMetrics)
+                  boolean enableRocksDBMetrics, boolean initializeReferenceCountedIterator)
 
       throws IOException {
     Preconditions.checkNotNull(dbFile, "DB file location cannot be null");
     Preconditions.checkNotNull(families);
     Preconditions.checkArgument(!families.isEmpty());
     this.maxDbUpdatesSizeThreshold = maxDbUpdatesSizeThreshold;
+    this.initializeReferenceCountedIterator = initializeReferenceCountedIterator;
     dbLocation = dbFile;
     this.dbOptions = dbOptions;
     this.statistics = statistics;
@@ -308,7 +310,7 @@ public RDBTable getTable(String name) throws IOException {
     if (handle == null) {
       throw new IOException("No such table in this DB. TableName : " + name);
     }
-    return new RDBTable(this.db, handle, rdbMetrics);
+    return new RDBTable(this.db, handle, rdbMetrics, initializeReferenceCountedIterator);
   }
 
   @Override
@@ -321,7 +323,7 @@ public <K, V> TypedTable<K, V> getTable(
   public ArrayList<Table> listTables() {
     ArrayList<Table> returnList = new ArrayList<>();
     for (ColumnFamily family : getColumnFamilies()) {
-      returnList.add(new RDBTable(db, family, rdbMetrics));
+      returnList.add(new RDBTable(db, family, rdbMetrics, initializeReferenceCountedIterator));
     }
     return returnList;
   }

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

@@ -20,6 +20,7 @@
 import java.io.IOException;
 import java.util.NoSuchElementException;
 import java.util.function.Consumer;
+import org.apache.hadoop.hdds.utils.db.iterator.BaseDBTableIterator;
 import org.apache.hadoop.hdds.utils.db.managed.ManagedRocksIterator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -30,14 +31,14 @@
  * @param <RAW> the raw type.
  */
 abstract class RDBStoreAbstractIterator<RAW>
-    implements TableIterator<RAW, Table.KeyValue<RAW, RAW>> {
+    implements BaseDBTableIterator<RAW, 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 RawKeyValue<RAW> currentEntry;
   // This is for schemas that use a fixed-length
   // prefix for each key.
   private final RAW prefix;
@@ -53,7 +54,7 @@ abstract class RDBStoreAbstractIterator<RAW>
   abstract RAW key();
 
   /** @return the {@link Table.KeyValue} for the current entry. */
-  abstract Table.KeyValue<RAW, RAW> getKeyValue();
+  abstract RawKeyValue<RAW> getKeyValue();
 
   /** Seek to the given key. */
   abstract void seek0(RAW key);
@@ -78,7 +79,7 @@ final RAW getPrefix() {
 
   @Override
   public final void forEachRemaining(
-      Consumer<? super Table.KeyValue<RAW, RAW>> action) {
+      Consumer<? super RawKeyValue<RAW>> action) {
     while (hasNext()) {
       action.accept(next());
     }
@@ -99,7 +100,7 @@ public final boolean hasNext() {
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> next() {
+  public final RawKeyValue<RAW> next() {
     setCurrentEntry();
     if (currentEntry != null) {
       rocksDBIterator.get().next();
@@ -129,7 +130,7 @@ public final void seekToLast() {
   }
 
   @Override
-  public final Table.KeyValue<RAW, RAW> seek(RAW key) {
+  public final RawKeyValue<RAW> seek(RAW key) {
     seek0(key);
     setCurrentEntry();
     return currentEntry;
@@ -151,4 +152,9 @@ public final void removeFromDB() throws IOException {
   public void close() {
     rocksDBIterator.close();
   }
+
+  @Override
+  public boolean isKVCloseable() {
+    return false;
+  }
 }

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

@@ -38,7 +38,7 @@ byte[] key() {
   }
 
   @Override
-  Table.KeyValue<byte[], byte[]> getKeyValue() {
+  RawKeyValue<byte[]> getKeyValue() {
     final ManagedRocksIterator i = getRocksDBIterator();
     return RawKeyValue.create(i.get().key(), i.get().value());
   }

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

@@ -57,9 +57,9 @@ CodecBuffer key() {
   }
 
   @Override
-  Table.KeyValue<CodecBuffer, CodecBuffer> getKeyValue() {
+  RawKeyValue<CodecBuffer> getKeyValue() {
     assertOpen();
-    return Table.newKeyValue(key(), valueBuffer.getFromDb());
+    return new RawKeyValue(key(), valueBuffer.getFromDb());
   }
 
   @Override
@@ -96,60 +96,4 @@ public void close() {
       valueBuffer.release();
     }
   }
-
-  static class Buffer {
-    private final CodecBuffer.Capacity initialCapacity;
-    private final PutToByteBuffer<RuntimeException> source;
-    private CodecBuffer buffer;
-
-    Buffer(CodecBuffer.Capacity initialCapacity,
-           PutToByteBuffer<RuntimeException> source) {
-      this.initialCapacity = initialCapacity;
-      this.source = source;
-    }
-
-    void release() {
-      if (buffer != null) {
-        buffer.release();
-      }
-    }
-
-    private void prepare() {
-      if (buffer == null) {
-        allocate();
-      } else {
-        buffer.clear();
-      }
-    }
-
-    private void allocate() {
-      if (buffer != null) {
-        buffer.release();
-      }
-      buffer = CodecBuffer.allocateDirect(-initialCapacity.get());
-    }
-
-    CodecBuffer getFromDb() {
-      for (prepare(); ; allocate()) {
-        final Integer required = buffer.putFromSource(source);
-        if (required == null) {
-          return null; // the source is unavailable
-        } else if (required == buffer.readableBytes()) {
-          return buffer; // buffer size is big enough
-        }
-        // buffer size too small, try increasing the capacity.
-        if (buffer.setCapacity(required)) {
-          buffer.clear();
-          // retry with the new capacity
-          final int retried = buffer.putFromSource(source);
-          Preconditions.assertSame(required.intValue(), retried, "required");
-          return buffer;
-        }
-
-        // failed to increase the capacity
-        // increase initial capacity and reallocate it
-        initialCapacity.increase(required);
-      }
-    }
-  }
 }