HDDS-16289. Reduce OM FSO apply-path write-lock hold and path-walk cost

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

@@ -24,12 +24,14 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
 import org.apache.commons.lang3.NotImplementedException;
 import org.apache.hadoop.hdds.annotation.InterfaceStability;
 import org.apache.hadoop.hdds.utils.MetadataKeyFilters.KeyPrefixFilter;
 import org.apache.hadoop.hdds.utils.TableCacheMetrics;
 import org.apache.hadoop.hdds.utils.db.cache.CacheKey;
 import org.apache.hadoop.hdds.utils.db.cache.CacheValue;
+import org.apache.ratis.util.function.CheckedFunction;
 
 /**
  * Interface for key-value store that stores ozone metadata. Ozone metadata is
@@ -127,6 +129,26 @@ default VALUE getReadCopy(KEY key) throws RocksDatabaseException, CodecException
     throw new NotImplementedException("getReadCopy is not implemented");
   }
 
+  /**
+   * Returns a projection of the value mapped to the given key, i.e. only the fields the caller needs,
+   * or null if the key is not found. Implementations may decode {@code fromPersistedValue} straight
+   * from the serialized value instead of materializing the whole VALUE object.
+   * <p>
+   * {@code fromCachedValue} is applied to the cached object itself, without the defensive copy
+   * {@link #get} makes, so it must neither mutate nor retain it.
+   *
+   * @param key metadata key
+   * @param fromCachedValue extracts the projection from a value found in the cache
+   * @param fromPersistedValue extracts the projection from the serialized value read from the store
+   * @return the projection, or null if the key is not found.
+   */
+  default <R> R getProjected(KEY key, Function<VALUE, R> fromCachedValue,
+      CheckedFunction<CodecBuffer, R, CodecException> fromPersistedValue)
+      throws RocksDatabaseException, CodecException {
+    final VALUE value = get(key);
+    return value == null ? null : fromCachedValue.apply(value);
+  }
+
   /**
    * Returns the value mapped to the given key in byte array or returns null
    * if the key is not found.

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

@@ -30,6 +30,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.function.Function;
 import org.apache.hadoop.hdds.utils.IOUtils;
 import org.apache.hadoop.hdds.utils.MetadataKeyFilters.KeyPrefixFilter;
 import org.apache.hadoop.hdds.utils.TableCacheMetrics;
@@ -43,6 +44,7 @@
 import org.apache.hadoop.hdds.utils.db.cache.TableNoCache;
 import org.apache.ratis.util.Preconditions;
 import org.apache.ratis.util.function.CheckedBiFunction;
+import org.apache.ratis.util.function.CheckedFunction;
 import org.rocksdb.ByteBufferGetStatus;
 
 /**
@@ -218,6 +220,34 @@ public VALUE get(KEY key) throws RocksDatabaseException, CodecException {
     }
   }
 
+  /**
+   * Same cache semantics as {@link #get(Object)}, except that a value read from the store is decoded
+   * by {@code fromPersistedValue} directly from the buffer it was read into, so the fields the caller
+   * does not need are never decoded and no VALUE object is built.
+   */
+  @Override
+  public <R> R getProjected(KEY key, Function<VALUE, R> fromCachedValue,
+      CheckedFunction<CodecBuffer, R, CodecException> fromPersistedValue)
+      throws RocksDatabaseException, CodecException {
+    final CacheResult<VALUE> cacheResult = cache.lookup(new CacheKey<>(key));
+
+    if (cacheResult.getCacheStatus() == EXISTS) {
+      return fromCachedValue.apply(cacheResult.getValue().getCacheValue());
+    } else if (cacheResult.getCacheStatus() == NOT_EXIST) {
+      return null;
+    } else if (supportCodecBuffer) {
+      return getFromTable(key, this::getFromTable, fromPersistedValue);
+    } else {
+      final byte[] valueBytes = rawTable.get(encodeKey(key));
+      if (valueBytes == null) {
+        return null;
+      }
+      try (CodecBuffer buffer = CodecBuffer.wrap(valueBytes)) {
+        return fromPersistedValue.apply(buffer);
+      }
+    }
+  }
+
   /**
    * Skip checking cache and get the value mapped to the given key in byte
    * array or returns null if the key is not found.
@@ -466,6 +496,13 @@ private Integer getFromTableIfExist(CodecBuffer key, CodecBuffer outValue) throw
   private VALUE getFromTable(KEY key,
       CheckedBiFunction<CodecBuffer, CodecBuffer, Integer, RocksDatabaseException> get)
       throws RocksDatabaseException, CodecException {
+    return getFromTable(key, get, valueCodec::fromCodecBuffer);
+  }
+
+  private <R> R getFromTable(KEY key,
+      CheckedBiFunction<CodecBuffer, CodecBuffer, Integer, RocksDatabaseException> get,
+      CheckedFunction<CodecBuffer, R, CodecException> decoder)
+      throws RocksDatabaseException, CodecException {
     try (CodecBuffer inKey = keyCodec.toDirectCodecBuffer(key)) {
       for (; ;) {
         final Integer required;
@@ -482,7 +519,7 @@ private VALUE getFromTable(KEY key,
           for (; ;) {
             if (required == outValue.readableBytes()) {
               // buffer size is big enough
-              return valueCodec.fromCodecBuffer(outValue);
+              return decoder.apply(outValue);
             }
             // buffer size too small, try increasing the capacity.
             if (!outValue.setCapacity(required)) {

hadoop-hdds/framework/src/test/java/org/apache/hadoop/hdds/utils/db/TestTypedTable.java

@@ -59,7 +59,7 @@
  */
 public class TestTypedTable {
   private final List<String> families = Arrays.asList(StringUtils.bytes2String(RocksDB.DEFAULT_COLUMN_FAMILY),
-      "First", "Second", "Third", "Fourth", "Fifth", "Sixth", "Seventh", "Eighth");
+      "First", "Second", "Third", "Fourth", "Fifth", "Sixth", "Seventh", "Eighth", "Ninth");
 
   private RDBStore rdb;
   private final List<UncheckedAutoCloseable> closeables = new ArrayList<>();
@@ -231,6 +231,60 @@ void runTestEmptyString(Codec<String> codec) throws Exception {
     runTestSingleKeyValue(nonEmpty, empty, table);
   }
 
+  @Test
+  public void testGetProjectedCodecBuffer() throws Exception {
+    runTestGetProjected(StringCodec.get());
+  }
+
+  @Test
+  public void testGetProjectedByteArray() throws Exception {
+    final Codec<String> codec = CodecTestUtil.newCodecWithoutCodecBuffer(StringCodec.get());
+    assertFalse(codec.supportCodecBuffer());
+    runTestGetProjected(codec);
+  }
+
+  /**
+   * The two projections are tagged differently so that each assertion also pins down which path
+   * produced the value: the table cache, or a decode of the value read from the store.
+   */
+  void runTestGetProjected(Codec<String> valueCodec) throws Exception {
+    final TypedTable<String, String> table = newTypedTable(9, StringCodec.get(), valueCodec);
+
+    // A value larger than the initial buffer capacity, to exercise the capacity-retry loop.
+    final StringBuilder large = new StringBuilder();
+    for (int i = 0; i < TypedTable.BUFFER_SIZE_DEFAULT; i++) {
+      large.append('x');
+    }
+
+    table.put("inDb", "dbValue");
+    table.put("large", large.toString());
+    table.put("tombstoned", "gone");
+    table.addCacheEntry("inCache", "cachedValue", 1L);
+    table.addCacheEntry("tombstoned", 2L);
+
+    assertNull(getProjected(table, "absent"));
+    assertEquals("db:dbValue", getProjected(table, "inDb"));
+    assertEquals("db:" + large, getProjected(table, "large"));
+    assertEquals("cache:cachedValue", getProjected(table, "inCache"));
+    // Deleted in the cache, still in the store: the projection must not fall through to the store.
+    assertNull(getProjected(table, "tombstoned"));
+  }
+
+  @Test
+  public void testGetProjectedInMemoryTable() throws Exception {
+    final Table<String, String> table = new InMemoryTestTable<>();
+    table.put("key", "value");
+
+    // The Table default has no serialized value to project from, so it always uses fromCachedValue.
+    assertEquals("cache:value", getProjected(table, "key"));
+    assertNull(getProjected(table, "absent"));
+  }
+
+  static String getProjected(Table<String, String> table, String key) throws Exception {
+    return table.getProjected(key, value -> "cache:" + value,
+        buffer -> "db:" + StringCodec.get().fromCodecBuffer(buffer));
+  }
+
   @Test
   public void testContainerIDvsLong() throws Exception {
     final Map<Long, ContainerID> keys = newMap(1000, ContainerID::valueOf);