HBASE-26281 DBB got from BucketCache would be freed unexpectedly before RPC completed (#3680)

Signed-off-by: Duo Zhang <zhangduo@apache.org>

Author: comnetwork

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

@@ -25,8 +25,10 @@
 import java.util.Comparator;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Function;
 
 import org.apache.hadoop.hbase.io.ByteBuffAllocator;
+import org.apache.hadoop.hbase.io.ByteBuffAllocator.Recycler;
 import org.apache.hadoop.hbase.io.hfile.BlockPriority;
 import org.apache.hadoop.hbase.io.hfile.Cacheable;
 import org.apache.hadoop.hbase.io.hfile.CacheableDeserializer;
@@ -88,16 +90,21 @@ class BucketEntry implements HBaseReferenceCounted {
   private final long cachedTime = System.nanoTime();
 
   BucketEntry(long offset, int length, long accessCounter, boolean inMemory) {
-    this(offset, length, accessCounter, inMemory, RefCnt.create(), ByteBuffAllocator.HEAP);
+    this(offset, length, accessCounter, inMemory, null, ByteBuffAllocator.HEAP);
   }
 
-  BucketEntry(long offset, int length, long accessCounter, boolean inMemory, RefCnt refCnt,
+  BucketEntry(long offset, int length, long accessCounter, boolean inMemory,
+      Function<BucketEntry, Recycler> createRecycler,
       ByteBuffAllocator allocator) {
     setOffset(offset);
     this.length = length;
     this.accessCounter = accessCounter;
     this.priority = inMemory ? BlockPriority.MEMORY : BlockPriority.MULTI;
-    this.refCnt = refCnt;
+    if (createRecycler == null) {
+      this.refCnt = RefCnt.create();
+    } else {
+      this.refCnt = RefCnt.create(createRecycler.apply(this));
+    }
     this.markedAsEvicted = new AtomicBoolean(false);
     this.allocator = allocator;
   }
@@ -165,19 +172,6 @@ boolean markAsEvicted() {
     return false;
   }
 
-  /**
-   * Mark as evicted only when NO RPC references. Mainly used for eviction when cache size exceed
-   * the max acceptable size.
-   * @return true if we deallocate this entry successfully.
-   */
-  boolean markStaleAsEvicted() {
-    if (!markedAsEvicted.get() && this.refCnt() == 1) {
-      // The only reference was coming from backingMap, now release the stale entry.
-      return this.markAsEvicted();
-    }
-    return false;
-  }
-
   /**
    * Check whether have some RPC patch referring this block. There're two case: <br>
    * 1. If current refCnt is greater than 1, there must be at least one referring RPC path; <br>

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

@@ -18,6 +18,7 @@
 package org.apache.hadoop.hbase.client;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -43,6 +44,8 @@
 import org.apache.hadoop.hbase.io.hfile.BlockCacheKey;
 import org.apache.hadoop.hbase.io.hfile.CacheConfig;
 import org.apache.hadoop.hbase.io.hfile.CachedBlock;
+import org.apache.hadoop.hbase.io.hfile.CombinedBlockCache;
+import org.apache.hadoop.hbase.io.hfile.bucket.BucketCache;
 import org.apache.hadoop.hbase.regionserver.DelegatingInternalScanner;
 import org.apache.hadoop.hbase.regionserver.HRegion;
 import org.apache.hadoop.hbase.regionserver.HStore;
@@ -372,7 +375,11 @@ public void run() {
               CachedBlock next = iterator.next();
               BlockCacheKey cacheKey = new BlockCacheKey(next.getFilename(), next.getOffset());
               cacheList.add(cacheKey);
-              cache.evictBlock(cacheKey);
+              /**
+               * There is only one Block referenced by rpc,here we evict blocks which have no rpc
+               * referenced.
+               */
+              evictBlock(cache, cacheKey);
             }
             try {
               Thread.sleep(1);
@@ -437,4 +444,20 @@ public void run() {
       table.close();
     }
   }
+
+  /**
+   * For {@link BucketCache},we only evict Block if there is no rpc referenced.
+   */
+  private void evictBlock(BlockCache blockCache, BlockCacheKey blockCacheKey) {
+    assertTrue(blockCache instanceof CombinedBlockCache);
+    BlockCache[] blockCaches = blockCache.getBlockCaches();
+    for (BlockCache currentBlockCache : blockCaches) {
+      if (currentBlockCache instanceof BucketCache) {
+        ((BucketCache) currentBlockCache).evictBlockIfNoRpcReferenced(blockCacheKey);
+      } else {
+        currentBlockCache.evictBlock(blockCacheKey);
+      }
+    }
+
+  }
 }

hbase-server/src/test/java/org/apache/hadoop/hbase/client/TestAvoidCellReferencesIntoShippedBlocks.java

@@ -249,8 +249,26 @@ public void run() {
     assertEquals(1, cache.getBlockCount());
     lock.writeLock().unlock();
     evictThread.join();
-    assertEquals(0, cache.getBlockCount());
-    assertEquals(cache.getCurrentSize(), 0L);
+    /**
+     * <pre>
+     * The asserts here before HBASE-21957 are:
+     * assertEquals(1L, cache.getBlockCount());
+     * assertTrue(cache.getCurrentSize() > 0L);
+     * assertTrue("We should have a block!", cache.iterator().hasNext());
+     *
+     * The asserts here after HBASE-21957 are:
+     * assertEquals(0, cache.getBlockCount());
+     * assertEquals(cache.getCurrentSize(), 0L);
+     *
+     * I think the asserts before HBASE-21957 is more reasonable,because
+     * {@link BucketCache#evictBlock} should only evict the {@link BucketEntry}
+     * it had seen, and newly added Block after the {@link BucketEntry}
+     * it had seen should not be evicted.
+     * </pre>
+     */
+    assertEquals(1L, cache.getBlockCount());
+    assertTrue(cache.getCurrentSize() > 0L);
+    assertTrue("We should have a block!", cache.iterator().hasNext());
   }
 
   @Test
@@ -503,8 +521,8 @@ public void testRAMCache() {
         HFileBlock.FILL_HEADER, -1, 52, -1, meta, ByteBuffAllocator.HEAP);
     HFileBlock blk2 = new HFileBlock(BlockType.DATA, size, size, -1, ByteBuff.wrap(buf),
         HFileBlock.FILL_HEADER, -1, -1, -1, meta, ByteBuffAllocator.HEAP);
-    RAMQueueEntry re1 = new RAMQueueEntry(key1, blk1, 1, false, ByteBuffAllocator.NONE);
-    RAMQueueEntry re2 = new RAMQueueEntry(key1, blk2, 1, false, ByteBuffAllocator.NONE);
+    RAMQueueEntry re1 = new RAMQueueEntry(key1, blk1, 1, false);
+    RAMQueueEntry re2 = new RAMQueueEntry(key1, blk2, 1, false);
 
     assertFalse(cache.containsKey(key1));
     assertNull(cache.putIfAbsent(key1, re1));
@@ -551,11 +569,11 @@ public void testFreeBlockWhenIOEngineWriteFailure() throws IOException {
     BucketAllocator allocator = new BucketAllocator(availableSpace, null);
 
     BlockCacheKey key = new BlockCacheKey("dummy", 1L);
-    RAMQueueEntry re = new RAMQueueEntry(key, block, 1, true, ByteBuffAllocator.NONE);
+    RAMQueueEntry re = new RAMQueueEntry(key, block, 1, true);
 
     Assert.assertEquals(0, allocator.getUsedSize());
     try {
-      re.writeToCache(ioEngine, allocator, null);
+      re.writeToCache(ioEngine, allocator, null, null);
       Assert.fail();
     } catch (Exception e) {
     }