HADOOP-18291. S3A prefetch - Implement thread-safe LRU cache for SingleFilePerBlockCache

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/impl/prefetch/CachingBlockManager.java

@@ -129,6 +129,7 @@ public CachingBlockManager(
     this.numReadErrors = new AtomicInteger();
     this.cachingDisabled = new AtomicBoolean();
     this.prefetchingStatistics = requireNonNull(prefetchingStatistics);
+    this.conf = requireNonNull(conf);
 
     if (this.getBlockData().getFileSize() > 0) {
       this.bufferPool = new BufferPool(bufferPoolSize, this.getBlockData().getBlockSize(),
@@ -138,7 +139,6 @@ public CachingBlockManager(
 
     this.ops = new BlockOperations();
     this.ops.setDebug(false);
-    this.conf = requireNonNull(conf);
     this.localDirAllocator = localDirAllocator;
   }
 
@@ -558,7 +558,7 @@ private void addToCacheAndRelease(BufferData data, Future<Void> blockFuture,
   }
 
   protected BlockCache createCache() {
-    return new SingleFilePerBlockCache(prefetchingStatistics);
+    return new SingleFilePerBlockCache(prefetchingStatistics, conf);
   }
 
   protected void cachePut(int blockNumber, ByteBuffer buffer) throws IOException {

hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/fs/impl/prefetch/SingleFilePerBlockCache.java

@@ -47,6 +47,7 @@
 
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.LocalDirAllocator;
+import org.apache.hadoop.util.Preconditions;
 
 import static java.util.Objects.requireNonNull;
 import static org.apache.hadoop.fs.impl.prefetch.Validate.checkNotNull;
@@ -61,7 +62,27 @@ public class SingleFilePerBlockCache implements BlockCache {
   /**
    * Blocks stored in this cache.
    */
-  private final Map<Integer, Entry> blocks = new ConcurrentHashMap<>();
+  private final Map<Integer, Entry> blocks;
+
+  /**
+   * Total max blocks count, to be considered as baseline for LRU cache.
+   */
+  private final int maxBlocksCount;
+
+  /**
+   * The lock to be shared by LRU based linked list updates.
+   */
+  private final ReentrantReadWriteLock blocksLock;
+
+  /**
+   * Head of the linked list.
+   */
+  private Entry head;
+
+  /**
+   * Tail of the lined list.
+   */
+  private Entry tail;
 
   /**
    * Number of times a block was read from this cache.
@@ -89,6 +110,16 @@ public class SingleFilePerBlockCache implements BlockCache {
   private static final Set<PosixFilePermission> TEMP_FILE_ATTRS =
       ImmutableSet.of(PosixFilePermission.OWNER_READ, PosixFilePermission.OWNER_WRITE);
 
+  /**
+   * Prefetch max blocks count config.
+   */
+  public static final String FS_PREFETCH_MAX_BLOCKS_COUNT = "fs.prefetch.max.blocks.count";
+
+  /**
+   * Default value for max blocks count config.
+   */
+  private static final int DEFAULT_FS_PREFETCH_MAX_BLOCKS_COUNT = 20;
+
   /**
    * Cache entry.
    * Each block is stored as a separate file.
@@ -103,13 +134,17 @@ private enum LockType {
       READ,
       WRITE
     }
+    private Entry previous;
+    private Entry next;
 
     Entry(int blockNumber, Path path, int size, long checksum) {
       this.blockNumber = blockNumber;
       this.path = path;
       this.size = size;
       this.checksum = checksum;
       this.lock = new ReentrantReadWriteLock();
+      this.previous = null;
+      this.next = null;
     }
 
     @Override
@@ -166,16 +201,39 @@ private boolean takeLock(LockType lockType, long timeout, TimeUnit unit) {
       }
       return false;
     }
+
+    private Entry getPrevious() {
+      return previous;
+    }
+
+    private void setPrevious(Entry previous) {
+      this.previous = previous;
+    }
+
+    private Entry getNext() {
+      return next;
+    }
+
+    private void setNext(Entry next) {
+      this.next = next;
+    }
   }
 
   /**
    * Constructs an instance of a {@code SingleFilePerBlockCache}.
    *
    * @param prefetchingStatistics statistics for this stream.
+   * @param conf the configuration object.
    */
-  public SingleFilePerBlockCache(PrefetchingStatistics prefetchingStatistics) {
+  public SingleFilePerBlockCache(PrefetchingStatistics prefetchingStatistics, Configuration conf) {
     this.prefetchingStatistics = requireNonNull(prefetchingStatistics);
     this.closed = new AtomicBoolean(false);
+    this.maxBlocksCount =
+        conf.getInt(FS_PREFETCH_MAX_BLOCKS_COUNT, DEFAULT_FS_PREFETCH_MAX_BLOCKS_COUNT);
+    Preconditions.checkArgument(this.maxBlocksCount > 0,
+        "prefetch blocks total capacity should be more than 0");
+    blocks = new ConcurrentHashMap<>();
+    blocksLock = new ReentrantReadWriteLock();
   }
 
   /**
@@ -247,9 +305,46 @@ private Entry getEntry(int blockNumber) {
       throw new IllegalStateException(String.format("block %d not found in cache", blockNumber));
     }
     numGets++;
+    addToHeadOfLinkedList(entry);
     return entry;
   }
 
+  /**
+   * Add the given entry to the head of the linked list.
+   *
+   * @param entry Block entry to add.
+   */
+  private void addToHeadOfLinkedList(Entry entry) {
+    blocksLock.writeLock().lock();
+    try {
+      if (head == null) {
+        head = entry;
+        tail = entry;
+      }
+      if (entry != head) {
+        Entry prev = entry.getPrevious();
+        Entry nxt = entry.getNext();
+        if (prev != null) {
+          prev.setNext(nxt);
+        }
+        if (nxt != null) {
+          nxt.setPrevious(prev);
+        }
+        entry.setPrevious(null);
+        entry.setNext(head);
+        head.setPrevious(entry);
+        head = entry;
+      }
+      if (tail != null) {
+        while (tail.getNext() != null) {
+          tail = tail.getNext();
+        }
+      }
+    } finally {
+      blocksLock.writeLock().unlock();
+    }
+  }
+
   /**
    * Puts the given block in this cache.
    *
@@ -278,6 +373,7 @@ public void put(int blockNumber, ByteBuffer buffer, Configuration conf,
       } finally {
         entry.releaseLock(Entry.LockType.READ);
       }
+      addToHeadOfLinkedList(entry);
       return;
     }
 
@@ -299,9 +395,62 @@ public void put(int blockNumber, ByteBuffer buffer, Configuration conf,
     // Update stream_read_blocks_in_cache stats only after blocks map is updated with new file
     // entry to avoid any discrepancy related to the value of stream_read_blocks_in_cache.
     // If stream_read_blocks_in_cache is updated before updating the blocks map here, closing of
-    // the input stream can lead to the removal of the cache file even before blocks is added with
-    // the new cache file, leading to incorrect value of stream_read_blocks_in_cache.
+    // the input stream can lead to the removal of the cache file even before blocks is added
+    // with the new cache file, leading to incorrect value of stream_read_blocks_in_cache.
     prefetchingStatistics.blockAddedToFileCache();
+    addToLinkedListAndEvictIfRequired(entry);
+  }
+
+  /**
+   * Add the given entry to the head of the linked list and if the LRU cache size
+   * exceeds the max limit, evict tail of the LRU linked list.
+   *
+   * @param entry Block entry to add.
+   */
+  private void addToLinkedListAndEvictIfRequired(Entry entry) {
+    addToHeadOfLinkedList(entry);
+    blocksLock.writeLock().lock();
+    try {
+      if (blocks.size() > maxBlocksCount && !closed.get()) {
+        Entry elementToPurge = tail;
+        tail = tail.getPrevious();
+        if (tail == null) {
+          tail = head;
+        }
+        tail.setNext(null);
+        elementToPurge.setPrevious(null);
+        deleteBlockFileAndEvictCache(elementToPurge);
+      }
+    } finally {
+      blocksLock.writeLock().unlock();
+    }
+  }
+
+  /**
+   * Delete cache file as part of the block cache LRU eviction.
+   *
+   * @param elementToPurge Block entry to evict.
+   */
+  private void deleteBlockFileAndEvictCache(Entry elementToPurge) {
+    boolean lockAcquired =
+        elementToPurge.takeLock(Entry.LockType.WRITE, PREFETCH_WRITE_LOCK_TIMEOUT,
+            PREFETCH_WRITE_LOCK_TIMEOUT_UNIT);
+    if (!lockAcquired) {
+      LOG.error("Cache file {} deletion would not be attempted as write lock could not"
+              + " be acquired within {} {}", elementToPurge.path, PREFETCH_WRITE_LOCK_TIMEOUT,
+          PREFETCH_WRITE_LOCK_TIMEOUT_UNIT);
+    } else {
+      try {
+        if (Files.deleteIfExists(elementToPurge.path)) {
+          prefetchingStatistics.blockRemovedFromFileCache();
+          blocks.remove(elementToPurge.blockNumber);
+        }
+      } catch (IOException e) {
+        LOG.warn("Failed to delete cache file {}", elementToPurge.path, e);
+      } finally {
+        elementToPurge.releaseLock(Entry.LockType.WRITE);
+      }
+    }
   }
 
   private static final Set<? extends OpenOption> CREATE_OPTIONS =
@@ -337,30 +486,36 @@ protected Path getCacheFilePath(final Configuration conf,
   public void close() throws IOException {
     if (closed.compareAndSet(false, true)) {
       LOG.debug(getStats());
-      int numFilesDeleted = 0;
+      deleteCacheFiles();
+    }
+  }
 
-      for (Entry entry : blocks.values()) {
-        boolean lockAcquired = entry.takeLock(Entry.LockType.WRITE, PREFETCH_WRITE_LOCK_TIMEOUT,
+  /**
+   * Delete cache files as part of the close call.
+   */
+  private void deleteCacheFiles() {
+    int numFilesDeleted = 0;
+    for (Entry entry : blocks.values()) {
+      boolean lockAcquired = entry.takeLock(Entry.LockType.WRITE, PREFETCH_WRITE_LOCK_TIMEOUT,
+          PREFETCH_WRITE_LOCK_TIMEOUT_UNIT);
+      if (!lockAcquired) {
+        LOG.error("Cache file {} deletion would not be attempted as write lock could not"
+                + " be acquired within {} {}", entry.path, PREFETCH_WRITE_LOCK_TIMEOUT,
             PREFETCH_WRITE_LOCK_TIMEOUT_UNIT);
-        if (!lockAcquired) {
-          LOG.error("Cache file {} deletion would not be attempted as write lock could not"
-                  + " be acquired within {} {}", entry.path, PREFETCH_WRITE_LOCK_TIMEOUT,
-              PREFETCH_WRITE_LOCK_TIMEOUT_UNIT);
-          continue;
-        }
-        try {
-          Files.deleteIfExists(entry.path);
+        continue;
+      }
+      try {
+        if (Files.deleteIfExists(entry.path)) {
           prefetchingStatistics.blockRemovedFromFileCache();
           numFilesDeleted++;
-        } catch (IOException e) {
-          LOG.warn("Failed to delete cache file {}", entry.path, e);
-        } finally {
-          entry.releaseLock(Entry.LockType.WRITE);
         }
+      } catch (IOException e) {
+        LOG.warn("Failed to delete cache file {}", entry.path, e);
+      } finally {
+        entry.releaseLock(Entry.LockType.WRITE);
       }
-
-      LOG.debug("Prefetch cache close: Deleted {} cache files", numFilesDeleted);
     }
+    LOG.debug("Prefetch cache close: Deleted {} cache files", numFilesDeleted);
   }
 
   @Override

hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/fs/impl/prefetch/TestBlockCache.java

@@ -45,7 +45,7 @@ public class TestBlockCache extends AbstractHadoopTestBase {
   public void testArgChecks() throws Exception {
     // Should not throw.
     BlockCache cache =
-        new SingleFilePerBlockCache(EmptyPrefetchingStatistics.getInstance());
+        new SingleFilePerBlockCache(EmptyPrefetchingStatistics.getInstance(), CONF);
 
     ByteBuffer buffer = ByteBuffer.allocate(16);
 
@@ -55,15 +55,15 @@ public void testArgChecks() throws Exception {
 
 
     intercept(NullPointerException.class, null,
-        () -> new SingleFilePerBlockCache(null));
+        () -> new SingleFilePerBlockCache(null, CONF));
 
   }
 
 
   @Test
   public void testPutAndGet() throws Exception {
     BlockCache cache =
-        new SingleFilePerBlockCache(EmptyPrefetchingStatistics.getInstance());
+        new SingleFilePerBlockCache(EmptyPrefetchingStatistics.getInstance(), CONF);
 
     ByteBuffer buffer1 = ByteBuffer.allocate(BUFFER_SIZE);
     for (byte i = 0; i < BUFFER_SIZE; i++) {