This is an automated email from the ASF dual-hosted git repository.

swamirishi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/ozone.git


The following commit(s) were added to refs/heads/master by this push:
     new cb9867bcd3 HDDS-13065. Refactor SnapshotCache to return 
AutoCloseSupplier instead of ReferenceCounted (#8473)
cb9867bcd3 is described below

commit cb9867bcd317ad0f4c4dfdaeeb737353f7b8bca5
Author: Swaminathan Balachandran <[email protected]>
AuthorDate: Wed May 21 18:47:50 2025 -0400

    HDDS-13065. Refactor SnapshotCache to return AutoCloseSupplier instead of 
ReferenceCounted (#8473)
---
 .../hadoop/ozone/freon/TestOMSnapshotDAG.java      |  8 ++--
 .../snapshot/TestSnapshotBackgroundServices.java   |  5 ++-
 ...TestSnapshotDeletingServiceIntegrationTest.java | 13 +++---
 .../hadoop/ozone/om/OmMetadataManagerImpl.java     |  4 +-
 .../apache/hadoop/ozone/om/OmSnapshotManager.java  | 16 ++++----
 .../org/apache/hadoop/ozone/om/OzoneManager.java   | 48 +++++++++++++---------
 .../hadoop/ozone/om/SstFilteringService.java       |  4 +-
 .../hadoop/ozone/om/request/OMClientRequest.java   |  6 +--
 .../key/OMDirectoriesPurgeResponseWithFSO.java     |  4 +-
 .../ozone/om/response/key/OMKeyPurgeResponse.java  |  4 +-
 .../OMSnapshotMoveDeletedKeysResponse.java         |  8 ++--
 .../snapshot/OMSnapshotMoveTableKeysResponse.java  |  6 +--
 .../ozone/om/service/DirectoryDeletingService.java |  4 +-
 .../ozone/om/service/KeyDeletingService.java       |  8 ++--
 .../ozone/om/service/SnapshotDeletingService.java  |  4 +-
 .../service/SnapshotDirectoryCleaningService.java  |  8 ++--
 .../hadoop/ozone/om/snapshot/ReferenceCounted.java | 12 +-----
 .../hadoop/ozone/om/snapshot/SnapshotCache.java    | 26 +++++++++++-
 .../ozone/om/snapshot/SnapshotDiffManager.java     |  5 ++-
 .../om/snapshot/filter/ReclaimableDirFilter.java   |  4 +-
 .../om/snapshot/filter/ReclaimableFilter.java      | 10 ++---
 .../om/snapshot/filter/ReclaimableKeyFilter.java   |  6 +--
 .../filter/ReclaimableRenameEntryFilter.java       |  4 +-
 .../TestOMDirectoriesPurgeRequestAndResponse.java  |  4 +-
 .../key/TestOMKeyPurgeRequestAndResponse.java      |  4 +-
 .../ozone/om/request/key/TestOMKeyRequest.java     |  6 +--
 .../s3/multipart/TestS3MultipartRequest.java       |  6 +--
 .../TestOMSnapshotMoveTableKeysResponse.java       |  8 ++--
 .../ozone/om/service/TestKeyDeletingService.java   |  4 +-
 .../ozone/om/snapshot/TestSnapshotCache.java       | 42 +++++++++----------
 .../ozone/om/snapshot/TestSnapshotDiffManager.java | 13 +++---
 .../ozone/om/snapshot/TestSstFilteringService.java |  3 +-
 .../filter/AbstractReclaimableFilterTest.java      |  6 +--
 .../snapshot/filter/TestReclaimableDirFilter.java  |  6 +--
 .../om/snapshot/filter/TestReclaimableFilter.java  |  4 +-
 .../snapshot/filter/TestReclaimableKeyFilter.java  |  8 ++--
 .../filter/TestReclaimableRenameEntryFilter.java   |  6 +--
 37 files changed, 184 insertions(+), 153 deletions(-)

diff --git 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/freon/TestOMSnapshotDAG.java
 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/freon/TestOMSnapshotDAG.java
index 332dab4fb4..4a2c030526 100644
--- 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/freon/TestOMSnapshotDAG.java
+++ 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/freon/TestOMSnapshotDAG.java
@@ -57,12 +57,12 @@
 import org.apache.hadoop.ozone.om.helpers.OmBucketInfo;
 import org.apache.hadoop.ozone.om.helpers.OmVolumeArgs;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.ozone.rocksdiff.DifferSnapshotInfo;
 import org.apache.ozone.rocksdiff.RocksDBCheckpointDiffer;
 import org.apache.ozone.test.GenericTestUtils;
 import org.apache.ratis.server.RaftServer;
 import org.apache.ratis.server.raftlog.RaftLog;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -212,9 +212,9 @@ public void testDAGReconstruction()
     OMMetadataManager omMetadataManager = ozoneManager.getMetadataManager();
     RDBStore rdbStore = (RDBStore) omMetadataManager.getStore();
     RocksDBCheckpointDiffer differ = rdbStore.getRocksDBCheckpointDiffer();
-    ReferenceCounted<OmSnapshot> snapDB1 = ozoneManager.getOmSnapshotManager()
+    UncheckedAutoCloseableSupplier<OmSnapshot> snapDB1 = 
ozoneManager.getOmSnapshotManager()
         .getActiveSnapshot(volumeName, bucketName, "snap1");
-    ReferenceCounted<OmSnapshot> snapDB2 = ozoneManager.getOmSnapshotManager()
+    UncheckedAutoCloseableSupplier<OmSnapshot> snapDB2 = 
ozoneManager.getOmSnapshotManager()
         .getActiveSnapshot(volumeName, bucketName, "snap2");
     DifferSnapshotInfo snap1 = getDifferSnapshotInfo(omMetadataManager,
         volumeName, bucketName, "snap1",
@@ -240,7 +240,7 @@ public void testDAGReconstruction()
 
     resp = store.createSnapshot(volumeName, bucketName, "snap3");
     LOG.debug("Snapshot created: {}", resp);
-    ReferenceCounted<OmSnapshot> snapDB3 = ozoneManager.getOmSnapshotManager()
+    UncheckedAutoCloseableSupplier<OmSnapshot> snapDB3 = 
ozoneManager.getOmSnapshotManager()
         .getActiveSnapshot(volumeName, bucketName, "snap3");
     DifferSnapshotInfo snap3 = getDifferSnapshotInfo(omMetadataManager,
         volumeName, bucketName, "snap3",
diff --git 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotBackgroundServices.java
 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotBackgroundServices.java
index 0983ff06a3..1f12697d66 100644
--- 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotBackgroundServices.java
+++ 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotBackgroundServices.java
@@ -79,6 +79,7 @@
 import org.apache.ozone.test.LambdaTestUtils;
 import org.apache.ozone.test.tag.Flaky;
 import org.apache.ratis.server.protocol.TermIndex;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.DisplayName;
@@ -260,7 +261,7 @@ public void testSnapshotAndKeyDeletionBackgroundServices()
 
     // get snapshot c
     OmSnapshot snapC;
-    try (ReferenceCounted<OmSnapshot> rcC = newLeaderOM
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcC = newLeaderOM
         .getOmSnapshotManager()
         .getSnapshot(volumeName, bucketName, snapshotInfoC.getName())) {
       assertNotNull(rcC);
@@ -284,7 +285,7 @@ public void testSnapshotAndKeyDeletionBackgroundServices()
 
     // get snapshot d
     OmSnapshot snapD;
-    try (ReferenceCounted<OmSnapshot> rcD = newLeaderOM
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcD = newLeaderOM
         .getOmSnapshotManager()
         .getSnapshot(volumeName, bucketName, snapshotInfoD.getName())) {
       assertNotNull(rcD);
diff --git 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDeletingServiceIntegrationTest.java
 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDeletingServiceIntegrationTest.java
index dddfd97177..7e8befa6a1 100644
--- 
a/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDeletingServiceIntegrationTest.java
+++ 
b/hadoop-ozone/integration-test/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDeletingServiceIntegrationTest.java
@@ -78,6 +78,7 @@
 import org.apache.hadoop.ozone.om.service.SnapshotDeletingService;
 import org.apache.ozone.test.GenericTestUtils;
 import org.apache.ozone.test.tag.Flaky;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -421,7 +422,7 @@ public void testSnapshotWithFSO() throws Exception {
     assertTableRowCount(renamedTable, 4);
     assertTableRowCount(deletedDirTable, 3);
 
-    ReferenceCounted<OmSnapshot> rcSnap1 =
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcSnap1 =
         om.getOmSnapshotManager().getSnapshot(
             VOLUME_NAME, BUCKET_NAME_FSO, "snap1");
     OmSnapshot snap1 = rcSnap1.get();
@@ -613,8 +614,9 @@ public void 
testParallelExcecutionOfKeyDeletionAndSnapshotDeletion() throws Exce
     List<Table.KeyValue<String, String>> renamesKeyEntries;
     List<Table.KeyValue<String, List<OmKeyInfo>>> deletedKeyEntries;
     List<Table.KeyValue<String, OmKeyInfo>> deletedDirEntries;
-    try (ReferenceCounted<OmSnapshot> snapshot = 
om.getOmSnapshotManager().getSnapshot(testBucket.getVolumeName(),
-        testBucket.getName(), testBucket.getName() + "snap2")) {
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> snapshot =
+             om.getOmSnapshotManager().getSnapshot(testBucket.getVolumeName(), 
testBucket.getName(),
+                 testBucket.getName() + "snap2")) {
       renamesKeyEntries = 
snapshot.get().getKeyManager().getRenamesKeyEntries(testBucket.getVolumeName(),
           testBucket.getName(), "", 1000);
       deletedKeyEntries = 
snapshot.get().getKeyManager().getDeletedKeyEntries(testBucket.getVolumeName(),
@@ -649,8 +651,9 @@ public void 
testParallelExcecutionOfKeyDeletionAndSnapshotDeletion() throws Exce
     Future<?> future = snapshotDeletingThread.submit(snapshotDeletionRunnable);
     GenericTestUtils.waitFor(snapshotDeletionStarted::get, 1000, 30000);
     future.get();
-    try (ReferenceCounted<OmSnapshot> snapshot = 
om.getOmSnapshotManager().getSnapshot(testBucket.getVolumeName(),
-        testBucket.getName(), testBucket.getName() + "snap2")) {
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> snapshot =
+             om.getOmSnapshotManager().getSnapshot(testBucket.getVolumeName(), 
testBucket.getName(),
+                 testBucket.getName() + "snap2")) {
       Assertions.assertEquals(Collections.emptyList(),
           
snapshot.get().getKeyManager().getRenamesKeyEntries(testBucket.getVolumeName(),
           testBucket.getName(), "", 1000));
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmMetadataManagerImpl.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmMetadataManagerImpl.java
index 7a4b6ecdb2..173f18313c 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmMetadataManagerImpl.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmMetadataManagerImpl.java
@@ -113,7 +113,6 @@
 import 
org.apache.hadoop.ozone.om.protocolPB.OzoneManagerProtocolClientSideTranslatorPB;
 import org.apache.hadoop.ozone.om.request.file.OMFileRequest;
 import org.apache.hadoop.ozone.om.request.util.OMMultipartUploadUtils;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.ExpiredMultipartUploadInfo;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.ExpiredMultipartUploadsBucket;
@@ -124,6 +123,7 @@
 import org.apache.hadoop.util.Time;
 import org.apache.ozone.compaction.log.CompactionLogEntry;
 import org.apache.ratis.util.ExitUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.eclipse.jetty.util.StringUtil;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -1351,7 +1351,7 @@ public PendingKeysDeletion getPendingDeletionKeys(final 
int keyCount,
             continue;
           }
           // Get the latest snapshot in snapshot path.
-          try (ReferenceCounted<OmSnapshot> rcLatestSnapshot = 
previousSnapshotInfo == null ? null :
+          try (UncheckedAutoCloseableSupplier<OmSnapshot> rcLatestSnapshot = 
previousSnapshotInfo == null ? null :
               
omSnapshotManager.getSnapshot(previousSnapshotInfo.getVolumeName(),
                   previousSnapshotInfo.getBucketName(), 
previousSnapshotInfo.getName())) {
 
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmSnapshotManager.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmSnapshotManager.java
index 7b6b9bab4f..e6f893ee72 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmSnapshotManager.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OmSnapshotManager.java
@@ -92,7 +92,6 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotDiffJob;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.service.SnapshotDiffCleanupService;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotCache;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotDiffManager;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
@@ -102,6 +101,7 @@
 import org.apache.hadoop.ozone.snapshot.SnapshotDiffResponse;
 import org.apache.ozone.rocksdiff.RocksDBCheckpointDiffer;
 import org.apache.ratis.util.function.CheckedFunction;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.rocksdb.ColumnFamilyDescriptor;
 import org.rocksdb.ColumnFamilyHandle;
 import org.rocksdb.RocksDBException;
@@ -604,7 +604,7 @@ private static void 
deleteKeysFromDelKeyTableInSnapshotScope(
 
   // Get OmSnapshot if the keyName has ".snapshot" key indicator
   @SuppressWarnings("unchecked")
-  public ReferenceCounted<IOmMetadataReader> getActiveFsMetadataOrSnapshot(
+  public UncheckedAutoCloseableSupplier<IOmMetadataReader> 
getActiveFsMetadataOrSnapshot(
       String volumeName,
       String bucketName,
       String keyName) throws IOException {
@@ -622,28 +622,28 @@ public ReferenceCounted<IOmMetadataReader> 
getActiveFsMetadataOrSnapshot(
           bucketName), false, false);
       volumeName = resolvedBucket.realVolume();
       bucketName = resolvedBucket.realBucket();
-      return (ReferenceCounted<IOmMetadataReader>) (ReferenceCounted<?>)
+      return (UncheckedAutoCloseableSupplier<IOmMetadataReader>) 
(UncheckedAutoCloseableSupplier<?>)
           getActiveSnapshot(volumeName, bucketName, snapshotName);
     } else {
       return ozoneManager.getOmMetadataReader();
     }
   }
 
-  public ReferenceCounted<OmSnapshot> getActiveSnapshot(
+  public UncheckedAutoCloseableSupplier<OmSnapshot> getActiveSnapshot(
       String volumeName,
       String bucketName,
       String snapshotName) throws IOException {
     return getSnapshot(volumeName, bucketName, snapshotName, false);
   }
 
-  public ReferenceCounted<OmSnapshot> getSnapshot(
+  public UncheckedAutoCloseableSupplier<OmSnapshot> getSnapshot(
       String volumeName,
       String bucketName,
       String snapshotName) throws IOException {
     return getSnapshot(volumeName, bucketName, snapshotName, true);
   }
 
-  private ReferenceCounted<OmSnapshot> getSnapshot(
+  private UncheckedAutoCloseableSupplier<OmSnapshot> getSnapshot(
       String volumeName,
       String bucketName,
       String snapshotName,
@@ -659,7 +659,7 @@ private ReferenceCounted<OmSnapshot> getSnapshot(
     return getSnapshot(snapshotTableKey, skipActiveCheck);
   }
 
-  private ReferenceCounted<OmSnapshot> getSnapshot(String snapshotTableKey, 
boolean skipActiveCheck)
+  private UncheckedAutoCloseableSupplier<OmSnapshot> getSnapshot(String 
snapshotTableKey, boolean skipActiveCheck)
       throws IOException {
     SnapshotInfo snapshotInfo = SnapshotUtils.getSnapshotInfo(ozoneManager, 
snapshotTableKey);
     // Block FS API reads when snapshot is not active.
@@ -710,7 +710,7 @@ public static boolean 
areSnapshotChangesFlushedToDB(OMMetadataManager metadataMa
    * This should only be used for API calls initiated by background service 
e.g. purgeKeys, purgeSnapshot,
    * snapshotMoveDeletedKeys, and SetSnapshotProperty.
    */
-  public ReferenceCounted<OmSnapshot> getSnapshot(UUID snapshotId) throws 
IOException {
+  public UncheckedAutoCloseableSupplier<OmSnapshot> getSnapshot(UUID 
snapshotId) throws IOException {
     return snapshotCache.get(snapshotId);
   }
 
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OzoneManager.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OzoneManager.java
index 526145d169..ad1741df0a 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OzoneManager.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/OzoneManager.java
@@ -283,7 +283,6 @@
 import org.apache.hadoop.ozone.om.service.OMRangerBGSyncService;
 import org.apache.hadoop.ozone.om.service.QuotaRepairTask;
 import org.apache.hadoop.ozone.om.snapshot.OmSnapshotUtils;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.upgrade.OMLayoutFeature;
 import org.apache.hadoop.ozone.om.upgrade.OMLayoutVersionManager;
 import org.apache.hadoop.ozone.om.upgrade.OMUpgradeFinalizer;
@@ -339,6 +338,7 @@
 import org.apache.ratis.util.ExitUtils;
 import org.apache.ratis.util.FileUtils;
 import org.apache.ratis.util.LifeCycle;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -490,9 +490,9 @@ public final class OzoneManager extends 
ServiceRuntimeInfoImpl
   private IAccessAuthorizer accessAuthorizer;
   // This metadata reader points to the active filesystem
   private OmMetadataReader omMetadataReader;
-  // Wrap active DB metadata reader in ReferenceCounted once to avoid
+  // Wrap active DB metadata reader in UncheckedAutoCloseableSupplier once to 
avoid
   // instance creation every single time.
-  private ReferenceCounted<IOmMetadataReader> rcOmMetadataReader;
+  private UncheckedAutoCloseableSupplier<IOmMetadataReader> rcOmMetadataReader;
   private OmSnapshotManager omSnapshotManager;
 
   @SuppressWarnings("methodlength")
@@ -963,7 +963,17 @@ private void instantiateServices(boolean withNewSnapshot) 
throws IOException {
         this, LOG, AUDIT, metrics, accessAuthorizer);
     // Active DB's OmMetadataReader instance does not need to be reference
     // counted, but it still needs to be wrapped to be consistent.
-    rcOmMetadataReader = new ReferenceCounted<>(omMetadataReader, true, null);
+    rcOmMetadataReader = new 
UncheckedAutoCloseableSupplier<IOmMetadataReader>() {
+      @Override
+      public IOmMetadataReader get() {
+        return omMetadataReader;
+      }
+
+      @Override
+      public void close() {
+
+      }
+    };
 
     // Reload snapshot feature config flag
     fsSnapshotEnabled = configuration.getBoolean(
@@ -2714,7 +2724,7 @@ public boolean getAllowListAllVolumes() {
     return allowListAllVolumes;
   }
 
-  public ReferenceCounted<IOmMetadataReader> getOmMetadataReader() {
+  public UncheckedAutoCloseableSupplier<IOmMetadataReader> 
getOmMetadataReader() {
     return rcOmMetadataReader;
   }
 
@@ -2984,7 +2994,7 @@ public OmBucketInfo getBucketInfo(String volume, String 
bucket)
    */
   @Override
   public OmKeyInfo lookupKey(OmKeyArgs args) throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader = getReader(args)) {
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader = 
getReader(args)) {
       return rcReader.get().lookupKey(args);
     }
   }
@@ -2996,7 +3006,7 @@ public OmKeyInfo lookupKey(OmKeyArgs args) throws 
IOException {
   public KeyInfoWithVolumeContext getKeyInfo(final OmKeyArgs args,
                                              boolean assumeS3Context)
       throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader = getReader(args)) {
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader = 
getReader(args)) {
       return rcReader.get().getKeyInfo(args, assumeS3Context);
     }
   }
@@ -3008,7 +3018,7 @@ public KeyInfoWithVolumeContext getKeyInfo(final 
OmKeyArgs args,
   public ListKeysResult listKeys(String volumeName, String bucketName,
                                  String startKey, String keyPrefix, int 
maxKeys)
       throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader =
              getReader(volumeName, bucketName, keyPrefix)) {
       return rcReader.get().listKeys(
           volumeName, bucketName, startKey, keyPrefix, maxKeys);
@@ -3841,7 +3851,7 @@ public OmMultipartUploadList listMultipartUploads(String 
volumeName,
    */
   @Override
   public OzoneFileStatus getFileStatus(OmKeyArgs args) throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader =
         getReader(args)) {
       return rcReader.get().getFileStatus(args);
     }
@@ -3852,7 +3862,7 @@ public OzoneFileStatus getFileStatus(OmKeyArgs args) 
throws IOException {
    */
   @Override
   public OmKeyInfo lookupFile(OmKeyArgs args) throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader =
         getReader(args)) {
       return rcReader.get().lookupFile(args);
     }
@@ -3872,7 +3882,7 @@ public List<OzoneFileStatus> listStatus(OmKeyArgs args, 
boolean recursive,
   public List<OzoneFileStatus> listStatus(OmKeyArgs args, boolean recursive,
       String startKey, long numEntries, boolean allowPartialPrefixes)
       throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader =
         getReader(args)) {
       return rcReader.get().listStatus(
           args, recursive, startKey, numEntries, allowPartialPrefixes);
@@ -3896,7 +3906,7 @@ public List<OzoneFileStatusLight> 
listStatusLight(OmKeyArgs args,
    */
   @Override
   public List<OzoneAcl> getAcl(OzoneObj obj) throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader =
         getReader(obj)) {
       return rcReader.get().getAcl(obj);
     }
@@ -4860,7 +4870,7 @@ public void startQuotaRepair(List<String> buckets) throws 
IOException {
   @Override
   public Map<String, String> getObjectTagging(final OmKeyArgs args)
       throws IOException {
-    try (ReferenceCounted<IOmMetadataReader> rcReader = getReader(args)) {
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcReader = 
getReader(args)) {
       return rcReader.get().getObjectTagging(args);
     }
   }
@@ -4947,9 +4957,9 @@ public static HddsProtos.OzoneManagerDetailsProto 
getOmDetailsProto(
    * Caller is responsible of closing the return value.
    * Using try-with-resources is recommended.
    * @param keyArgs OmKeyArgs
-   * @return ReferenceCounted<IOmMetadataReader, SnapshotCache>
+   * @return UncheckedAutoCloseableSupplier<IOmMetadataReader, SnapshotCache>
    */
-  private ReferenceCounted<IOmMetadataReader> getReader(OmKeyArgs keyArgs)
+  private UncheckedAutoCloseableSupplier<IOmMetadataReader> 
getReader(OmKeyArgs keyArgs)
       throws IOException {
     return omSnapshotManager.getActiveFsMetadataOrSnapshot(
         keyArgs.getVolumeName(), keyArgs.getBucketName(), 
keyArgs.getKeyName());
@@ -4962,9 +4972,9 @@ private ReferenceCounted<IOmMetadataReader> 
getReader(OmKeyArgs keyArgs)
    * @param volumeName volume name
    * @param bucketName bucket name
    * @param key key path
-   * @return ReferenceCounted<IOmMetadataReader, SnapshotCache>
+   * @return UncheckedAutoCloseableSupplier<IOmMetadataReader, SnapshotCache>
    */
-  private ReferenceCounted<IOmMetadataReader> getReader(
+  private UncheckedAutoCloseableSupplier<IOmMetadataReader> getReader(
           String volumeName, String bucketName, String key) throws IOException 
{
     return omSnapshotManager.getActiveFsMetadataOrSnapshot(
         volumeName, bucketName, key);
@@ -4975,9 +4985,9 @@ private ReferenceCounted<IOmMetadataReader> getReader(
    * Caller is responsible of closing the return value.
    * Using try-with-resources is recommended.
    * @param ozoneObj OzoneObj
-   * @return ReferenceCounted<IOmMetadataReader, SnapshotCache>
+   * @return UncheckedAutoCloseableSupplier<IOmMetadataReader, SnapshotCache>
    */
-  private ReferenceCounted<IOmMetadataReader> getReader(OzoneObj ozoneObj)
+  private UncheckedAutoCloseableSupplier<IOmMetadataReader> getReader(OzoneObj 
ozoneObj)
       throws IOException {
     return omSnapshotManager.getActiveFsMetadataOrSnapshot(
         ozoneObj.getVolumeName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/SstFilteringService.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/SstFilteringService.java
index 56bfc45c1e..ea46366d91 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/SstFilteringService.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/SstFilteringService.java
@@ -46,7 +46,7 @@
 import org.apache.hadoop.ozone.om.exceptions.OMException;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.OMLockDetails;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -190,7 +190,7 @@ public BackgroundTaskResult call() throws Exception {
                     snapshotInfo.getBucketName());
 
             try (
-                ReferenceCounted<OmSnapshot> snapshotMetadataReader =
+                UncheckedAutoCloseableSupplier<OmSnapshot> 
snapshotMetadataReader =
                     snapshotManager.get().getActiveSnapshot(
                         snapshotInfo.getVolumeName(),
                         snapshotInfo.getBucketName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/OMClientRequest.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/OMClientRequest.java
index 592e375ae2..fa18c919eb 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/OMClientRequest.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/request/OMClientRequest.java
@@ -49,7 +49,6 @@
 import org.apache.hadoop.ozone.om.protocolPB.grpc.GrpcClientConstants;
 import org.apache.hadoop.ozone.om.ratis.utils.OzoneManagerRatisUtils;
 import org.apache.hadoop.ozone.om.response.OMClientResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.LayoutVersion;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
@@ -60,6 +59,7 @@
 import org.apache.hadoop.ozone.security.acl.RequestContext;
 import org.apache.hadoop.security.UserGroupInformation;
 import 
org.apache.hadoop.security.authentication.client.AuthenticationException;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -300,7 +300,7 @@ protected void checkACLsWithFSO(OzoneManager ozoneManager, 
String volumeName,
         contextBuilder.setOwnerName(bucketOwner);
       }
 
-      try (ReferenceCounted<IOmMetadataReader> rcMetadataReader =
+      try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcMetadataReader =
           ozoneManager.getOmMetadataReader()) {
         OmMetadataReader omMetadataReader =
             (OmMetadataReader) rcMetadataReader.get();
@@ -366,7 +366,7 @@ public void checkAcls(OzoneManager ozoneManager,
       String bucketOwner)
       throws IOException {
 
-    try (ReferenceCounted<IOmMetadataReader> rcMetadataReader =
+    try (UncheckedAutoCloseableSupplier<IOmMetadataReader> rcMetadataReader =
         ozoneManager.getOmMetadataReader()) {
       OzoneAclUtils.checkAllAcls((OmMetadataReader) rcMetadataReader.get(),
           resType, storeType, aclType,
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMDirectoriesPurgeResponseWithFSO.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMDirectoriesPurgeResponseWithFSO.java
index d932603b62..721ce8401c 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMDirectoriesPurgeResponseWithFSO.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMDirectoriesPurgeResponseWithFSO.java
@@ -42,9 +42,9 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.request.key.OMDirectoriesPurgeRequestWithFSO;
 import org.apache.hadoop.ozone.om.response.CleanupTableInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMResponse;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -86,7 +86,7 @@ public void addToDBBatch(OMMetadataManager metadataManager,
           ((OmMetadataManagerImpl) metadataManager)
               .getOzoneManager().getOmSnapshotManager();
 
-      try (ReferenceCounted<OmSnapshot>
+      try (UncheckedAutoCloseableSupplier<OmSnapshot>
           rcFromSnapshotInfo = 
omSnapshotManager.getSnapshot(fromSnapshotInfo.getSnapshotId())) {
         OmSnapshot fromSnapshot = rcFromSnapshotInfo.get();
         DBStore fromSnapshotStore = fromSnapshot.getMetadataManager()
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMKeyPurgeResponse.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMKeyPurgeResponse.java
index 39584f1692..7a1aebe6a4 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMKeyPurgeResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/key/OMKeyPurgeResponse.java
@@ -34,10 +34,10 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.request.key.OMKeyPurgeRequest;
 import org.apache.hadoop.ozone.om.response.CleanupTableInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyInfo;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMResponse;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotMoveKeyInfos;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Response for {@link OMKeyPurgeRequest} request.
@@ -75,7 +75,7 @@ public void addToDBBatch(OMMetadataManager omMetadataManager,
       OmSnapshotManager omSnapshotManager =
           ((OmMetadataManagerImpl) 
omMetadataManager).getOzoneManager().getOmSnapshotManager();
 
-      try (ReferenceCounted<OmSnapshot> rcOmFromSnapshot =
+      try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmFromSnapshot =
           omSnapshotManager.getSnapshot(fromSnapshot.getSnapshotId())) {
 
         OmSnapshot fromOmSnapshot = rcOmFromSnapshot.get();
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveDeletedKeysResponse.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveDeletedKeysResponse.java
index 1cf45b17e7..6cc39d7ecd 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveDeletedKeysResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveDeletedKeysResponse.java
@@ -35,10 +35,10 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.response.CleanupTableInfo;
 import org.apache.hadoop.ozone.om.response.OMClientResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyInfo;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMResponse;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotMoveKeyInfos;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Response for OMSnapshotMoveDeletedKeysRequest.
@@ -89,13 +89,13 @@ protected void addToDBBatch(OMMetadataManager 
omMetadataManager,
         ((OmMetadataManagerImpl) omMetadataManager)
             .getOzoneManager().getOmSnapshotManager();
 
-    try (ReferenceCounted<OmSnapshot> rcOmFromSnapshot =
-        omSnapshotManager.getSnapshot(fromSnapshot.getSnapshotId())) {
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmFromSnapshot =
+             omSnapshotManager.getSnapshot(fromSnapshot.getSnapshotId())) {
 
       OmSnapshot fromOmSnapshot = rcOmFromSnapshot.get();
 
       if (nextSnapshot != null) {
-        try (ReferenceCounted<OmSnapshot>
+        try (UncheckedAutoCloseableSupplier<OmSnapshot>
             rcOmNextSnapshot = 
omSnapshotManager.getSnapshot(nextSnapshot.getSnapshotId())) {
 
           OmSnapshot nextOmSnapshot = rcOmNextSnapshot.get();
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveTableKeysResponse.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveTableKeysResponse.java
index d1b09bbeb6..2b00556d63 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveTableKeysResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/response/snapshot/OMSnapshotMoveTableKeysResponse.java
@@ -35,9 +35,9 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.response.CleanupTableInfo;
 import org.apache.hadoop.ozone.om.response.OMClientResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMResponse;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotMoveKeyInfos;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Response for OMSnapshotMoveDeletedKeysRequest.
@@ -78,13 +78,13 @@ protected void addToDBBatch(OMMetadataManager 
omMetadataManager, BatchOperation
     OmSnapshotManager omSnapshotManager = ((OmMetadataManagerImpl) 
omMetadataManager)
         .getOzoneManager().getOmSnapshotManager();
 
-    try (ReferenceCounted<OmSnapshot> rcOmFromSnapshot =
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmFromSnapshot =
              omSnapshotManager.getSnapshot(fromSnapshot.getSnapshotId())) {
 
       OmSnapshot fromOmSnapshot = rcOmFromSnapshot.get();
 
       if (nextSnapshot != null) {
-        try (ReferenceCounted<OmSnapshot>
+        try (UncheckedAutoCloseableSupplier<OmSnapshot>
             rcOmNextSnapshot = 
omSnapshotManager.getSnapshot(nextSnapshot.getSnapshotId())) {
 
           OmSnapshot nextOmSnapshot = rcOmNextSnapshot.get();
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/DirectoryDeletingService.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/DirectoryDeletingService.java
index 3790bb286b..7451032492 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/DirectoryDeletingService.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/DirectoryDeletingService.java
@@ -45,10 +45,10 @@
 import org.apache.hadoop.ozone.om.helpers.OmDirectoryInfo;
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.PurgePathRequest;
 import org.apache.hadoop.util.Time;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -301,7 +301,7 @@ private boolean previousSnapshotHasDir(
                   previousSnapshotInfo)) {
         return true;
       }
-      try (ReferenceCounted<OmSnapshot> rcLatestSnapshot =
+      try (UncheckedAutoCloseableSupplier<OmSnapshot> rcLatestSnapshot =
           omSnapshotManager.getSnapshot(
               deletedDirInfo.getVolumeName(),
               deletedDirInfo.getBucketName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/KeyDeletingService.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/KeyDeletingService.java
index 04d5db3104..7ec8ed71e6 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/KeyDeletingService.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/KeyDeletingService.java
@@ -61,13 +61,13 @@
 import org.apache.hadoop.ozone.om.helpers.RepeatedOmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.ratis.utils.OzoneManagerRatisUtils;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SetSnapshotPropertyRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotSize;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Type;
 import org.apache.ratis.protocol.ClientId;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -284,7 +284,7 @@ private void processSnapshotDeepClean(int delCount)
             continue;
           }
 
-          try (ReferenceCounted<OmSnapshot>
+          try (UncheckedAutoCloseableSupplier<OmSnapshot>
               rcCurrOmSnapshot = omSnapshotManager.getSnapshot(
                   currSnapInfo.getVolumeName(),
                   currSnapInfo.getBucketName(),
@@ -323,7 +323,7 @@ private void processSnapshotDeepClean(int delCount)
 
             Table<String, OmKeyInfo> previousKeyTable = null;
             Table<String, String> prevRenamedTable = null;
-            ReferenceCounted<OmSnapshot> rcPrevOmSnapshot = null;
+            UncheckedAutoCloseableSupplier<OmSnapshot> rcPrevOmSnapshot = null;
 
             // Split RepeatedOmKeyInfo and update current snapshot
             // deletedKeyTable and next snapshot deletedKeyTable.
@@ -341,7 +341,7 @@ private void processSnapshotDeepClean(int delCount)
             }
 
             Table<String, OmKeyInfo> previousToPrevKeyTable = null;
-            ReferenceCounted<OmSnapshot> rcPrevToPrevOmSnapshot = null;
+            UncheckedAutoCloseableSupplier<OmSnapshot> rcPrevToPrevOmSnapshot 
= null;
             if (previousToPrevSnapshot != null) {
               rcPrevToPrevOmSnapshot = omSnapshotManager.getSnapshot(
                   previousToPrevSnapshot.getVolumeName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDeletingService.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDeletingService.java
index a570a452c2..4c2151baaf 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDeletingService.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDeletingService.java
@@ -53,7 +53,6 @@
 import org.apache.hadoop.ozone.om.SnapshotChainManager;
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotMoveKeyInfos;
@@ -62,6 +61,7 @@
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Status;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Type;
 import org.apache.ratis.protocol.ClientId;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -186,7 +186,7 @@ public BackgroundTaskResult call() throws 
InterruptedException {
             LOG.info("Snapshot: {} entries will be moved to next active 
snapshot: {}",
                 snapInfo.getTableKey(), nextSnapshot.getTableKey());
           }
-          try (ReferenceCounted<OmSnapshot> snapshot = 
omSnapshotManager.getSnapshot(
+          try (UncheckedAutoCloseableSupplier<OmSnapshot> snapshot = 
omSnapshotManager.getSnapshot(
               snapInfo.getVolumeName(), snapInfo.getBucketName(), 
snapInfo.getName())) {
             KeyManager snapshotKeyManager = snapshot.get().getKeyManager();
             int moveCount = 0;
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDirectoryCleaningService.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDirectoryCleaningService.java
index bb985da559..4c80d1600f 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDirectoryCleaningService.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/service/SnapshotDirectoryCleaningService.java
@@ -55,12 +55,12 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.ratis.utils.OzoneManagerRatisUtils;
 import org.apache.hadoop.ozone.om.request.file.OMFileRequest;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SetSnapshotPropertyRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.SnapshotSize;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Type;
 import org.apache.ratis.protocol.ClientId;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Snapshot BG Service for deleted directory deep clean and exclusive size
@@ -152,8 +152,8 @@ public BackgroundTaskResult call() {
             continue;
           }
 
-          ReferenceCounted<OmSnapshot> rcPrevOmSnapshot = null;
-          ReferenceCounted<OmSnapshot> rcPrevToPrevOmSnapshot = null;
+          UncheckedAutoCloseableSupplier<OmSnapshot> rcPrevOmSnapshot = null;
+          UncheckedAutoCloseableSupplier<OmSnapshot> rcPrevToPrevOmSnapshot = 
null;
           try {
             long volumeId = metadataManager
                 .getVolumeId(currSnapInfo.getVolumeName());
@@ -208,7 +208,7 @@ public BackgroundTaskResult call() {
 
             String dbBucketKeyForDir = metadataManager.getBucketKeyPrefixFSO(
                 currSnapInfo.getVolumeName(), currSnapInfo.getBucketName());
-            try (ReferenceCounted<OmSnapshot>
+            try (UncheckedAutoCloseableSupplier<OmSnapshot>
                      rcCurrOmSnapshot = omSnapshotManager.getActiveSnapshot(
                 currSnapInfo.getVolumeName(),
                 currSnapInfo.getBucketName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/ReferenceCounted.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/ReferenceCounted.java
index 2f85f6d21b..6937b5b4e1 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/ReferenceCounted.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/ReferenceCounted.java
@@ -24,8 +24,7 @@
 /**
  * Add reference counter to an object instance.
  */
-public class ReferenceCounted<T>
-    implements AutoCloseable {
+class ReferenceCounted<T> {
 
   /**
    * Object that is being reference counted. e.g. OmSnapshot
@@ -52,7 +51,7 @@ public class ReferenceCounted<T>
    */
   private final ReferenceCountedCallback parentWithCallback;
 
-  public ReferenceCounted(T obj, boolean disableCounter,
+  ReferenceCounted(T obj, boolean disableCounter,
       ReferenceCountedCallback parentWithCallback) {
     // A param to allow disabling ref counting to reduce active DB
     //  access penalties due to AtomicLong operations.
@@ -153,11 +152,4 @@ public long getCurrentThreadRefCount() {
     long tid = Thread.currentThread().getId();
     return threadMap.getOrDefault(tid, 0L);
   }
-
-  @Override
-  public void close() {
-    // Decrease ref count by 1 when close() is called on this object
-    // so it is eligible to be used with try-with-resources.
-    decrementRefCount();
-  }
 }
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotCache.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotCache.java
index d2b64296dc..c6003af3c3 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotCache.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotCache.java
@@ -26,10 +26,12 @@
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
 import org.apache.hadoop.hdds.utils.Scheduler;
 import org.apache.hadoop.ozone.om.OMMetrics;
 import org.apache.hadoop.ozone.om.OmSnapshot;
 import org.apache.hadoop.ozone.om.exceptions.OMException;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -141,7 +143,7 @@ public enum Reason {
    * @param key SnapshotId
    * @return an OmSnapshot instance, or null on error
    */
-  public ReferenceCounted<OmSnapshot> get(UUID key) throws IOException {
+  public UncheckedAutoCloseableSupplier<OmSnapshot> get(UUID key) throws 
IOException {
     // Warn if actual cache size exceeds the soft limit already.
     if (size() > cacheSizeLimit) {
       LOG.warn("Snapshot cache size ({}) exceeds configured soft-limit ({}).",
@@ -181,7 +183,23 @@ 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);
     }
-    return rcOmSnapshot;
+    return new UncheckedAutoCloseableSupplier<OmSnapshot>() {
+      private AtomicReference<Boolean> closed = new AtomicReference<>(false);
+      @Override
+      public OmSnapshot get() {
+        return rcOmSnapshot.get();
+      }
+
+      @Override
+      public void close() {
+        closed.updateAndGet(alreadyClosed -> {
+          if (!alreadyClosed) {
+            rcOmSnapshot.decrementRefCount();
+          }
+          return true;
+        });
+      }
+    };
   }
 
   /**
@@ -245,4 +263,8 @@ public void callback(ReferenceCounted referenceCounted) {
           .getSnapshotID());
     }
   }
+
+  long totalRefCount(UUID key) {
+    return dbMap.containsKey(key) ? dbMap.get(key).getTotalRefCount() : 0;
+  }
 }
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotDiffManager.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotDiffManager.java
index f2e239615b..1e0b5f7d87 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotDiffManager.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/SnapshotDiffManager.java
@@ -130,6 +130,7 @@
 import org.apache.ozone.rocksdiff.DifferSnapshotInfo;
 import org.apache.ozone.rocksdiff.RocksDBCheckpointDiffer;
 import org.apache.ozone.rocksdiff.RocksDiffUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.rocksdb.ColumnFamilyDescriptor;
 import org.rocksdb.ColumnFamilyHandle;
 import org.rocksdb.RocksDBException;
@@ -837,8 +838,8 @@ void generateSnapshotDiffReport(final String jobKey,
     // job by RocksDBCheckpointDiffer#pruneOlderSnapshotsWithCompactionHistory.
     Path path = Paths.get(sstBackupDirForSnapDiffJobs + "/" + jobId);
 
-    ReferenceCounted<OmSnapshot> rcFromSnapshot = null;
-    ReferenceCounted<OmSnapshot> rcToSnapshot = null;
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcFromSnapshot = null;
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcToSnapshot = null;
 
     try {
       if (!areDiffJobAndSnapshotsActive(volumeName, bucketName,
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableDirFilter.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableDirFilter.java
index 7384ed88a9..9dcb49ac3e 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableDirFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableDirFilter.java
@@ -29,7 +29,7 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Class to filter out deleted directories which are reclaimable based on 
their presence in previous snapshot in
@@ -56,7 +56,7 @@ protected String getBucketName(Table.KeyValue<String, 
OmKeyInfo> keyValue) throw
 
   @Override
   protected Boolean isReclaimable(Table.KeyValue<String, OmKeyInfo> 
deletedDirInfo) throws IOException {
-    ReferenceCounted<OmSnapshot> previousSnapshot = getPreviousOmSnapshot(0);
+    UncheckedAutoCloseableSupplier<OmSnapshot> previousSnapshot = 
getPreviousOmSnapshot(0);
     KeyManager prevKeyManager = previousSnapshot == null ? null : 
previousSnapshot.get().getKeyManager();
     return isDirReclaimable(getVolumeId(), getBucketInfo(), 
deletedDirInfo.getValue(), getKeyManager(), prevKeyManager);
   }
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableFilter.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableFilter.java
index 1f97fe10d8..0bb53e6280 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableFilter.java
@@ -37,9 +37,9 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
 import org.apache.hadoop.ozone.om.snapshot.MultiSnapshotLocks;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import org.apache.ratis.util.function.CheckedFunction;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -60,7 +60,7 @@ public abstract class ReclaimableFilter<V>
   private final List<SnapshotInfo> tmpValidationSnapshotInfos;
   private final List<UUID> lockedSnapshotIds;
   private final List<SnapshotInfo> previousSnapshotInfos;
-  private final List<ReferenceCounted<OmSnapshot>> previousOmSnapshots;
+  private final List<UncheckedAutoCloseableSupplier<OmSnapshot>> 
previousOmSnapshots;
   private final MultiSnapshotLocks snapshotIdLocks;
   private Long volumeId;
   private OmBucketInfo bucketInfo;
@@ -131,7 +131,7 @@ private boolean 
validateExistingLastNSnapshotsInChain(String volume, String buck
     }
     for (int i = 0; i < expectedLastNSnapshotsInChain.size(); i++) {
       SnapshotInfo snapshotInfo = expectedLastNSnapshotsInChain.get(i);
-      ReferenceCounted<OmSnapshot> omSnapshot = previousOmSnapshots.get(i);
+      UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot = 
previousOmSnapshots.get(i);
       UUID snapshotId = snapshotInfo == null ? null : 
snapshotInfo.getSnapshotId();
       UUID existingOmSnapshotId = omSnapshot == null ? null : 
omSnapshot.get().getSnapshotID();
       if (!Objects.equals(snapshotId, existingOmSnapshotId)) {
@@ -212,7 +212,7 @@ private void cleanup() {
     lockedSnapshotIds.clear();
   }
 
-  protected ReferenceCounted<OmSnapshot> getPreviousOmSnapshot(int index) {
+  protected UncheckedAutoCloseableSupplier<OmSnapshot> 
getPreviousOmSnapshot(int index) {
     return previousOmSnapshots.get(index);
   }
 
@@ -240,7 +240,7 @@ List<SnapshotInfo> getPreviousSnapshotInfos() {
     return previousSnapshotInfos;
   }
 
-  List<ReferenceCounted<OmSnapshot>> getPreviousOmSnapshots() {
+  List<UncheckedAutoCloseableSupplier<OmSnapshot>> getPreviousOmSnapshots() {
     return previousOmSnapshots;
   }
 }
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableKeyFilter.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableKeyFilter.java
index 6b558b5854..d357f66097 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableKeyFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableKeyFilter.java
@@ -34,9 +34,9 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.ratis.util.MemoizedCheckedSupplier;
 import org.apache.ratis.util.function.CheckedSupplier;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Filter to return deleted keys which are reclaimable based on their presence 
in previous snapshot in
@@ -85,7 +85,7 @@ protected String getBucketName(Table.KeyValue<String, 
OmKeyInfo> keyValue) throw
    * @throws IOException If an error occurs while accessing snapshot data or 
key information.
    */
   protected Boolean isReclaimable(Table.KeyValue<String, OmKeyInfo> 
deletedKeyInfo) throws IOException {
-    ReferenceCounted<OmSnapshot> previousSnapshot = getPreviousOmSnapshot(1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> previousSnapshot = 
getPreviousOmSnapshot(1);
 
 
     KeyManager previousKeyManager = Optional.ofNullable(previousSnapshot)
@@ -101,7 +101,7 @@ protected Boolean isReclaimable(Table.KeyValue<String, 
OmKeyInfo> deletedKeyInfo
       return true;
     }
 
-    ReferenceCounted<OmSnapshot> previousToPreviousSnapshot = 
getPreviousOmSnapshot(0);
+    UncheckedAutoCloseableSupplier<OmSnapshot> previousToPreviousSnapshot = 
getPreviousOmSnapshot(0);
     KeyManager previousToPreviousKeyManager = 
Optional.ofNullable(previousToPreviousSnapshot)
         .map(i -> i.get().getKeyManager()).orElse(null);
 
diff --git 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableRenameEntryFilter.java
 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableRenameEntryFilter.java
index 563f71d81f..cad4ebc6df 100644
--- 
a/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableRenameEntryFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/main/java/org/apache/hadoop/ozone/om/snapshot/filter/ReclaimableRenameEntryFilter.java
@@ -29,7 +29,7 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.helpers.WithObjectID;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 
 /**
  * Class to filter out rename table entries which are reclaimable based on the 
key presence in previous snapshot's
@@ -53,7 +53,7 @@ public ReclaimableRenameEntryFilter(OzoneManager ozoneManager,
    */
   @Override
   protected Boolean isReclaimable(Table.KeyValue<String, String> renameEntry) 
throws IOException {
-    ReferenceCounted<OmSnapshot> previousSnapshot = getPreviousOmSnapshot(0);
+    UncheckedAutoCloseableSupplier<OmSnapshot> previousSnapshot = 
getPreviousOmSnapshot(0);
     Table<String, OmKeyInfo> previousKeyTable = null;
     Table<String, OmDirectoryInfo> prevDirTable = null;
     if (previousSnapshot != null) {
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMDirectoriesPurgeRequestAndResponse.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMDirectoriesPurgeRequestAndResponse.java
index 928c375c70..d44b835847 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMDirectoriesPurgeRequestAndResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMDirectoriesPurgeRequestAndResponse.java
@@ -43,9 +43,9 @@
 import org.apache.hadoop.ozone.om.request.OMRequestTestUtils;
 import 
org.apache.hadoop.ozone.om.response.key.OMDirectoriesPurgeResponseWithFSO;
 import org.apache.hadoop.ozone.om.response.key.OMKeyPurgeResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -224,7 +224,7 @@ public void 
testValidateAndUpdateCacheSnapshotLastTransactionInfoUpdated() throw
 
     String snapshotName = "snap1";
     SnapshotInfo snapshotInfo = createSnapshot(snapshotName);
-    ReferenceCounted<OmSnapshot> rcOmSnapshot = 
ozoneManager.getOmSnapshotManager()
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
ozoneManager.getOmSnapshotManager()
         .getSnapshot(snapshotInfo.getVolumeName(), 
snapshotInfo.getBucketName(), snapshotInfo.getName());
     // Keys should be present in snapshot
     validateDeletedKeysTable(rcOmSnapshot.get().getMetadataManager(), 
deletedKeyInfos, true);
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyPurgeRequestAndResponse.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyPurgeRequestAndResponse.java
index 46f442c45e..39c3995343 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyPurgeRequestAndResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyPurgeRequestAndResponse.java
@@ -32,7 +32,6 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.request.OMRequestTestUtils;
 import org.apache.hadoop.ozone.om.response.key.OMKeyPurgeResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.DeletedKeys;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMResponse;
@@ -40,6 +39,7 @@
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.PurgeKeysResponse;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Status;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Type;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.Test;
 
 /**
@@ -176,7 +176,7 @@ public void testKeyPurgeInSnapshot() throws Exception {
       assertFalse(omMetadataManager.getDeletedTable().isExist(deletedKey));
     }
 
-    ReferenceCounted<OmSnapshot> rcOmSnapshot = 
ozoneManager.getOmSnapshotManager()
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
ozoneManager.getOmSnapshotManager()
         .getSnapshot(snapInfo.getVolumeName(), snapInfo.getBucketName(), 
snapInfo.getName());
     OmSnapshot omSnapshot = rcOmSnapshot.get();
 
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyRequest.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyRequest.java
index 8ddb1a23d1..b45625b1bd 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyRequest.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/key/TestOMKeyRequest.java
@@ -83,7 +83,6 @@
 import org.apache.hadoop.ozone.om.request.snapshot.OMSnapshotCreateRequest;
 import org.apache.hadoop.ozone.om.request.snapshot.TestOMSnapshotCreateRequest;
 import org.apache.hadoop.ozone.om.response.snapshot.OMSnapshotCreateResponse;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.upgrade.OMLayoutVersionManager;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyArgs;
@@ -91,6 +90,7 @@
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.hadoop.util.Time;
 import org.apache.ozone.test.GenericTestUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -192,8 +192,8 @@ public void setup() throws Exception {
     when(ozoneManager.getAccessAuthorizer())
         .thenReturn(new OzoneNativeAuthorizer());
 
-    ReferenceCounted<IOmMetadataReader> rcOmMetadataReader =
-        mock(ReferenceCounted.class);
+    UncheckedAutoCloseableSupplier<IOmMetadataReader> rcOmMetadataReader =
+        mock(UncheckedAutoCloseableSupplier.class);
     when(ozoneManager.getOmMetadataReader()).thenReturn(rcOmMetadataReader);
     // Init OmMetadataReader to let the test pass
     OmMetadataReader omMetadataReader = mock(OmMetadataReader.class);
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/s3/multipart/TestS3MultipartRequest.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/s3/multipart/TestS3MultipartRequest.java
index ce224d9895..d0f8ac9f55 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/s3/multipart/TestS3MultipartRequest.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/request/s3/multipart/TestS3MultipartRequest.java
@@ -49,13 +49,13 @@
 import org.apache.hadoop.ozone.om.helpers.KeyValueUtil;
 import org.apache.hadoop.ozone.om.request.OMClientRequest;
 import org.apache.hadoop.ozone.om.request.OMRequestTestUtils;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.upgrade.OMLayoutVersionManager;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.KeyArgs;
 import 
org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.OMRequest;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos.Part;
 import org.apache.hadoop.ozone.security.acl.OzoneNativeAuthorizer;
 import org.apache.hadoop.security.UserGroupInformation;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.io.TempDir;
@@ -85,8 +85,8 @@ public void setup() throws Exception {
     when(ozoneManager.getMetrics()).thenReturn(omMetrics);
     when(ozoneManager.getMetadataManager()).thenReturn(omMetadataManager);
     auditLogger = mock(AuditLogger.class);
-    ReferenceCounted<IOmMetadataReader> rcOmMetadataReader =
-        mock(ReferenceCounted.class);
+    UncheckedAutoCloseableSupplier<IOmMetadataReader> rcOmMetadataReader =
+        mock(UncheckedAutoCloseableSupplier.class);
     when(ozoneManager.getOmMetadataReader()).thenReturn(rcOmMetadataReader);
     // Init OmMetadataReader to let the test pass
     OmMetadataReader omMetadataReader = mock(OmMetadataReader.class);
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/response/snapshot/TestOMSnapshotMoveTableKeysResponse.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/response/snapshot/TestOMSnapshotMoveTableKeysResponse.java
index dff833aa5b..aabacd8c50 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/response/snapshot/TestOMSnapshotMoveTableKeysResponse.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/response/snapshot/TestOMSnapshotMoveTableKeysResponse.java
@@ -36,10 +36,10 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.RepeatedOmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import org.apache.hadoop.ozone.om.snapshot.TestSnapshotRequestAndResponse;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -107,9 +107,9 @@ private <V> void addDataToTable(Table<String, V> table, 
List<Pair<String, V>> va
   public void testMoveTableKeysToNextSnapshot(boolean nextSnapshotExists) 
throws Exception {
     createSnapshots(nextSnapshotExists);
 
-    try (ReferenceCounted<OmSnapshot> snapshot1 = 
getOmSnapshotManager().getSnapshot(getVolumeName(), getBucketName(),
-        snapshotName1);
-         ReferenceCounted<OmSnapshot> snapshot2 = nextSnapshotExists ? 
getOmSnapshotManager().getSnapshot(
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> snapshot1 = 
getOmSnapshotManager().getSnapshot(
+        getVolumeName(), getBucketName(), snapshotName1);
+         UncheckedAutoCloseableSupplier<OmSnapshot> snapshot2 = 
nextSnapshotExists ? getOmSnapshotManager().getSnapshot(
              getVolumeName(), getBucketName(), snapshotName2) : null) {
       OmSnapshot snapshot = snapshot1.get();
       List<OzoneManagerProtocolProtos.SnapshotMoveKeyInfos> deletedTable = new 
ArrayList<>();
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/service/TestKeyDeletingService.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/service/TestKeyDeletingService.java
index 46538492c6..a306ae1cf1 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/service/TestKeyDeletingService.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/service/TestKeyDeletingService.java
@@ -87,11 +87,11 @@
 import org.apache.hadoop.ozone.om.protocol.OzoneManagerProtocol;
 import org.apache.hadoop.ozone.om.ratis.utils.OzoneManagerRatisUtils;
 import org.apache.hadoop.ozone.om.request.OMRequestTestUtils;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.protocol.proto.OzoneManagerProtocolProtos;
 import org.apache.ozone.test.GenericTestUtils;
 import org.apache.ozone.test.OzoneTestBase;
 import org.apache.ratis.util.ExitUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
@@ -483,7 +483,7 @@ void testSnapshotDeepClean() throws Exception {
 
       keyDeletingService.resume();
 
-      try (ReferenceCounted<OmSnapshot> rcOmSnapshot =
+      try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot =
                om.getOmSnapshotManager().getSnapshot(volumeName, bucketName, 
snap3)) {
         OmSnapshot snapshot3 = rcOmSnapshot.get();
 
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotCache.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotCache.java
index 0fcc7c3de1..4a9bc0f857 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotCache.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotCache.java
@@ -33,6 +33,7 @@
 import org.apache.hadoop.ozone.om.OMMetrics;
 import org.apache.hadoop.ozone.om.OmSnapshot;
 import org.apache.ozone.test.GenericTestUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
@@ -93,7 +94,7 @@ void tearDown() {
   void testGet() throws IOException {
     final UUID dbKey1 = UUID.randomUUID();
     assertEquals(0, omMetrics.getNumSnapshotCacheSize());
-    ReferenceCounted<OmSnapshot> omSnapshot = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot = 
snapshotCache.get(dbKey1);
     assertNotNull(omSnapshot);
     assertNotNull(omSnapshot.get());
     assertInstanceOf(OmSnapshot.class, omSnapshot.get());
@@ -105,14 +106,13 @@ void testGet() throws IOException {
   @DisplayName("get() same entry twice yields one cache entry only")
   void testGetTwice() throws IOException {
     final UUID dbKey1 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot1 = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot1 = 
snapshotCache.get(dbKey1);
     assertNotNull(omSnapshot1);
     assertEquals(1, snapshotCache.size());
     assertEquals(1, omMetrics.getNumSnapshotCacheSize());
 
-    ReferenceCounted<OmSnapshot> omSnapshot1again = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot1again = 
snapshotCache.get(dbKey1);
     // Should be the same instance
-    assertEquals(omSnapshot1, omSnapshot1again);
     assertEquals(omSnapshot1.get(), omSnapshot1again.get());
     assertEquals(1, snapshotCache.size());
     assertEquals(1, omMetrics.getNumSnapshotCacheSize());
@@ -122,7 +122,7 @@ void testGetTwice() throws IOException {
   @DisplayName("release(String)")
   void testReleaseByDbKey() throws IOException {
     final UUID dbKey1 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot1 = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot1 = 
snapshotCache.get(dbKey1);
     assertNotNull(omSnapshot1);
     assertNotNull(omSnapshot1.get());
     assertEquals(1, snapshotCache.size());
@@ -138,7 +138,7 @@ void testReleaseByDbKey() throws IOException {
   @DisplayName("invalidate()")
   void testInvalidate() throws IOException {
     final UUID dbKey1 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot = 
snapshotCache.get(dbKey1);
     assertNotNull(omSnapshot);
     assertEquals(1, snapshotCache.size());
     assertEquals(1, omMetrics.getNumSnapshotCacheSize());
@@ -157,13 +157,13 @@ void testInvalidate() throws IOException {
   @DisplayName("invalidateAll()")
   void testInvalidateAll() throws IOException {
     final UUID dbKey1 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot1 = snapshotCache.get(dbKey1);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot1 = 
snapshotCache.get(dbKey1);
     assertNotNull(omSnapshot1);
     assertEquals(1, snapshotCache.size());
     assertEquals(1, omMetrics.getNumSnapshotCacheSize());
 
     final UUID dbKey2 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot2 = snapshotCache.get(dbKey2);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot2 = 
snapshotCache.get(dbKey2);
     assertNotNull(omSnapshot2);
     assertEquals(2, snapshotCache.size());
     assertEquals(2, omMetrics.getNumSnapshotCacheSize());
@@ -171,7 +171,7 @@ void testInvalidateAll() throws IOException {
     assertNotEquals(omSnapshot1, omSnapshot2);
 
     final UUID dbKey3 = UUID.randomUUID();
-    ReferenceCounted<OmSnapshot> omSnapshot3 = snapshotCache.get(dbKey3);
+    UncheckedAutoCloseableSupplier<OmSnapshot> omSnapshot3 = 
snapshotCache.get(dbKey3);
     assertNotNull(omSnapshot3);
     assertEquals(3, snapshotCache.size());
     assertEquals(3, omMetrics.getNumSnapshotCacheSize());
@@ -283,8 +283,8 @@ void testEviction2() throws IOException, 
InterruptedException, TimeoutException
   void testEviction3WithClose() throws IOException, InterruptedException, 
TimeoutException {
 
     final UUID dbKey1 = UUID.randomUUID();
-    try (ReferenceCounted<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey1)) {
-      assertEquals(1L, rcOmSnapshot.getTotalRefCount());
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey1)) {
+      assertEquals(1L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
       assertEquals(1, snapshotCache.size());
       assertEquals(1, omMetrics.getNumSnapshotCacheSize());
     }
@@ -295,26 +295,26 @@ void testEviction3WithClose() throws IOException, 
InterruptedException, TimeoutE
     assertEquals(1, omMetrics.getNumSnapshotCacheSize());
 
     final UUID dbKey2 = UUID.randomUUID();
-    try (ReferenceCounted<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey2)) {
-      assertEquals(1L, rcOmSnapshot.getTotalRefCount());
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey2)) {
+      assertEquals(1L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
       assertEquals(2, snapshotCache.size());
       assertEquals(2, omMetrics.getNumSnapshotCacheSize());
       // Get dbKey2 entry a second time
-      try (ReferenceCounted<OmSnapshot> rcOmSnapshot2 = 
snapshotCache.get(dbKey2)) {
-        assertEquals(2L, rcOmSnapshot.getTotalRefCount());
-        assertEquals(2L, rcOmSnapshot2.getTotalRefCount());
+      try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot2 = 
snapshotCache.get(dbKey2)) {
+        assertEquals(2L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
+        assertEquals(2L, 
snapshotCache.totalRefCount(rcOmSnapshot2.get().getSnapshotID()));
         assertEquals(2, snapshotCache.size());
         assertEquals(2, omMetrics.getNumSnapshotCacheSize());
       }
-      assertEquals(1L, rcOmSnapshot.getTotalRefCount());
+      assertEquals(1L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
     }
     assertEquals(0L, snapshotCache.getDbMap().get(dbKey2).getTotalRefCount());
     assertEquals(2, snapshotCache.size());
     assertEquals(2, omMetrics.getNumSnapshotCacheSize());
 
     final UUID dbKey3 = UUID.randomUUID();
-    try (ReferenceCounted<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey3)) {
-      assertEquals(1L, rcOmSnapshot.getTotalRefCount());
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey3)) {
+      assertEquals(1L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
       assertEquals(3, snapshotCache.size());
       assertEquals(3, omMetrics.getNumSnapshotCacheSize());
     }
@@ -323,9 +323,9 @@ void testEviction3WithClose() throws IOException, 
InterruptedException, TimeoutE
     assertEquals(3, omMetrics.getNumSnapshotCacheSize());
 
     final UUID dbKey4 = UUID.randomUUID();
-    try (ReferenceCounted<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey4)) {
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> rcOmSnapshot = 
snapshotCache.get(dbKey4)) {
       GenericTestUtils.waitFor(() -> snapshotCache.size() == 1, 50, 3000);
-      assertEquals(1L, rcOmSnapshot.getTotalRefCount());
+      assertEquals(1L, 
snapshotCache.totalRefCount(rcOmSnapshot.get().getSnapshotID()));
       assertEquals(1, snapshotCache.size());
       assertEquals(1, omMetrics.getNumSnapshotCacheSize());
     }
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDiffManager.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDiffManager.java
index ed4df7d505..71ce98c387 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDiffManager.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSnapshotDiffManager.java
@@ -153,6 +153,7 @@
 import org.apache.ozone.test.tag.Flaky;
 import org.apache.ratis.util.ExitUtils;
 import org.apache.ratis.util.TimeDuration;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeAll;
@@ -437,9 +438,9 @@ public void testGetDeltaFilesWithDag(int numberOfFiles) 
throws IOException {
         eq(diffDir))
     ).thenReturn(Optional.of(Lists.newArrayList(randomStrings)));
 
-    ReferenceCounted<OmSnapshot> rcFromSnapshot =
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcFromSnapshot =
         omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap1.toString());
-    ReferenceCounted<OmSnapshot> rcToSnapshot =
+    UncheckedAutoCloseableSupplier<OmSnapshot> rcToSnapshot =
         omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap2.toString());
     OmSnapshot fromSnapshot = rcFromSnapshot.get();
     OmSnapshot toSnapshot = rcToSnapshot.get();
@@ -517,9 +518,9 @@ public void testGetDeltaFilesWithFullDiff(int numberOfFiles,
             .thenReturn(Optional.ofNullable(Collections.emptyList()));
       }
 
-      ReferenceCounted<OmSnapshot> rcFromSnapshot =
+      UncheckedAutoCloseableSupplier<OmSnapshot> rcFromSnapshot =
           omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap1.toString());
-      ReferenceCounted<OmSnapshot> rcToSnapshot =
+      UncheckedAutoCloseableSupplier<OmSnapshot> rcToSnapshot =
           omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap2.toString());
       OmSnapshot fromSnapshot = rcFromSnapshot.get();
       OmSnapshot toSnapshot = rcToSnapshot.get();
@@ -588,9 +589,9 @@ public void testGetDeltaFilesWithDifferThrowException(int 
numberOfFiles)
               any(DifferSnapshotInfo.class),
               anyString());
 
-      ReferenceCounted<OmSnapshot> rcFromSnapshot =
+      UncheckedAutoCloseableSupplier<OmSnapshot> rcFromSnapshot =
           omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap1.toString());
-      ReferenceCounted<OmSnapshot> rcToSnapshot =
+      UncheckedAutoCloseableSupplier<OmSnapshot> rcToSnapshot =
           omSnapshotManager.getActiveSnapshot(VOLUME_NAME, BUCKET_NAME, 
snap2.toString());
       OmSnapshot fromSnapshot = rcFromSnapshot.get();
       OmSnapshot toSnapshot = rcToSnapshot.get();
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSstFilteringService.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSstFilteringService.java
index 7e03fe7df2..e523f32ef7 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSstFilteringService.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/TestSstFilteringService.java
@@ -67,6 +67,7 @@
 import org.apache.hadoop.ozone.om.request.OMRequestTestUtils;
 import org.apache.hadoop.security.UserGroupInformation;
 import org.apache.ratis.util.ExitUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.Test;
@@ -489,7 +490,7 @@ private Set<String> getKeysFromSnapshot(String volume,
                                           String snapshot) throws IOException {
     SnapshotInfo snapshotInfo = om.getMetadataManager().getSnapshotInfoTable()
         .get(SnapshotInfo.getTableKey(volume, bucket, snapshot));
-    try (ReferenceCounted<OmSnapshot> snapshotMetadataReader =
+    try (UncheckedAutoCloseableSupplier<OmSnapshot> snapshotMetadataReader =
              om.getOmSnapshotManager().getActiveSnapshot(
                  snapshotInfo.getVolumeName(),
                  snapshotInfo.getBucketName(),
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/AbstractReclaimableFilterTest.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/AbstractReclaimableFilterTest.java
index 4f133868f9..fc7a53422c 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/AbstractReclaimableFilterTest.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/AbstractReclaimableFilterTest.java
@@ -59,11 +59,11 @@
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
 import org.apache.hadoop.ozone.om.lock.OMLockDetails;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotCache;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotDiffManager;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
 import org.apache.ozone.rocksdiff.RocksDBCheckpointDiffer;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestInstance;
@@ -181,7 +181,7 @@ private void mockOmSnapshotManager(OzoneManager om) throws 
RocksDBException, IOE
                  doNothing().when(mock).close());
          MockedConstruction<SnapshotCache> mockedCache = 
Mockito.mockConstruction(SnapshotCache.class,
              (mock, context) -> {
-               Map<UUID, ReferenceCounted<OmSnapshot>> map = new HashMap<>();
+               Map<UUID, UncheckedAutoCloseableSupplier<OmSnapshot>> map = new 
HashMap<>();
                when(mock.get(any(UUID.class))).thenAnswer(i -> {
                  if (snapshotInfos.values().stream().flatMap(List::stream)
                      .map(SnapshotInfo::getSnapshotId)
@@ -189,7 +189,7 @@ private void mockOmSnapshotManager(OzoneManager om) throws 
RocksDBException, IOE
                    throw new IOException("Snapshot " + i.getArgument(0, 
UUID.class) + " not found");
                  }
                  return map.computeIfAbsent(i.getArgument(0, UUID.class), (k) 
-> {
-                   ReferenceCounted<OmSnapshot> ref = 
mock(ReferenceCounted.class);
+                   UncheckedAutoCloseableSupplier<OmSnapshot> ref = 
mock(UncheckedAutoCloseableSupplier.class);
                    OmSnapshot omSnapshot = mock(OmSnapshot.class);
                    when(omSnapshot.getSnapshotID()).thenReturn(k);
                    when(ref.get()).thenReturn(omSnapshot);
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableDirFilter.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableDirFilter.java
index a8c6a11b2e..a85da9900a 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableDirFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableDirFilter.java
@@ -37,7 +37,7 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -76,7 +76,7 @@ private void testReclaimableDirFilter(String volume, String 
bucket, int index,
     long volumeId = getOzoneManager().getMetadataManager().getVolumeId(volume);
     KeyManager keyManager = getKeyManager();
     if (prevSnapshotInfo != null) {
-      ReferenceCounted<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
+      UncheckedAutoCloseableSupplier<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
           .map(info -> {
             try {
               return getOmSnapshotManager().getActiveSnapshot(volume, bucket, 
info.getName());
@@ -106,7 +106,7 @@ private OmDirectoryInfo getMockedOmDirInfo(long objectId) {
     return keyInfo;
   }
 
-  private KeyManager mockOmSnapshot(ReferenceCounted<OmSnapshot> snapshot) {
+  private KeyManager mockOmSnapshot(UncheckedAutoCloseableSupplier<OmSnapshot> 
snapshot) {
     if (snapshot != null) {
       OmSnapshot omSnapshot = snapshot.get();
       KeyManager keyManager = mock(KeyManager.class);
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableFilter.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableFilter.java
index 9fdc874391..2b986f8fb3 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableFilter.java
@@ -42,8 +42,8 @@
 import org.apache.hadoop.ozone.om.helpers.BucketLayout;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.TestInstance;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -108,7 +108,7 @@ private void testSnapshotInitAndLocking(
     Assertions.assertEquals(infos.size(), 
getReclaimableFilter().getPreviousOmSnapshots().size());
     Assertions.assertEquals(infos.stream().map(si -> si == null ? null : 
si.getSnapshotId())
         .collect(Collectors.toList()), 
getReclaimableFilter().getPreviousOmSnapshots().stream()
-        .map(i -> i == null ? null : ((ReferenceCounted<OmSnapshot>) 
i).get().getSnapshotID())
+        .map(i -> i == null ? null : 
((UncheckedAutoCloseableSupplier<OmSnapshot>) i).get().getSnapshotID())
         .collect(Collectors.toList()));
     infos.add(currentSnapshotInfo);
     
Assertions.assertEquals(infos.stream().filter(Objects::nonNull).map(SnapshotInfo::getSnapshotId).collect(
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableKeyFilter.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableKeyFilter.java
index 021a44eafd..9db680c18f 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableKeyFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableKeyFilter.java
@@ -38,8 +38,8 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
 import org.apache.hadoop.ozone.om.snapshot.SnapshotUtils;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -82,7 +82,7 @@ List<Arguments> testReclaimableFilterArguments() {
     return arguments;
   }
 
-  private KeyManager mockOmSnapshot(ReferenceCounted<OmSnapshot> snapshot) {
+  private KeyManager mockOmSnapshot(UncheckedAutoCloseableSupplier<OmSnapshot> 
snapshot) {
     if (snapshot != null) {
       OmSnapshot omSnapshot = snapshot.get();
       KeyManager keyManager = mock(KeyManager.class);
@@ -104,7 +104,7 @@ private void testReclaimableKeyFilter(String volume, String 
bucket, int index,
     OmBucketInfo bucketInfo = getOzoneManager().getBucketInfo(volume, bucket);
     long volumeId = getOzoneManager().getMetadataManager().getVolumeId(volume);
 
-    ReferenceCounted<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
+    UncheckedAutoCloseableSupplier<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
         .map(info -> {
           try {
             return getOmSnapshotManager().getActiveSnapshot(volume, bucket, 
info.getName());
@@ -112,7 +112,7 @@ private void testReclaimableKeyFilter(String volume, String 
bucket, int index,
             throw new RuntimeException(e);
           }
         }).orElse(null);
-    ReferenceCounted<OmSnapshot> prevToPrevSnap = 
Optional.ofNullable(previousToPreviousSapshotInfo)
+    UncheckedAutoCloseableSupplier<OmSnapshot> prevToPrevSnap = 
Optional.ofNullable(previousToPreviousSapshotInfo)
         .map(info -> {
           try {
             return getOmSnapshotManager().getActiveSnapshot(volume, bucket, 
info.getName());
diff --git 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableRenameEntryFilter.java
 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableRenameEntryFilter.java
index cb0dabc6cf..4fad10f248 100644
--- 
a/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableRenameEntryFilter.java
+++ 
b/hadoop-ozone/ozone-manager/src/test/java/org/apache/hadoop/ozone/om/snapshot/filter/TestReclaimableRenameEntryFilter.java
@@ -44,7 +44,7 @@
 import org.apache.hadoop.ozone.om.helpers.OmKeyInfo;
 import org.apache.hadoop.ozone.om.helpers.SnapshotInfo;
 import org.apache.hadoop.ozone.om.lock.IOzoneManagerLock;
-import org.apache.hadoop.ozone.om.snapshot.ReferenceCounted;
+import org.apache.ratis.util.function.UncheckedAutoCloseableSupplier;
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.Arguments;
 import org.junit.jupiter.params.provider.MethodSource;
@@ -82,7 +82,7 @@ private void testReclaimableRenameEntryFilter(String volume, 
String bucket, int
     SnapshotInfo prevSnapshotInfo = snapshotInfos.get(0);
     OmBucketInfo bucketInfo = getOzoneManager().getBucketInfo(volume, bucket);
     if (prevSnapshotInfo != null) {
-      ReferenceCounted<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
+      UncheckedAutoCloseableSupplier<OmSnapshot> prevSnap = 
Optional.ofNullable(prevSnapshotInfo)
           .map(info -> {
             try {
               return getOmSnapshotManager().getActiveSnapshot(volume, bucket, 
info.getName());
@@ -115,7 +115,7 @@ private <T> Table<String, T> getFailingMockedTable() throws 
IOException {
     return table;
   }
 
-  private void mockOmSnapshot(ReferenceCounted<OmSnapshot> snapshot,
+  private void mockOmSnapshot(UncheckedAutoCloseableSupplier<OmSnapshot> 
snapshot,
                               OmBucketInfo bucketInfo, Table<String, 
OmKeyInfo> keyTable,
                               Table<String, OmDirectoryInfo> dirTable) {
     if (snapshot != null) {


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]


Reply via email to