HDDS-12779. Parallelize table Spliterator with multiple multiple iterators across different ranges

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

@@ -50,7 +50,7 @@
  * A buffer used by {@link Codec}
  * for supporting RocksDB direct {@link ByteBuffer} APIs.
  */
-public class CodecBuffer implements UncheckedAutoCloseable {
+public class CodecBuffer implements UncheckedAutoCloseable, Comparable<byte[]> {
   private static final Logger LOG = LoggerFactory.getLogger(CodecBuffer.class);
 
   private static final ByteBufAllocator POOL = PooledByteBufAllocator.DEFAULT;
@@ -376,6 +376,19 @@ public boolean startsWith(CodecBuffer prefix) {
     return buf.slice(buf.readerIndex(), length).equals(prefix.buf);
   }
 
+  public int compareTo(byte[] other) {
+    Objects.requireNonNull(other, "other == null");
+    final int size = Math.min(readableBytes(), other.length);
+    for (int i = 0; i < size; i++) {
+      final int b1 = buf.getByte(buf.readerIndex() + i) & 0xff;
+      final int b2 = other[i] & 0xff;
+      if (b1 != b2) {
+        return b1 - b2;
+      }
+    }
+    return readableBytes() - other.length;
+  }
+
   /** @return an {@link InputStream} reading from this buffer. */
   public InputStream getInputStream() {
     return new ByteBufInputStream(buf.duplicate());

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

@@ -0,0 +1,29 @@
+/*
+ * 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 java.io.IOException;
+import java.util.List;
+import org.rocksdb.LiveFileMetaData;
+
+/**
+ * Base table interface for Rocksdb.
+ */
+public interface BaseRDBTable<KEY, VALUE> extends Table<KEY, VALUE> {
+  List<LiveFileMetaData> getTableSstFiles() throws IOException;
+}

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

@@ -17,17 +17,24 @@
 
 package org.apache.hadoop.hdds.utils.db;
 
+import java.util.Comparator;
+
 /**
  * No-op codec for byte arrays.
  */
 public final class ByteArrayCodec implements Codec<byte[]> {
 
   private static final Codec<byte[]> INSTANCE = new ByteArrayCodec();
+  private static final Comparator<byte[]> COMPARATOR = new ByteWiseComparator();
 
   public static Codec<byte[]> get() {
     return INSTANCE;
   }
 
+  public static Comparator<byte[]> getComparator() {
+    return COMPARATOR;
+  }
+
   private ByteArrayCodec() {
     // singleton
   }
@@ -51,4 +58,18 @@ public byte[] fromPersistedFormat(byte[] bytes) {
   public byte[] copyObject(byte[] bytes) {
     return bytes;
   }
+
+  private static class ByteWiseComparator implements Comparator<byte[]> {
+    @Override
+    public int compare(byte[] o1, byte[] o2) {
+      int length = Math.min(o1.length, o2.length);
+      for (int i = 0; i < length; i++) {
+        int compareValue = Integer.compareUnsigned(Byte.toUnsignedInt(o1[i]), Byte.toUnsignedInt(o2[i]));
+        if (compareValue != 0) {
+          return compareValue;
+        }
+      }
+      return Integer.compare(o1.length, o2.length);
+    }
+  }
 }

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