HDFS-17529. Improve router state store cache update

Author: kokonguyen191

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/router/RBFConfigKeys.java

@@ -218,6 +218,8 @@ public class RBFConfigKeys extends CommonConfigurationKeysPublic {
       FEDERATION_STORE_PREFIX + "driver.class";
   public static final Class<? extends StateStoreDriver>
       FEDERATION_STORE_DRIVER_CLASS_DEFAULT = StateStoreZooKeeperImpl.class;
+  public static final String FEDERATION_STORE_MEMBERSHIP_ASYNC_OVERRIDE_CLASSES =
+      FEDERATION_STORE_PREFIX + "async.override.classes";
 
   public static final String FEDERATION_STORE_CONNECTION_TEST_MS =
       FEDERATION_STORE_PREFIX + "connection.test";

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/store/CachedRecordStore.java

@@ -21,6 +21,12 @@
 import java.util.ArrayList;
 import java.util.LinkedList;
 import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.Future;
+import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.ThreadPoolExecutor;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
@@ -47,7 +53,8 @@ public abstract class CachedRecordStore<R extends BaseRecord>
 
   /** Prevent loading the cache more than once every 500 ms. */
   private static final long MIN_UPDATE_MS = 500;
-
+  /** Should spawn 2 separate threads for overwriting and deleting records or not? */
+  private boolean asyncOverride = false;
 
   /** Cached entries. */
   private List<R> records = new ArrayList<>();
@@ -94,6 +101,10 @@ protected CachedRecordStore(
     this.override = over;
   }
 
+  public void toggleAsyncOverride(boolean flag) {
+    this.asyncOverride = flag;
+  }
+
   /**
    * Check that the cache of the State Store information is available.
    *
@@ -121,7 +132,7 @@ public boolean loadCache(boolean force) throws IOException {
 
         // If we have any expired record, update the State Store
         if (this.override) {
-          overrideExpiredRecords(result);
+          overrideExpiredRecords(result, this.asyncOverride);
         }
       } catch (IOException e) {
         LOG.error("Cannot get \"{}\" records from the State Store",
@@ -168,11 +179,13 @@ private boolean isUpdateTime() {
    * removed.
    *
    * @param query RecordQueryResult containing the data to be inspected.
+   * @param async should spawn threads or not, one for overwriting, one for deleting
    * @throws IOException If the values cannot be updated.
    */
-  public void overrideExpiredRecords(QueryResult<R> query) throws IOException {
+  public void overrideExpiredRecords(QueryResult<R> query, boolean async) throws IOException {
     List<R> commitRecords = new ArrayList<>();
-    List<R> deleteRecords = new ArrayList<>();
+    List<R> toDeleteRecords = new ArrayList<>();
+    List<R> deletedRecords = new ArrayList<>();
     List<R> newRecords = query.getRecords();
     long currentDriverTime = query.getTimestamp();
     if (newRecords == null || currentDriverTime <= 0) {
@@ -182,24 +195,58 @@ public void overrideExpiredRecords(QueryResult<R> query) throws IOException {
     for (R record : newRecords) {
       if (record.shouldBeDeleted(currentDriverTime)) {
         String recordName = StateStoreUtils.getRecordName(record.getClass());
-        if (getDriver().remove(record)) {
-          deleteRecords.add(record);
-          LOG.info("Deleted State Store record {}: {}", recordName, record);
-        } else {
-          LOG.warn("Couldn't delete State Store record {}: {}", recordName,
-              record);
-        }
+        LOG.info("State Store record to delete {}: {}", recordName, record);
+        toDeleteRecords.add(record);
       } else if (!record.isExpired() && record.checkExpired(currentDriverTime)) {
         String recordName = StateStoreUtils.getRecordName(record.getClass());
         LOG.info("Override State Store record {}: {}", recordName, record);
         commitRecords.add(record);
       }
     }
-    if (commitRecords.size() > 0) {
-      getDriver().putAll(commitRecords, true, false);
-    }
-    if (deleteRecords.size() > 0) {
-      newRecords.removeAll(deleteRecords);
+
+    List<Callable<Void>> callables = new ArrayList<>();
+    callables.add(() -> {
+      if (!commitRecords.isEmpty()) {
+        getDriver().putAll(commitRecords, true, false);
+      }
+      return null;
+    });
+
+    callables.add(() -> {
+      if (!toDeleteRecords.isEmpty()) {
+        deletedRecords.addAll(getDriver().removeMultiple(toDeleteRecords));
+      }
+      if (!deletedRecords.isEmpty()) {
+        newRecords.removeAll(deletedRecords);
+      }
+      return null;
+    });
+
+    if (async) {
+      ThreadPoolExecutor executor = new ThreadPoolExecutor(2, 2,
+          0L, TimeUnit.MILLISECONDS, new LinkedBlockingQueue<>());
+      List<Future<Void>> futures = new ArrayList<>();
+      futures.add(executor.submit(callables.get(0)));
+      futures.add(executor.submit(callables.get(1)));
+      try {
+        for (Future<Void> future : futures) {
+          future.get();
+        }
+      } catch (InterruptedException e) {
+        LOG.error("Failed to override expired records.", e);
+        throw new IOException(e);
+      } catch (ExecutionException e) {
+        throw new IOException(e);
+      } finally {
+        executor.shutdown();
+      }
+    } else {
+      try {
+        callables.get(0).call();
+        callables.get(1).call();
+      } catch (Exception e) {
+        throw new IOException(e);
+      }
     }
   }
 
@@ -215,7 +262,7 @@ public void overrideExpiredRecord(R record) throws IOException {
     newRecords.add(record);
     long time = getDriver().getTime();
     QueryResult<R> query = new QueryResult<>(newRecords, time);
-    overrideExpiredRecords(query);
+    overrideExpiredRecords(query, false);
   }
 
   /**

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/store/StateStoreService.java

@@ -251,6 +251,11 @@ private <T extends RecordStore<?>> void addRecordStore(
 
     T recordStore = RecordStore.newInstance(clazz, this.getDriver());
     Class<? extends BaseRecord> recordClass = recordStore.getRecordClass();
+    if (recordStore instanceof CachedRecordStore && conf.getStringCollection(
+            RBFConfigKeys.FEDERATION_STORE_MEMBERSHIP_ASYNC_OVERRIDE_CLASSES)
+        .contains(clazz.getCanonicalName())) {
+      ((CachedRecordStore<?>) recordStore).toggleAsyncOverride(true);
+    }
     this.recordStores.put(recordClass, recordStore);
 
     // Subscribe for cache updates

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/store/driver/StateStoreRecordOperations.java

@@ -127,6 +127,17 @@ <T extends BaseRecord> StateStoreOperationResult putAll(
   @AtMostOnce
   <T extends BaseRecord> boolean remove(T record) throws IOException;
 
+  /**
+   * Remove multiple records.
+   *
+   * @param <T> Record class of the records.
+   * @param records Records to be removed.
+   * @return Records that were removed.
+   * @throws IOException Throws exception if unable to query the data store.
+   */
+  @AtMostOnce
+  <T extends BaseRecord> List<T> removeMultiple(List<T> records) throws IOException;
+
   /**
    * Remove all records of this class from the store.
    *
@@ -152,4 +163,17 @@ <T extends BaseRecord> StateStoreOperationResult putAll(
   <T extends BaseRecord> int remove(Class<T> clazz, Query<T> query)
       throws IOException;
 
+  /**
+   * Remove all records of a specific class that match any query in a list of queries.
+   * Requires the getAll implementation to fetch fresh records on each call.
+   *
+   * @param clazz The class to match the records with.
+   * @param queries Queries (logical OR) to filter what to remove.
+   * @param <T> Record class of the records.
+   * @return Removed records. Not true removed records, but rather queries that deleted records.
+   * @throws IOException Throws exception if unable to query the data store.
+   */
+  @AtMostOnce
+  <T extends BaseRecord> List<T> remove(Class<T> clazz, List<Query<T>> queries)
+      throws IOException;
 }

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/store/driver/impl/StateStoreBaseImpl.java

@@ -86,4 +86,39 @@ public <T extends BaseRecord> boolean remove(T record) throws IOException {
     Class<T> recordClass = (Class<T>)StateStoreUtils.getRecordClass(clazz);
     return remove(recordClass, query) == 1;
   }
+
+  @Override
+  public <T extends BaseRecord> List<T> removeMultiple(List<T> records) throws IOException {
+    assert !records.isEmpty();
+    // Fall back to iterative remove() calls if all records don't share 1 class
+    Class<? extends BaseRecord> expectedClazz = records.get(0).getClass();
+    if (!records.stream().allMatch(x -> x.getClass() == expectedClazz)) {
+      List<T> result = new ArrayList<>();
+      for (T record : records) {
+        if (remove(record)) {
+          result.add(record);
+        }
+      }
+      return result;
+    }
+
+    final List<Query<T>> queries = new ArrayList<>();
+    for (T record: records) {
+      queries.add(new Query<>(record));
+    }
+    @SuppressWarnings("unchecked")
+    Class<T> recordClass = (Class<T>) StateStoreUtils.getRecordClass(expectedClazz);
+    return remove(recordClass, queries);
+  }
+
+  public <T extends BaseRecord> List<T> remove(Class<T> clazz, List<Query<T>> queries)
+      throws IOException {
+    List<T> result = new ArrayList<>();
+    for (Query<T> query : queries) {
+      if (remove(clazz, query) > 0) {
+        result.add(query.getPartial());
+      }
+    }
+    return result;
+  }
 }

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/java/org/apache/hadoop/hdfs/server/federation/store/driver/impl/StateStoreZooKeeperImpl.java

@@ -25,7 +25,11 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.Callable;
 import java.util.concurrent.Future;
 import java.util.concurrent.LinkedBlockingQueue;
@@ -284,53 +288,89 @@ public <T extends BaseRecord> StateStoreOperationResult putAll(
   }
 
   @Override
-  public <T extends BaseRecord> int remove(
-      Class<T> clazz, Query<T> query) throws IOException {
+  public <T extends BaseRecord> List<T> remove(Class<T> clazz, List<Query<T>> queries)
+      throws IOException {
     verifyDriverReady();
-    if (query == null) {
-      return 0;
+    final List<T> removed = new ArrayList<>();
+    final List<T> trueRemoved = Collections.synchronizedList(new ArrayList<>());
+    if (queries.isEmpty()) {
+      return removed;
     }
 
     // Read the current data
     long start = monotonicNow();
-    List<T> records = null;
+    List<T> records;
     try {
       QueryResult<T> result = get(clazz);
       records = result.getRecords();
     } catch (IOException ex) {
       LOG.error("Cannot get existing records", ex);
       getMetrics().addFailure(monotonicNow() - start);
-      return 0;
+      return removed;
     }
 
     // Check the records to remove
     String znode = getZNodeForClass(clazz);
-    List<T> recordsToRemove = filterMultiple(query, records);
+    Set<T> recordsToRemove = new HashSet<>();
+    Map<Query<T>, List<T>> queryToRecords = new HashMap<>();
+    for (Query<T> query : queries) {
+      List<T> filtered = filterMultiple(query, records);
+      queryToRecords.put(query, filtered);
+      recordsToRemove.addAll(filtered);
+    }
 
     // Remove the records
-    int removed = 0;
-    for (T existingRecord : recordsToRemove) {
+    List<Callable<Void>> callables = new ArrayList<>();
+    recordsToRemove.forEach(existingRecord -> callables.add(() -> {
       LOG.info("Removing \"{}\"", existingRecord);
       try {
         String primaryKey = getPrimaryKey(existingRecord);
         String path = getNodePath(znode, primaryKey);
         if (zkManager.delete(path)) {
-          removed++;
+          trueRemoved.add(existingRecord);
         } else {
           LOG.error("Did not remove \"{}\"", existingRecord);
         }
       } catch (Exception e) {
         LOG.error("Cannot remove \"{}\"", existingRecord, e);
         getMetrics().addFailure(monotonicNow() - start);
       }
+      return null;
+    }));
+    try {
+      if (enableConcurrent) {
+        executorService.invokeAll(callables);
+      } else {
+        for (Callable<Void> callable : callables) {
+          callable.call();
+        }
+      }
+    } catch (Exception e) {
+      LOG.error("Record removal failed : {}", e.getMessage(), e);
+      throw new IOException(e);
     }
     long end = monotonicNow();
-    if (removed > 0) {
+    if (!trueRemoved.isEmpty()) {
       getMetrics().addRemove(end - start);
     }
+    // Convert back to partials
+    for (Map.Entry<Query<T>, List<T>> entry : queryToRecords.entrySet()) {
+      for (T record : entry.getValue()) {
+        if (trueRemoved.contains(record)) {
+          removed.add(entry.getKey().getPartial());
+          break;
+        }
+      }
+    }
     return removed;
   }
 
+  @Override
+  public <T extends BaseRecord> int remove(Class<T> clazz, Query<T> query)
+      throws IOException {
+    return remove(clazz, Collections.singletonList(query)).size();
+  }
+
   @Override
   public <T extends BaseRecord> boolean removeAll(Class<T> clazz)
       throws IOException {

hadoop-hdfs-project/hadoop-hdfs-rbf/src/main/resources/hdfs-rbf-default.xml

@@ -370,6 +370,20 @@
     </description>
   </property>
 
+  <property>
+    <name>dfs.federation.router.store.async.override.classes</name>
+    <value></value>
+    <description>
+      CachedRecordStore periodic updates contain an overwrite part and a deletion part.
+      Classes in this config will do these 2 steps in parallel when updating cache.
+      Supported classes:
+      org.apache.hadoop.hdfs.server.federation.store.impl.DisabledNameserviceStoreImpl
+      org.apache.hadoop.hdfs.server.federation.store.impl.MembershipStoreImpl
+      org.apache.hadoop.hdfs.server.federation.store.impl.MountTableStoreImpl
+      org.apache.hadoop.hdfs.server.federation.store.impl.RouterStoreImpl
+    </description>
+  </property>
+
   <property>
     <name>dfs.federation.router.store.connection.test</name>
     <value>60000</value>

hadoop-hdfs-project/hadoop-hdfs-rbf/src/site/markdown/HDFSRouterFederation.md

@@ -469,6 +469,7 @@ The connection to the State Store and the internal caching at the Router.
 | dfs.federation.router.store.connection.test | 60000 | How often to check for the connection to the State Store in milliseconds. |
 | dfs.federation.router.cache.ttl | 60000 | How often to refresh the State Store caches in milliseconds. |
 | dfs.federation.router.store.membership.expiration | 300000 | Expiration time in milliseconds for a membership record. |
+| dfs.federation.router.store.async.override.classes |  | These cache State Stores, when updating cache, will overwrite and delete records in parallel. |
 | dfs.federation.router.mount-table.cache.update | false | If true, Mount table cache is updated whenever a mount table entry is added, modified or removed for all the routers. |
 | dfs.federation.router.mount-table.cache.update.timeout | 1m | Max time to wait for all the routers to finish their mount table cache update. |
 | dfs.federation.router.mount-table.cache.update.client.max.time | 5m | Max time a RouterClient connection can be cached. |