HBASE-30285 BucketCache shutdown/disable does not release backingMap-owned BucketEntry references

hbase-server/src/main/java/org/apache/hadoop/hbase/io/hfile/bucket/BucketCache.java

@@ -92,6 +92,7 @@
 import org.apache.hadoop.hbase.util.IdReadWriteLockWithObjectPool;
 import org.apache.hadoop.hbase.util.IdReadWriteLockWithObjectPool.ReferenceType;
 import org.apache.hadoop.hbase.util.Pair;
+import org.apache.hadoop.hbase.util.Threads;
 import org.apache.hadoop.util.StringUtils;
 import org.apache.yetus.audience.InterfaceAudience;
 import org.slf4j.Logger;
@@ -206,6 +207,9 @@ protected enum CacheState {
    */
   private volatile CacheState cacheState;
 
+  /** The single cleanup thread shared by disable and explicit shutdown calls. */
+  private volatile Thread cacheCleanupThread;
+
   /**
    * A list of writer queues. We have a queue per {@link WriterThread} we have running. In other
    * words, the work adding blocks to the BucketCache is divided up amongst the running
@@ -690,6 +694,7 @@ public Cacheable getBlock(BlockCacheKey key, boolean caching, boolean repeat,
     if (bucketEntry != null) {
       long start = System.nanoTime();
       ReentrantReadWriteLock lock = offsetLock.getLock(bucketEntry.offset());
+      boolean failedBucketEntryRead = false;
       try {
         lock.readLock().lock();
         // We can not read here even if backingMap does contain the given key because its offset
@@ -722,22 +727,48 @@ public Cacheable getBlock(BlockCacheKey key, boolean caching, boolean repeat,
         // When using file io engine persistent cache,
         // the cache map state might differ from the actual cache. If we reach this block,
         // we should remove the cache key entry from the backing map
-        backingMap.remove(key);
-        fileNotFullyCached(key, bucketEntry);
+        failedBucketEntryRead = true;
         LOG.debug("Failed to fetch block for cache key: {}.", key, hioex);
       } catch (IOException ioex) {
         LOG.error("Failed reading block " + key + " from bucket cache", ioex);
         checkIOErrorIsTolerated();
       } finally {
         lock.readLock().unlock();
       }
+      if (failedBucketEntryRead) {
+        removeFailedBucketEntry(bucketEntry);
+      }
     }
     if (!repeat && updateCacheMetrics) {
       cacheStats.miss(caching, key.isPrimary(), key.getBlockType());
     }
     return null;
   }
 
+  private void removeFailedBucketEntry(BucketEntry bucketEntry) {
+    BlockCacheKey cacheKey = findBackingMapKey(bucketEntry);
+    if (cacheKey == null) {
+      return;
+    }
+    bucketEntry.withWriteLock(offsetLock, () -> {
+      if (backingMap.remove(cacheKey, bucketEntry)) {
+        blockEvicted(cacheKey, bucketEntry, true, false);
+      }
+      return null;
+    });
+  }
+
+  private BlockCacheKey findBackingMapKey(BucketEntry bucketEntry) {
+    // Reference file lookups use a different key from the one stored in backingMap. This only runs
+    // after a failed read, so find the stored key to preserve its region metadata during eviction.
+    for (Map.Entry<BlockCacheKey, BucketEntry> entry : backingMap.entrySet()) {
+      if (entry.getValue() == bucketEntry) {
+        return entry.getKey();
+      }
+    }
+    return null;
+  }
+
   /**
    * This method is invoked after the bucketEntry is removed from {@link BucketCache#backingMap}
    */
@@ -1798,70 +1829,110 @@ private void checkIOErrorIsTolerated() {
       if (isCacheEnabled() && (now - ioErrorStartTimeTmp) > this.ioErrorsTolerationDuration) {
         LOG.error("IO errors duration time has exceeded " + ioErrorsTolerationDuration
           + "ms, disabling cache, please check your IOEngine");
-        disableCache();
+        disableCache(false);
       }
     } else {
       this.ioErrorStartTime = now;
     }
   }
 
   /**
-   * Used to shut down the cache -or- turn it off in the case of something broken.
+   * Used to shut down the cache -or- turn it off in the case of something broken. Cleanup runs
+   * separately because IO errors can invoke this method from a writer thread or while holding an
+   * offset read lock.
    */
-  private void disableCache() {
+  private synchronized void disableCache(boolean persistOnCleanup) {
     if (!isCacheEnabled()) {
       return;
     }
     LOG.info("Disabling cache");
     cacheState = CacheState.DISABLED;
     ioEngine.shutdown();
     this.scheduleThreadPool.shutdown();
-    for (int i = 0; i < writerThreads.length; ++i)
-      writerThreads[i].interrupt();
-    this.ramCache.clear();
-    if (!ioEngine.isPersistent() || persistencePath == null) {
-      // If persistent ioengine and a path, we will serialize out the backingMap.
-      this.backingMap.clear();
-      this.blocksByHFile.clear();
-      this.fullyCachedFiles.clear();
-      this.regionCachedSize.clear();
+    for (WriterThread writerThread : writerThreads) {
+      writerThread.interrupt();
     }
+    ramCache.clear();
     if (cacheStats.getMetricsRollerScheduler() != null) {
       cacheStats.getMetricsRollerScheduler().shutdownNow();
     }
+    cacheCleanupThread =
+      Threads.setDaemonThreadRunning(new Thread(() -> cleanupCache(persistOnCleanup)),
+        "BucketCacheCleanup-" + System.identityHashCode(this), Threads.LOGGING_EXCEPTION_HANDLER);
   }
 
-  private void join() throws InterruptedException {
-    for (int i = 0; i < writerThreads.length; ++i)
-      writerThreads[i].join();
-  }
-
-  @Override
-  public void shutdown() {
-    if (isCacheEnabled()) {
-      disableCache();
-      LOG.info("Shutdown bucket cache: IO persistent=" + ioEngine.isPersistent()
-        + "; path to write=" + persistencePath);
-      if (ioEngine.isPersistent() && persistencePath != null) {
+  private void cleanupCache(boolean persistOnCleanup) {
+    try {
+      joinWriterThreads();
+      stopCachePersister();
+      if (persistOnCleanup && isCachePersistent()) {
         try {
-          join();
-          if (cachePersister != null) {
-            LOG.info("Shutting down cache persister thread.");
-            cachePersister.shutdown();
-            while (cachePersister.isAlive()) {
-              Thread.sleep(10);
-            }
-          }
           persistToFile();
         } catch (IOException ex) {
           LOG.error("Unable to persist data on exit: " + ex.toString(), ex);
-        } catch (InterruptedException e) {
-          LOG.warn("Failed to persist data on exit", e);
         }
       }
+    } finally {
+      releaseBackingMapReferences();
+    }
+  }
+
+  private void joinWriterThreads() {
+    for (WriterThread writerThread : writerThreads) {
+      Threads.shutdown(writerThread);
+    }
+  }
+
+  private void stopCachePersister() {
+    if (cachePersister != null) {
+      LOG.info("Shutting down cache persister thread.");
+      cachePersister.shutdown();
+      Threads.shutdown(cachePersister);
+    }
+  }
+
+  private void releaseBackingMapReferences() {
+    for (Map.Entry<BlockCacheKey, BucketEntry> entry : backingMap.entrySet()) {
+      BlockCacheKey cacheKey = entry.getKey();
+      BucketEntry bucketEntry = entry.getValue();
+      bucketEntry.withWriteLock(offsetLock, () -> {
+        if (backingMap.remove(cacheKey, bucketEntry)) {
+          bucketEntry.markAsEvicted();
+        }
+        return null;
+      });
+    }
+    blocksByHFile.clear();
+    fullyCachedFiles.clear();
+    regionCachedSize.clear();
+  }
+
+  private void waitForCacheCleanup() {
+    Thread cleanupThread = cacheCleanupThread;
+    if (cleanupThread == null || cleanupThread == Thread.currentThread()) {
+      return;
+    }
+    boolean interrupted = false;
+    while (cleanupThread.isAlive()) {
+      try {
+        cleanupThread.join();
+      } catch (InterruptedException e) {
+        interrupted = true;
+      }
+    }
+    if (interrupted) {
+      Thread.currentThread().interrupt();
     }
   }
 
+  @Override
+  public void shutdown() {
+    disableCache(true);
+    waitForCacheCleanup();
+    LOG.info("Shutdown bucket cache: IO persistent=" + ioEngine.isPersistent() + "; path to write="
+      + persistencePath);
+  }
+
   /**
    * Needed mostly for UTs that might run in the same VM and create different BucketCache instances
    * on different UT methods.

hbase-server/src/test/java/org/apache/hadoop/hbase/io/hfile/TestBlockEvictionOnRegionMovement.java

@@ -27,6 +27,7 @@
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.Path;
 import org.apache.hadoop.hbase.HBaseTestingUtil;
+import org.apache.hadoop.hbase.ServerName;
 import org.apache.hadoop.hbase.SingleProcessHBaseCluster;
 import org.apache.hadoop.hbase.StartTestingClusterOption;
 import org.apache.hadoop.hbase.TableName;
@@ -126,20 +127,26 @@ public void testBlockEvictionOnGracefulStop() throws Exception {
         : cluster.getRegionServer(0);
 
     assertTrue(regionServingRS.getBlockCache().isPresent());
-    long oldUsedCacheSize =
-      regionServingRS.getBlockCache().get().getBlockCaches()[1].getCurrentSize();
-    assertNotEquals(0, regionServingRS.getBlockCache().get().getBlockCaches()[1].getBlockCount());
-
-    cluster.stopRegionServer(regionServingRS.getServerName());
-    Thread.sleep(500);
-    cluster.startRegionServer();
-    Thread.sleep(500);
-
-    regionServingRS.getBlockCache().get().waitForCacheInitialization(10000);
-    long newUsedCacheSize =
-      regionServingRS.getBlockCache().get().getBlockCaches()[1].getCurrentSize();
-    assertEquals(oldUsedCacheSize, newUsedCacheSize);
-    assertNotEquals(0, regionServingRS.getBlockCache().get().getBlockCaches()[1].getBlockCount());
+    BlockCache oldBucketCache = regionServingRS.getBlockCache().get().getBlockCaches()[1];
+    long oldUsedCacheSize = oldBucketCache.getCurrentSize();
+    assertNotEquals(0, oldUsedCacheSize);
+    assertNotEquals(0, oldBucketCache.getBlockCount());
+
+    ServerName serverName = regionServingRS.getServerName();
+    cluster.stopRegionServer(serverName);
+    cluster.waitForRegionServerToStop(serverName, 10000);
+
+    // Persistent state is restored into a new cache; the stopped cache must release its
+    // references.
+    assertEquals(0, oldBucketCache.getCurrentSize());
+
+    HRegionServer restartedRegionServer = cluster.startRegionServer().getRegionServer();
+    assertTrue(restartedRegionServer.getBlockCache().isPresent());
+    BlockCache restoredBucketCache =
+      restartedRegionServer.getBlockCache().get().getBlockCaches()[1];
+    assertTrue(restoredBucketCache.waitForCacheInitialization(10000));
+    assertEquals(oldUsedCacheSize, restoredBucketCache.getCurrentSize());
+    assertNotEquals(0, restoredBucketCache.getBlockCount());
   }
 
   public TableName writeDataToTable(String testName) throws IOException, InterruptedException {

hbase-server/src/test/java/org/apache/hadoop/hbase/io/hfile/bucket/TestBucketCache.java

@@ -841,10 +841,15 @@ public void testFreeBucketEntryRestoredFromFile() throws Exception {
         cacheAndWaitUntilFlushedToBucket(bucketCache, hfileBlockPair.getBlockName(),
           hfileBlockPair.getBlock(), false);
       }
+      BucketEntry bucketEntryBeforeShutdown =
+        bucketCache.backingMap.get(hfileBlockPairs[0].getBlockName());
+      assertNotNull(bucketEntryBeforeShutdown);
+      assertEquals(1, bucketEntryBeforeShutdown.refCnt());
       usedByteSize = bucketCache.getAllocator().getUsedSize();
       assertNotEquals(0, usedByteSize);
       // persist cache to file
       bucketCache.shutdown();
+      assertEquals(0, bucketEntryBeforeShutdown.refCnt());
       assertTrue(new File(persistencePath).exists());
 
       // restore cache from file