HDDS-1829 On OM reload/restart OmMetrics#numKeys should be updated (#1164)

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

@@ -183,4 +183,14 @@ public String getName() throws IOException {
   public void close() throws Exception {
     // Nothing do for a Column Family.
   }
+
+  @Override
+  public long getEstimatedKeyCount() throws IOException {
+    try {
+      return db.getLongProperty(handle, "rocksdb.estimate-num-keys");
+    } catch (RocksDBException e) {
+      throw toIOException(
+          "Failed to get estimated key count of table " + getName(), e);
+    }
+  }
 }

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

@@ -111,6 +111,13 @@ void putWithBatch(BatchOperation batch, KEY key, VALUE value)
    */
   String getName() throws IOException;
 
+  /**
+   * Returns the key count of this Table.  Note the result can be inaccurate.
+   * @return Estimated key count of this Table
+   * @throws IOException on failure
+   */
+  long getEstimatedKeyCount() throws IOException;
+
   /**
    * Add entry to the table cache.
    *

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

@@ -205,6 +205,11 @@ public String getName() throws IOException {
     return rawTable.getName();
   }
 
+  @Override
+  public long getEstimatedKeyCount() throws IOException {
+    return rawTable.getEstimatedKeyCount();
+  }
+
   @Override
   public void close() throws Exception {
     rawTable.close();

hadoop-hdds/common/src/test/java/org/apache/hadoop/utils/db/TestRDBTableStore.java

@@ -51,7 +51,8 @@ public class TestRDBTableStore {
       Arrays.asList(DFSUtil.bytes2String(RocksDB.DEFAULT_COLUMN_FAMILY),
           "First", "Second", "Third",
           "Fourth", "Fifth",
-          "Sixth", "Seventh");
+          "Sixth", "Seventh",
+          "Eighth");
   @Rule
   public TemporaryFolder folder = new TemporaryFolder();
   private RDBStore rdbStore = null;
@@ -247,4 +248,22 @@ public void testIsExist() throws Exception {
       Assert.assertFalse(testTable.isExist(invalidKey));
     }
   }
+
+  @Test
+  public void testCountEstimatedRowsInTable() throws Exception {
+    try (Table<byte[], byte[]> testTable = rdbStore.getTable("Eighth")) {
+      // Add a few keys
+      final int numKeys = 12345;
+      for (int i = 0; i < numKeys; i++) {
+        byte[] key =
+            RandomStringUtils.random(10).getBytes(StandardCharsets.UTF_8);
+        byte[] value =
+            RandomStringUtils.random(10).getBytes(StandardCharsets.UTF_8);
+        testTable.put(key, value);
+      }
+      long count = testTable.getEstimatedKeyCount();
+      // The result should be larger than zero but not exceed(?) numKeys
+      Assert.assertTrue(count > 0 && count <= numKeys);
+    }
+  }
 }

hadoop-hdds/common/src/test/java/org/apache/hadoop/utils/db/TestTypedRDBTableStore.java

@@ -55,7 +55,8 @@ public class TestTypedRDBTableStore {
       Arrays.asList(DFSUtil.bytes2String(RocksDB.DEFAULT_COLUMN_FAMILY),
           "First", "Second", "Third",
           "Fourth", "Fifth",
-          "Sixth", "Seven", "Eighth");
+          "Sixth", "Seven", "Eighth",
+          "Ninth");
   @Rule
   public TemporaryFolder folder = new TemporaryFolder();
   private RDBStore rdbStore = null;
@@ -351,4 +352,22 @@ public void testIsExistCache() throws Exception {
       Assert.assertFalse(testTable.isExist(key));
     }
   }
+
+  @Test
+  public void testCountEstimatedRowsInTable() throws Exception {
+    try (Table<String, String> testTable = createTypedTable(
+        "Ninth")) {
+      // Add a few keys
+      final int numKeys = 12345;
+      for (int i = 0; i < numKeys; i++) {
+        String key =
+            RandomStringUtils.random(10);
+        String value = RandomStringUtils.random(10);
+        testTable.put(key, value);
+      }
+      long count = testTable.getEstimatedKeyCount();
+      // The result should be larger than zero but not exceed(?) numKeys
+      Assert.assertTrue(count > 0 && count <= numKeys);
+    }
+  }
 }

hadoop-ozone/common/src/main/java/org/apache/hadoop/ozone/om/OMMetadataManager.java

@@ -316,4 +316,15 @@ String getMultipartKey(String volume, String bucket, String key, String
    */
   <KEY, VALUE> long countRowsInTable(Table<KEY, VALUE> table)
       throws IOException;
+
+  /**
+   * Returns an estimated number of rows in a table.  This is much quicker
+   * than {@link OMMetadataManager#countRowsInTable} but the result can be
+   * inaccurate.
+   * @param table Table
+   * @return long Estimated number of rows in the table.
+   * @throws IOException
+   */
+  <KEY, VALUE> long countEstimatedRowsInTable(Table<KEY, VALUE> table)
+      throws IOException;
 }

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmMetadataManagerImpl.java

@@ -802,6 +802,16 @@ public <KEY, VALUE> long countRowsInTable(Table<KEY, VALUE> table)
     return count;
   }
 
+  @Override
+  public <KEY, VALUE> long countEstimatedRowsInTable(Table<KEY, VALUE> table)
+      throws IOException {
+    long count = 0;
+    if (table != null) {
+      count = table.getEstimatedKeyCount();
+    }
+    return count;
+  }
+
   @Override
   public Table<String, S3SecretValue> getS3SecretTable() {
     return s3SecretTable;

hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OzoneManager.java

@@ -3267,6 +3267,8 @@ void reloadOMState(long newSnapshotIndex) throws IOException {
         .getVolumeTable()));
     metrics.setNumBuckets(metadataManager.countRowsInTable(metadataManager
         .getBucketTable()));
+    metrics.setNumKeys(metadataManager.countEstimatedRowsInTable(metadataManager
+        .getKeyTable()));
 
     // Delete the omMetrics file if it exists and save the a new metrics file
     // with new data