HDDS-10156. Optimize Snapshot Cache get and eviction

hadoop-hdds/common/src/main/resources/ozone-default.xml

@@ -4302,6 +4302,15 @@
   </property>
 
   <property>
+    <name>ozone.om.snapshot.cache.cleanup.service.run.internal</name>
+    <value>1m</value>
+    <tag>OZONE, OM</tag>
+    <description>
+      Interval at which snapshot cache clean up will run.
+      Uses millisecond by default when no time unit is specified.
+    </description>
+  </property>
+
     <name>ozone.om.snapshot.load.native.lib</name>
     <value>true</value>
     <tag>OZONE, OM</tag>

hadoop-hdds/rocksdb-checkpoint-differ/src/main/java/org/apache/ozone/rocksdiff/RocksDBCheckpointDiffer.java

@@ -301,7 +301,7 @@ private void createSstBackUpDir() {
   }
 
   @Override
-  public void close() throws Exception {
+  public void close() {
     if (!closed) {
       synchronized (this) {
         if (!closed) {

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

@@ -570,12 +570,19 @@ private OMConfigKeys() {
 
   public static final String OZONE_OM_SNAPSHOT_DIFF_CLEANUP_SERVICE_RUN_INTERVAL
       = "ozone.om.snapshot.diff.cleanup.service.run.internal";
+  public static final String
+      OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL
+      = "ozone.om.snapshot.cache.cleanup.service.run.internal";
   public static final long
       OZONE_OM_SNAPSHOT_DIFF_CLEANUP_SERVICE_RUN_INTERVAL_DEFAULT
       = TimeUnit.MINUTES.toMillis(1);
+  public static final long
+      OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL_DEFAULT
+      = TimeUnit.MINUTES.toMillis(1);
 
   public static final String OZONE_OM_SNAPSHOT_DIFF_CLEANUP_SERVICE_TIMEOUT
       = "ozone.om.snapshot.diff.cleanup.service.timeout";
+
   public static final long
       OZONE_OM_SNAPSHOT_DIFF_CLEANUP_SERVICE_TIMEOUT_DEFAULT
       = TimeUnit.MINUTES.toMillis(5);

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

@@ -47,6 +47,7 @@
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
+import java.util.UUID;
 import java.util.stream.Collectors;
 
 /**
@@ -74,6 +75,7 @@ public class OmSnapshot implements IOmMetadataReader, Closeable {
   private final String volumeName;
   private final String bucketName;
   private final String snapshotName;
+  private final UUID snapshotID;
   // To access snapshot checkpoint DB metadata
   private final OMMetadataManager omMetadataManager;
   private final KeyManager keyManager;
@@ -83,7 +85,8 @@ public OmSnapshot(KeyManager keyManager,
                     OzoneManager ozoneManager,
                     String volumeName,
                     String bucketName,
-                    String snapshotName) {
+                    String snapshotName,
+                    UUID snapshotID) {
     IAccessAuthorizer accessAuthorizer =
         OzoneAuthorizerFactory.forSnapshot(ozoneManager,
             keyManager, prefixManager);
@@ -93,6 +96,7 @@ public OmSnapshot(KeyManager keyManager,
     this.snapshotName = snapshotName;
     this.bucketName = bucketName;
     this.volumeName = volumeName;
+    this.snapshotID = snapshotID;
     this.keyManager = keyManager;
     this.omMetadataManager = keyManager.getMetadataManager();
   }
@@ -295,6 +299,10 @@ public String getName() {
     return snapshotName;
   }
 
+  public UUID getSnapshotID() {
+    return snapshotID;
+  }
+
   @Override
   public void close() throws IOException {
     // Close DB

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

@@ -83,6 +83,8 @@
 import static org.apache.hadoop.ozone.OzoneConsts.OM_SNAPSHOT_CHECKPOINT_DIR;
 import static org.apache.hadoop.ozone.OzoneConsts.OM_SNAPSHOT_DIFF_DB_NAME;
 import static org.apache.hadoop.ozone.OzoneConsts.OM_SNAPSHOT_INDICATOR;
+import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL;
+import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL_DEFAULT;
 import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_SNAPSHOT_CACHE_MAX_SIZE;
 import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_SNAPSHOT_CACHE_MAX_SIZE_DEFAULT;
 import static org.apache.hadoop.ozone.om.OMConfigKeys.OZONE_OM_SNAPSHOT_DB_MAX_OPEN_FILES;
@@ -272,7 +274,12 @@ public OmSnapshotManager(OzoneManager ozoneManager) {
     };
 
     // Init snapshot cache
-    this.snapshotCache = new SnapshotCache(loader, softCacheSize, ozoneManager.getMetrics());
+    long cacheCleanupServiceInterval = ozoneManager.getConfiguration()
+        .getTimeDuration(OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL,
+            OZONE_OM_SNAPSHOT_CACHE_CLEANUP_SERVICE_RUN_INTERVAL_DEFAULT,
+            TimeUnit.MILLISECONDS);
+    this.snapshotCache = new SnapshotCache(loader, softCacheSize, ozoneManager.getMetrics(),
+        cacheCleanupServiceInterval);
 
     this.snapshotDiffManager = new SnapshotDiffManager(snapshotDiffDb, differ,
         ozoneManager, snapDiffJobCf, snapDiffReportCf,
@@ -382,7 +389,8 @@ public OmSnapshot load(@Nonnull UUID snapshotId) throws IOException {
           return new OmSnapshot(km, pm, ozoneManager,
               snapshotInfo.getVolumeName(),
               snapshotInfo.getBucketName(),
-              snapshotInfo.getName());
+              snapshotInfo.getName(),
+              snapshotInfo.getSnapshotId());
         } catch (Exception e) {
           // Close RocksDB if there is any failure.
           if (!snapshotMetadataManager.getStore().isClosed()) {
@@ -418,6 +426,7 @@ public int getSnapshotCacheSize() {
   public void invalidateCache() {
     if (snapshotCache != null) {
       snapshotCache.invalidateAll();
+      snapshotCache.close();
     }
   }
 

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

@@ -51,10 +51,10 @@ public class ReferenceCounted<T>
   /**
    * Parent instance whose callback will be triggered upon this RC closure.
    */
-  private final Object parentWithCallback;
+  private final ReferenceCountedCallback parentWithCallback;
 
   public ReferenceCounted(T obj, boolean disableCounter,
-      Object parentWithCallback) {
+      ReferenceCountedCallback parentWithCallback) {
     // A param to allow disabling ref counting to reduce active DB
     //  access penalties due to AtomicLong operations.
     this.obj = obj;
@@ -126,7 +126,9 @@ public long decrementRefCount() {
       Preconditions.checkState(newValTotal >= 0L,
           "Total reference count underflow");
     }
-
+    if (refCount.get() == 0) {
+      this.parentWithCallback.callback(this);
+    }
     return refCount.get();
   }
 

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

@@ -0,0 +1,25 @@
+/*
+ * 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.ozone.om.snapshot;
+
+/**
+ * Callback interface for ReferenceCounted.
+ */
+public interface ReferenceCountedCallback {
+  void callback(ReferenceCounted referenceCounted);
+}

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

@@ -20,6 +20,7 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.cache.CacheLoader;
 import org.apache.hadoop.ozone.om.OMMetrics;
+import org.apache.hadoop.hdds.utils.Scheduler;
 import org.apache.hadoop.ozone.om.OmSnapshot;
 import org.apache.hadoop.ozone.om.exceptions.OMException;
 import org.slf4j.Logger;
@@ -28,15 +29,17 @@
 import java.io.IOException;
 import java.util.Iterator;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.TimeUnit;
 
 import static org.apache.hadoop.ozone.om.exceptions.OMException.ResultCodes.FILE_NOT_FOUND;
 
 /**
  * Thread-safe custom unbounded LRU cache to manage open snapshot DB instances.
  */
-public class SnapshotCache {
+public class SnapshotCache implements ReferenceCountedCallback, AutoCloseable {
 
   static final Logger LOG = LoggerFactory.getLogger(SnapshotCache.class);
 
@@ -51,14 +54,28 @@ public class SnapshotCache {
   // Soft-limit of the total number of snapshot DB instances allowed to be
   // opened on the OM.
   private final int cacheSizeLimit;
+  private final Set<UUID> pendingEvictionQueue;
+  private final Scheduler scheduler;
+  private static final String SNAPSHOT_CACHE_CLEANUP_SERVICE =
+      "SnapshotCacheCleanupService";
 
   private final OMMetrics omMetrics;
 
-  public SnapshotCache(CacheLoader<UUID, OmSnapshot> cacheLoader, int cacheSizeLimit, OMMetrics omMetrics) {
+  public SnapshotCache(CacheLoader<UUID, OmSnapshot> cacheLoader, int cacheSizeLimit, OMMetrics omMetrics,
+                       long cleanupInterval) {
     this.dbMap = new ConcurrentHashMap<>();
     this.cacheLoader = cacheLoader;
     this.cacheSizeLimit = cacheSizeLimit;
     this.omMetrics = omMetrics;
+    this.pendingEvictionQueue = ConcurrentHashMap.newKeySet();
+    if (cleanupInterval > 0) {
+      this.scheduler = new Scheduler(SNAPSHOT_CACHE_CLEANUP_SERVICE,
+          true, 1);
+      this.scheduler.scheduleWithFixedDelay(this::cleanup, cleanupInterval,
+          cleanupInterval, TimeUnit.MILLISECONDS);
+    } else {
+      this.scheduler = null;
+    }
   }
 
   @VisibleForTesting
@@ -113,6 +130,13 @@ public void invalidateAll() {
     }
   }
 
+  @Override
+  public void close() {
+    if (this.scheduler != null) {
+      this.scheduler.close();
+    }
+  }
+
   /**
    * State the reason the current thread is getting the OmSnapshot instance.
    * Unused for now.
@@ -170,13 +194,6 @@ public ReferenceCounted<OmSnapshot> get(UUID key) throws IOException {
       throw new OMException("SnapshotId: '" + key + "' not found, or the snapshot is no longer active.",
           OMException.ResultCodes.FILE_NOT_FOUND);
     }
-
-    // Check if any entries can be cleaned up.
-    // At this point, cache size might temporarily exceed cacheSizeLimit
-    // even if there are entries that can be evicted, which is fine since it
-    // is a soft limit.
-    cleanup();
-
     return rcOmSnapshot;
   }
 
@@ -185,58 +202,58 @@ public ReferenceCounted<OmSnapshot> get(UUID key) throws IOException {
    * @param key SnapshotId
    */
   public void release(UUID key) {
-    dbMap.compute(key, (k, v) -> {
-      if (v == null) {
-        throw new IllegalArgumentException("SnapshotId '" + key + "' does not exist in cache.");
-      } else {
-        v.decrementRefCount();
-      }
-      return v;
-    });
-
-    // The cache size might have already exceeded the soft limit
-    // Thus triggering cleanup() to check and evict if applicable
-    cleanup();
+    ReferenceCounted<OmSnapshot> val = dbMap.get(key);
+    if (val == null) {
+      throw new IllegalArgumentException("Key '" + key + "' does not " +
+          "exist in cache.");
+    }
+    val.decrementRefCount();
   }
 
   /**
-   * Wrapper for cleanupInternal() that is synchronized to prevent multiple
-   * threads from interleaving into the cleanup method.
+   * If cache size exceeds soft limit, attempt to clean up and close the
+     instances that has zero reference count.
    */
-  private synchronized void cleanup() {
+  private void cleanup() {
     if (dbMap.size() > cacheSizeLimit) {
-      cleanupInternal();
+      for (UUID evictionKey : pendingEvictionQueue) {
+        dbMap.compute(evictionKey, (k, v) -> {
+          pendingEvictionQueue.remove(k);
+          if (v == null) {
+            throw new IllegalStateException("SnapshotId '" + k + "' does not exist in cache. The RocksDB " +
+                "instance of the Snapshot may not be closed properly.");
+          }
+
+          if (v.getTotalRefCount() > 0) {
+            LOG.debug("SnapshotId {} is still being referenced ({}), skipping its clean up.", k, v.getTotalRefCount());
+            return v;
+          } else {
+            LOG.debug("Closing SnapshotId {}. It is not being referenced anymore.", k);
+            // Close the instance, which also closes its DB handle.
+            try {
+              v.get().close();
+            } catch (IOException ex) {
+              throw new IllegalStateException("Error while closing snapshot DB.", ex);
+            }
+            omMetrics.decNumSnapshotCacheSize();
+            return null;
+          }
+        });
+      }
     }
   }
 
   /**
-   * If cache size exceeds soft limit, attempt to clean up and close the
-   * instances that has zero reference count.
-   * TODO: [SNAPSHOT] Add new ozone debug CLI command to trigger this directly.
+   * Callback method used to enqueue or dequeue ReferenceCounted from
+   * pendingEvictionList.
+   * @param referenceCounted ReferenceCounted object
    */
-  private void cleanupInternal() {
-    for (Map.Entry<UUID, ReferenceCounted<OmSnapshot>> entry : dbMap.entrySet()) {
-      dbMap.compute(entry.getKey(), (k, v) -> {
-        if (v == null) {
-          throw new IllegalStateException("SnapshotId '" + k + "' does not exist in cache. The RocksDB " +
-              "instance of the Snapshot may not be closed properly.");
-        }
-
-        if (v.getTotalRefCount() > 0) {
-          LOG.debug("SnapshotId {} is still being referenced ({}), skipping its clean up.", k, v.getTotalRefCount());
-          return v;
-        } else {
-          LOG.debug("Closing SnapshotId {}. It is not being referenced anymore.", k);
-          // Close the instance, which also closes its DB handle.
-          try {
-            v.get().close();
-          } catch (IOException ex) {
-            throw new IllegalStateException("Error while closing snapshot DB.", ex);
-          }
-          omMetrics.decNumSnapshotCacheSize();
-          return null;
-        }
-      });
+  @Override
+  public void callback(ReferenceCounted referenceCounted) {
+    if (referenceCounted.getTotalRefCount() == 0L) {
+      // Reference count reaches zero, add to pendingEvictionList
+      pendingEvictionQueue.add(((OmSnapshot) referenceCounted.get())
+          .getSnapshotID());
     }
   }
 }

hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotCache.java

@@ -78,7 +78,7 @@ static void beforeAll() throws Exception {
   void setUp() {
     // Reset cache for each test case
     omMetrics = OMMetrics.create();
-    snapshotCache = new SnapshotCache(cacheLoader, CACHE_SIZE_LIMIT, omMetrics);
+    snapshotCache = new SnapshotCache(cacheLoader, CACHE_SIZE_LIMIT, omMetrics, 0);
   }
 
   @AfterEach

hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDiffManager.java

@@ -364,7 +364,7 @@ public void init() throws RocksDBException, IOException, ExecutionException {
     omSnapshotManager = mock(OmSnapshotManager.class);
     when(ozoneManager.getOmSnapshotManager()).thenReturn(omSnapshotManager);
     when(omSnapshotManager.isSnapshotStatus(any(), any())).thenReturn(true);
-    SnapshotCache snapshotCache = new SnapshotCache(mockCacheLoader(), 10, omMetrics);
+    SnapshotCache snapshotCache = new SnapshotCache(mockCacheLoader(), 10, omMetrics, 0);
 
     when(omSnapshotManager.getActiveSnapshot(anyString(), anyString(), anyString()))
         .thenAnswer(invocationOnMock -> {