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

Author: mwkang

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;
@@ -160,6 +161,7 @@ public class BucketCache implements BlockCache, HeapSize {
 
   final static int DEFAULT_WRITER_THREADS = 3;
   final static int DEFAULT_WRITER_QUEUE_ITEMS = 64;
+  private static final long WRITER_QUEUE_POLL_INTERVAL_MS = 100;
 
   final static long DEFAULT_BACKING_MAP_PERSISTENCE_CHUNK_SIZE = 10000;
 
@@ -206,6 +208,16 @@ protected enum CacheState {
    */
   private volatile CacheState cacheState;
 
+  /**
+   * Prevent publishing new RAM cache entries while the cache is transitioning to disabled. The read
+   * lock is only held for the RAM cache and writer queue publication path, so it is never nested
+   * with an offset lock.
+   */
+  private final ReentrantReadWriteLock cacheLifecycleLock = new ReentrantReadWriteLock();
+
+  /** 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
@@ -589,47 +601,76 @@ protected boolean shouldReplaceExistingCacheBlock(BlockCacheKey cacheKey, Cachea
 
   protected void cacheBlockWithWaitInternal(BlockCacheKey cacheKey, Cacheable cachedItem,
     boolean inMemory, boolean wait) {
-    if (!isCacheEnabled()) {
-      return;
-    }
-    if (cacheKey.getBlockType() == null && cachedItem.getBlockType() != null) {
-      cacheKey.setBlockType(cachedItem.getBlockType());
-    }
-    LOG.debug("Caching key={}, item={}, key heap size={}", cacheKey, cachedItem,
-      cacheKey.heapSize());
-    // Stuff the entry into the RAM cache so it can get drained to the persistent store
-    RAMQueueEntry re = new RAMQueueEntry(cacheKey, cachedItem, accessCount.incrementAndGet(),
-      inMemory, isCachePersistent() && ioEngine instanceof FileIOEngine, wait);
-    /**
-     * Don't use ramCache.put(cacheKey, re) here. because there may be a existing entry with same
-     * key in ramCache, the heap size of bucket cache need to update if replacing entry from
-     * ramCache. But WriterThread will also remove entry from ramCache and update heap size, if
-     * using ramCache.put(), It's possible that the removed entry in WriterThread is not the correct
-     * one, then the heap size will mess up (HBASE-20789)
-     */
-    if (ramCache.putIfAbsent(cacheKey, re) != null) {
-      return;
+    Lock lifecycleReadLock = cacheLifecycleLock.readLock();
+    RAMQueueEntry re;
+    lifecycleReadLock.lock();
+    try {
+      if (!isCacheEnabled()) {
+        return;
+      }
+      if (cacheKey.getBlockType() == null && cachedItem.getBlockType() != null) {
+        cacheKey.setBlockType(cachedItem.getBlockType());
+      }
+      LOG.debug("Caching key={}, item={}, key heap size={}", cacheKey, cachedItem,
+        cacheKey.heapSize());
+      // Stuff the entry into the RAM cache so it can get drained to the persistent store
+      re = new RAMQueueEntry(cacheKey, cachedItem, accessCount.incrementAndGet(), inMemory,
+        isCachePersistent() && ioEngine instanceof FileIOEngine, wait);
+      /**
+       * Don't use ramCache.put(cacheKey, re) here. because there may be a existing entry with same
+       * key in ramCache, the heap size of bucket cache need to update if replacing entry from
+       * ramCache. But WriterThread will also remove entry from ramCache and update heap size, if
+       * using ramCache.put(), It's possible that the removed entry in WriterThread is not the
+       * correct one, then the heap size will mess up (HBASE-20789)
+       */
+      if (ramCache.putIfAbsent(cacheKey, re) != null) {
+        return;
+      }
+    } finally {
+      lifecycleReadLock.unlock();
     }
+
     int queueNum = (cacheKey.hashCode() & 0x7FFFFFFF) % writerQueues.size();
     BlockingQueue<RAMQueueEntry> bq = writerQueues.get(queueNum);
     boolean successfulAddition = false;
-    if (wait) {
-      try {
-        successfulAddition = bq.offer(re, queueAdditionWaitTime, TimeUnit.MILLISECONDS);
-      } catch (InterruptedException e) {
-        LOG.error("Thread interrupted: ", e);
-        Thread.currentThread().interrupt();
-      }
-    } else {
-      successfulAddition = bq.offer(re);
+    try {
+      successfulAddition = offerToWriterQueue(bq, re, cachedItem, wait);
+    } catch (InterruptedException e) {
+      LOG.error("Thread interrupted: ", e);
+      Thread.currentThread().interrupt();
     }
     if (!successfulAddition) {
       LOG.debug("Failed to insert block {} into the cache writers queue", cacheKey);
       ramCache.remove(cacheKey);
       cacheStats.failInsert();
-    } else {
-      this.blockNumber.increment();
-      this.heapSize.add(cachedItem.heapSize());
+    }
+  }
+
+  private boolean offerToWriterQueue(BlockingQueue<RAMQueueEntry> queue, RAMQueueEntry entry,
+    Cacheable cachedItem, boolean wait) throws InterruptedException {
+    long timeoutNanos = TimeUnit.MILLISECONDS.toNanos(queueAdditionWaitTime);
+    long startTime = System.nanoTime();
+    while (true) {
+      Lock lifecycleReadLock = cacheLifecycleLock.readLock();
+      lifecycleReadLock.lock();
+      try {
+        if (!isCacheEnabled()) {
+          return false;
+        }
+        if (queue.offer(entry)) {
+          blockNumber.increment();
+          heapSize.add(cachedItem.heapSize());
+          return true;
+        }
+      } finally {
+        lifecycleReadLock.unlock();
+      }
+      if (!wait || System.nanoTime() - startTime >= timeoutNanos) {
+        return false;
+      }
+      long remainingNanos = timeoutNanos - (System.nanoTime() - startTime);
+      TimeUnit.NANOSECONDS.sleep(
+        Math.min(remainingNanos, TimeUnit.MILLISECONDS.toNanos(WRITER_QUEUE_POLL_INTERVAL_MS)));
     }
   }
 
@@ -690,6 +731,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,15 +764,17 @@ 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) {
+        removeBucketEntry(key, bucketEntry, true, false);
+      }
     }
     if (!repeat && updateCacheMetrics) {
       cacheStats.miss(caching, key.isPrimary(), key.getBlockType());
@@ -849,18 +893,24 @@ private boolean doEvictBlock(BlockCacheKey cacheKey, BucketEntry bucketEntry,
       }
       return existedInRamCache;
     } else {
-      return bucketEntryToUse.withWriteLock(offsetLock, () -> {
-        if (backingMap.remove(cacheKey, bucketEntryToUse)) {
-          LOG.debug("removed key {} from back map with offset lock {} in the evict process",
-            cacheKey, bucketEntryToUse.offset());
-          blockEvicted(cacheKey, bucketEntryToUse, !existedInRamCache, evictedByEvictionProcess);
-          return true;
-        }
-        return false;
-      });
+      return removeBucketEntry(cacheKey, bucketEntryToUse, !existedInRamCache,
+        evictedByEvictionProcess);
     }
   }
 
+  private boolean removeBucketEntry(BlockCacheKey cacheKey, BucketEntry bucketEntry,
+    boolean decrementBlockNumber, boolean evictedByEvictionProcess) {
+    return bucketEntry.withWriteLock(offsetLock, () -> {
+      if (backingMap.remove(cacheKey, bucketEntry)) {
+        LOG.debug("removed key {} from back map with offset lock {} in the evict process", cacheKey,
+          bucketEntry.offset());
+        blockEvicted(cacheKey, bucketEntry, decrementBlockNumber, evictedByEvictionProcess);
+        return true;
+      }
+      return false;
+    });
+  }
+
   /**
    * <pre>
    * Create the {@link Recycler} for {@link BucketEntry#refCnt},which would be used as
@@ -1317,13 +1367,15 @@ public void run() {
         while (isCacheEnabled() && writerEnabled) {
           try {
             try {
-              // Blocks
-              entries = getRAMQueueEntries(inputQueue, entries);
+              entries = pollRAMQueueEntries(inputQueue, entries);
             } catch (InterruptedException ie) {
               if (!isCacheEnabled() || !writerEnabled) {
                 break;
               }
             }
+            if (entries.isEmpty()) {
+              continue;
+            }
             doDrain(entries, metaBuff);
           } catch (Exception ioe) {
             LOG.error("WriterThread encountered error", ioe);
@@ -1545,6 +1597,17 @@ static List<RAMQueueEntry> getRAMQueueEntries(BlockingQueue<RAMQueueEntry> q,
     return receptacle;
   }
 
+  private static List<RAMQueueEntry> pollRAMQueueEntries(BlockingQueue<RAMQueueEntry> q,
+    List<RAMQueueEntry> receptacle) throws InterruptedException {
+    receptacle.clear();
+    RAMQueueEntry entry = q.poll(WRITER_QUEUE_POLL_INTERVAL_MS, TimeUnit.MILLISECONDS);
+    if (entry != null) {
+      receptacle.add(entry);
+      q.drainTo(receptacle);
+    }
+    return receptacle;
+  }
+
   /**
    * @see #retrieveFromFile(int[])
    */
@@ -1809,59 +1872,107 @@ private void checkIOErrorIsTolerated() {
    * Used to shut down the cache -or- turn it off in the case of something broken.
    */
   private void disableCache() {
-    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();
-    }
-    if (cacheStats.getMetricsRollerScheduler() != null) {
-      cacheStats.getMetricsRollerScheduler().shutdownNow();
+    Lock lifecycleWriteLock = cacheLifecycleLock.writeLock();
+    lifecycleWriteLock.lock();
+    try {
+      if (!isCacheEnabled()) {
+        return;
+      }
+      LOG.info("Disabling cache");
+      cacheState = CacheState.DISABLED;
+      this.scheduleThreadPool.shutdown();
+      if (cacheStats.getMetricsRollerScheduler() != null) {
+        cacheStats.getMetricsRollerScheduler().shutdownNow();
+      }
+      cacheCleanupThread = Threads.setDaemonThreadRunning(new Thread(this::cleanupCache),
+        "BucketCacheCleanup-" + System.identityHashCode(this), Threads.LOGGING_EXCEPTION_HANDLER);
+    } finally {
+      lifecycleWriteLock.unlock();
     }
   }
 
-  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() {
+    try {
+      joinWriterThreads();
+      ramCache.clear();
+      for (BlockingQueue<RAMQueueEntry> writerQueue : writerQueues) {
+        writerQueue.clear();
+      }
+      stopCachePersister();
+      if (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 {
+      try {
+        releaseBackingMapReferences();
+      } finally {
+        ioEngine.shutdown();
+      }
     }
   }
 
+  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();
+    blockNumber.reset();
+    heapSize.reset();
+  }
+
+  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();
+    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.