HADOOP-19622: [ABFS][ReadAheadV2] Implement Read Buffer Manager V2 with improved aggressiveness

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/AbfsConfiguration.java

@@ -391,6 +391,11 @@ public class AbfsConfiguration{
       DefaultValue = DEFAULT_ENABLE_READAHEAD_V2)
   private boolean isReadAheadV2Enabled;
 
+  @BooleanConfigurationValidatorAnnotation(
+      ConfigurationKey = FS_AZURE_ENABLE_READAHEAD_V2_DYNAMIC_SCALING,
+      DefaultValue = DEFAULT_ENABLE_READAHEAD_V2_DYNAMIC_SCALING)
+  private boolean isReadAheadV2DynamicScalingEnabled;
+
   @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
       FS_AZURE_READAHEAD_V2_MIN_THREAD_POOL_SIZE,
       DefaultValue = DEFAULT_READAHEAD_V2_MIN_THREAD_POOL_SIZE)
@@ -411,6 +416,26 @@ public class AbfsConfiguration{
       DefaultValue = DEFAULT_READAHEAD_V2_MAX_BUFFER_POOL_SIZE)
   private int maxReadAheadV2BufferPoolSize;
 
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_CPU_MONITORING_INTERVAL_MILLIS,
+      DefaultValue = DEFAULT_READAHEAD_V2_CPU_MONITORING_INTERVAL_MILLIS)
+  private int readAheadV2CpuMonitoringIntervalMillis;
+
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_THREAD_POOL_UPSCALE_PERCENTAGE,
+      DefaultValue = DEFAULT_READAHEAD_V2_THREAD_POOL_UPSCALE_PERCENTAGE)
+  private int readAheadV2ThreadPoolUpscalePercentage;
+
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_THREAD_POOL_DOWNSCALE_PERCENTAGE,
+      DefaultValue = DEFAULT_READAHEAD_V2_THREAD_POOL_DOWNSCALE_PERCENTAGE)
+  private int readAheadV2ThreadPoolDownscalePercentage;
+
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_MEMORY_MONITORING_INTERVAL_MILLIS,
+      DefaultValue = DEFAULT_READAHEAD_V2_MEMORY_MONITORING_INTERVAL_MILLIS)
+  private int readAheadV2MemoryMonitoringIntervalMillis;
+
   @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
       FS_AZURE_READAHEAD_V2_EXECUTOR_SERVICE_TTL_MILLIS,
       DefaultValue = DEFAULT_READAHEAD_V2_EXECUTOR_SERVICE_TTL_MILLIS)
@@ -421,6 +446,16 @@ public class AbfsConfiguration{
       DefaultValue = DEFAULT_READAHEAD_V2_CACHED_BUFFER_TTL_MILLIS)
   private int readAheadV2CachedBufferTTLMillis;
 
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_CPU_USAGE_THRESHOLD_PERCENT,
+      DefaultValue = DEFAULT_READAHEAD_V2_CPU_USAGE_THRESHOLD_PERCENT)
+  private int readAheadV2CpuUsageThresholdPercent;
+
+  @IntegerConfigurationValidatorAnnotation(ConfigurationKey =
+      FS_AZURE_READAHEAD_V2_MEMORY_USAGE_THRESHOLD_PERCENT,
+      DefaultValue = DEFAULT_READAHEAD_V2_MEMORY_USAGE_THRESHOLD_PERCENT)
+  private int readAheadV2MemoryUsageThresholdPercent;
+
   @LongConfigurationValidatorAnnotation(ConfigurationKey = FS_AZURE_SAS_TOKEN_RENEW_PERIOD_FOR_STREAMS,
       MinValue = 0,
       DefaultValue = DEFAULT_SAS_TOKEN_RENEW_PERIOD_FOR_STREAMS_IN_SECONDS)
@@ -1416,6 +1451,18 @@ public boolean isReadAheadEnabled() {
     return this.enabledReadAhead;
   }
 
+  /**
+   * Checks if the read-ahead v2 feature is enabled by user.
+   * @return true if read-ahead v2 is enabled, false otherwise.
+   */
+  public boolean isReadAheadV2Enabled() {
+    return this.isReadAheadV2Enabled;
+  }
+
+  public boolean isReadAheadV2DynamicScalingEnabled() {
+    return isReadAheadV2DynamicScalingEnabled;
+  }
+
   public int getMinReadAheadV2ThreadPoolSize() {
     if (minReadAheadV2ThreadPoolSize <= 0) {
       // If the minReadAheadV2ThreadPoolSize is not set, use the default value
@@ -1448,6 +1495,22 @@ public int getMaxReadAheadV2BufferPoolSize() {
     return maxReadAheadV2BufferPoolSize;
   }
 
+  public int getReadAheadV2CpuMonitoringIntervalMillis() {
+    return readAheadV2CpuMonitoringIntervalMillis;
+  }
+
+  public int getReadAheadV2ThreadPoolUpscalePercentage() {
+    return readAheadV2ThreadPoolUpscalePercentage;
+  }
+
+  public int getReadAheadV2ThreadPoolDownscalePercentage() {
+    return readAheadV2ThreadPoolDownscalePercentage;
+  }
+
+  public int getReadAheadV2MemoryMonitoringIntervalMillis() {
+    return readAheadV2MemoryMonitoringIntervalMillis;
+  }
+
   public int getReadAheadExecutorServiceTTLInMillis() {
     return readAheadExecutorServiceTTLMillis;
   }
@@ -1456,12 +1519,12 @@ public int getReadAheadV2CachedBufferTTLMillis() {
     return readAheadV2CachedBufferTTLMillis;
   }
 
-  /**
-   * Checks if the read-ahead v2 feature is enabled by user.
-   * @return true if read-ahead v2 is enabled, false otherwise.
-   */
-  public boolean isReadAheadV2Enabled() {
-    return this.isReadAheadV2Enabled;
+  public int getReadAheadV2CpuUsageThresholdPercent() {
+    return readAheadV2CpuUsageThresholdPercent;
+  }
+
+  public int getReadAheadV2MemoryUsageThresholdPercent() {
+    return readAheadV2MemoryUsageThresholdPercent;
   }
 
   @VisibleForTesting

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/AzureBlobFileSystemStore.java

@@ -110,6 +110,7 @@
 import org.apache.hadoop.fs.azurebfs.services.AuthType;
 import org.apache.hadoop.fs.azurebfs.services.ExponentialRetryPolicy;
 import org.apache.hadoop.fs.azurebfs.services.ListingSupport;
+import org.apache.hadoop.fs.azurebfs.services.ReadBufferManagerV2;
 import org.apache.hadoop.fs.azurebfs.services.SharedKeyCredentials;
 import org.apache.hadoop.fs.azurebfs.services.StaticRetryPolicy;
 import org.apache.hadoop.fs.azurebfs.services.VersionedFileStatus;
@@ -324,6 +325,12 @@ public void close() throws IOException {
       HadoopExecutors.shutdown(boundedThreadPool, LOG,
           30, TimeUnit.SECONDS);
       boundedThreadPool = null;
+      if (getAbfsConfiguration().isReadAheadV2Enabled()) {
+        ReadBufferManagerV2 bufferManagerV2 = ReadBufferManagerV2.getInstance();
+        if (bufferManagerV2 != null) {
+          bufferManagerV2.close();
+        }
+      }
     } catch (InterruptedException e) {
       LOG.error("Interrupted freeing leases", e);
       Thread.currentThread().interrupt();

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/constants/ConfigurationKeys.java

@@ -269,6 +269,12 @@ public final class ConfigurationKeys {
    */
   public static final String FS_AZURE_ENABLE_READAHEAD_V2 = "fs.azure.enable.readahead.v2";
 
+  /**
+   * Enable or disable dynamic scaling of thread pool and buffer pool of readahead V2.
+   * Value: {@value}.
+   */
+  public static final String FS_AZURE_ENABLE_READAHEAD_V2_DYNAMIC_SCALING = "fs.azure.enable.readahead.v2.dynamic.scaling";
+
   /**
    * Minimum number of prefetch threads in the thread pool for readahead V2.
    * {@value }
@@ -290,6 +296,28 @@ public final class ConfigurationKeys {
    */
   public static final String FS_AZURE_READAHEAD_V2_MAX_BUFFER_POOL_SIZE = "fs.azure.readahead.v2.max.buffer.pool.size";
 
+  /**
+   * Interval in milliseconds for periodic monitoring of CPU usage and up/down scaling thread pool size accordingly.
+   * {@value }
+   */
+  public static final String FS_AZURE_READAHEAD_V2_CPU_MONITORING_INTERVAL_MILLIS = "fs.azure.readahead.v2.cpu.monitoring.interval.millis";
+
+  /**
+   * Percentage by which the thread pool size should be upscaled when CPU usage is low.
+   */
+  public static final String FS_AZURE_READAHEAD_V2_THREAD_POOL_UPSCALE_PERCENTAGE = "fs.azure.readahead.v2.thread.pool.upscale.percentage";
+
+  /**
+   * Percentage by which the thread pool size should be downscaled when CPU usage is high.
+   */
+  public static final String FS_AZURE_READAHEAD_V2_THREAD_POOL_DOWNSCALE_PERCENTAGE = "fs.azure.readahead.v2.thread.pool.downscale.percentage";
+
+  /**
+   * Interval in milliseconds for periodic monitoring of memory usage and up/down scaling buffer pool size accordingly.
+   * {@value }
+   */
+  public static final String FS_AZURE_READAHEAD_V2_MEMORY_MONITORING_INTERVAL_MILLIS = "fs.azure.readahead.v2.memory.monitoring.interval.millis";
+
   /**
    * TTL in milliseconds for the idle threads in executor service used by read ahead v2.
    */
@@ -300,6 +328,16 @@ public final class ConfigurationKeys {
    */
   public static final String FS_AZURE_READAHEAD_V2_CACHED_BUFFER_TTL_MILLIS = "fs.azure.readahead.v2.cached.buffer.ttl.millis";
 
+  /**
+   * Threshold percentage for CPU usage to scale up/down the thread pool size in read ahead v2.
+   */
+  public static final String FS_AZURE_READAHEAD_V2_CPU_USAGE_THRESHOLD_PERCENT = "fs.azure.readahead.v2.cpu.usage.threshold.percent";
+
+  /**
+   * Threshold percentage for memory usage to scale up/down the buffer pool size in read ahead v2.
+   */
+  public static final String FS_AZURE_READAHEAD_V2_MEMORY_USAGE_THRESHOLD_PERCENT = "fs.azure.readahead.v2.memory.usage.threshold.percent";
+
   /** Setting this true will make the driver use it's own RemoteIterator implementation */
   public static final String FS_AZURE_ENABLE_ABFS_LIST_ITERATOR = "fs.azure.enable.abfslistiterator";
   /** Server side encryption key encoded in Base6format {@value}.*/

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/constants/FileSystemConfigurations.java

@@ -129,12 +129,19 @@ public final class FileSystemConfigurations {
 
   public static final boolean DEFAULT_ENABLE_READAHEAD = true;
   public static final boolean DEFAULT_ENABLE_READAHEAD_V2 = false;
+  public static final boolean DEFAULT_ENABLE_READAHEAD_V2_DYNAMIC_SCALING = true;
   public static final int DEFAULT_READAHEAD_V2_MIN_THREAD_POOL_SIZE = -1;
   public static final int DEFAULT_READAHEAD_V2_MAX_THREAD_POOL_SIZE = -1;
   public static final int DEFAULT_READAHEAD_V2_MIN_BUFFER_POOL_SIZE = -1;
   public static final int DEFAULT_READAHEAD_V2_MAX_BUFFER_POOL_SIZE = -1;
-  public static final int DEFAULT_READAHEAD_V2_EXECUTOR_SERVICE_TTL_MILLIS = 3_000;
+  public static final int DEFAULT_READAHEAD_V2_CPU_MONITORING_INTERVAL_MILLIS = 6_000;
+  public static final int DEFAULT_READAHEAD_V2_THREAD_POOL_UPSCALE_PERCENTAGE = 20;
+  public static final int DEFAULT_READAHEAD_V2_THREAD_POOL_DOWNSCALE_PERCENTAGE = 30;
+  public static final int DEFAULT_READAHEAD_V2_MEMORY_MONITORING_INTERVAL_MILLIS = 6_000;
+  public static final int DEFAULT_READAHEAD_V2_EXECUTOR_SERVICE_TTL_MILLIS = 6_000;
   public static final int DEFAULT_READAHEAD_V2_CACHED_BUFFER_TTL_MILLIS = 6_000;
+  public static final int DEFAULT_READAHEAD_V2_CPU_USAGE_THRESHOLD_PERCENT = 50;
+  public static final int DEFAULT_READAHEAD_V2_MEMORY_USAGE_THRESHOLD_PERCENT = 50;
 
   public static final String DEFAULT_FS_AZURE_USER_AGENT_PREFIX = EMPTY_STRING;
   public static final String DEFAULT_VALUE_UNKNOWN = "UNKNOWN";
@@ -210,6 +217,7 @@ public final class FileSystemConfigurations {
 
   public static final int ZERO = 0;
   public static final int HUNDRED = 100;
+  public static final double ONE_HUNDRED = 100.0;
   public static final long THOUSAND = 1000L;
 
   public static final HttpOperationType DEFAULT_NETWORKING_LIBRARY

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/AbfsInputStream.java

@@ -825,6 +825,10 @@ public String getStreamID() {
     return inputStreamId;
   }
 
+  public String getETag() {
+    return eTag;
+  }
+
   /**
    * Getter for AbfsInputStreamStatistics.
    *
@@ -922,6 +926,10 @@ long getLimit() {
     return this.limit;
   }
 
+  boolean isFirstRead() {
+    return this.firstRead;
+  }
+
   @VisibleForTesting
   BackReference getFsBackRef() {
     return fsBackRef;

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/ReadBuffer.java

@@ -20,15 +20,18 @@
 
 import java.io.IOException;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import org.apache.hadoop.fs.azurebfs.contracts.services.ReadBufferStatus;
 import org.apache.hadoop.fs.azurebfs.utils.TracingContext;
 
 import static org.apache.hadoop.fs.azurebfs.contracts.services.ReadBufferStatus.READ_FAILED;
 
-class ReadBuffer {
+public class ReadBuffer {
 
   private AbfsInputStream stream;
+  private String eTag;
+  private String path;                   // path of the file this buffer is for
   private long offset;                   // offset within the file for the buffer
   private int length;                    // actual length, set after the buffer is filles
   private int requestedLength;           // requested length of the read
@@ -44,17 +47,34 @@ class ReadBuffer {
   private boolean isFirstByteConsumed = false;
   private boolean isLastByteConsumed = false;
   private boolean isAnyByteConsumed = false;
+  private AtomicInteger refCount = new AtomicInteger(0);
 
   private IOException errException = null;
 
   public AbfsInputStream getStream() {
     return stream;
   }
 
+  public String getETag() {
+    return eTag;
+  }
+
+  public String getPath() {
+    return path;
+  }
+
   public void setStream(AbfsInputStream stream) {
     this.stream = stream;
   }
 
+  public void setETag(String eTag) {
+    this.eTag = eTag;
+  }
+
+  public void setPath(String path) {
+    this.path = path;
+  }
+
   public void setTracingContext(TracingContext tracingContext) {
     this.tracingContext = tracingContext;
   }
@@ -122,6 +142,20 @@ public void setStatus(ReadBufferStatus status) {
     }
   }
 
+  public void startReading() {
+    refCount.getAndIncrement();
+  }
+
+  public void endReading() {
+    if (refCount.decrementAndGet() < 0) {
+      throw new IllegalStateException("ReadBuffer refCount cannot be negative");
+    }
+  }
+
+  public int getRefCount() {
+    return refCount.get();
+  }
+
   public CountDownLatch getLatch() {
     return latch;
   }
@@ -162,4 +196,7 @@ public void setAnyByteConsumed(boolean isAnyByteConsumed) {
     this.isAnyByteConsumed = isAnyByteConsumed;
   }
 
+  public boolean isFullyConsumed() {
+    return isFirstByteConsumed() && isLastByteConsumed();
+  }
 }

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/ReadBufferManager.java

@@ -119,7 +119,6 @@ abstract void doneReading(ReadBuffer buffer,
    */
   abstract void purgeBuffersForStream(AbfsInputStream stream);
 
-
   // Following Methods are for testing purposes only and should not be used in production code.
 
   /**
@@ -264,7 +263,7 @@ protected synchronized List<ReadBuffer> getReadAheadQueueCopy() {
    * @return a list of in-progress {@link ReadBuffer} objects
    */
   @VisibleForTesting
-  protected synchronized List<ReadBuffer> getInProgressCopiedList() {
+  protected synchronized List<ReadBuffer> getInProgressListCopy() {
     return new ArrayList<>(inProgressList);
   }
 

hadoop-tools/hadoop-azure/src/main/java/org/apache/hadoop/fs/azurebfs/services/ReadBufferManagerV1.java

@@ -33,7 +33,7 @@
  * The Read Buffer Manager for Rest AbfsClient.
  * V1 implementation of ReadBufferManager.
  */
-final class ReadBufferManagerV1 extends ReadBufferManager {
+public final class ReadBufferManagerV1 extends ReadBufferManager {
 
   private static final int NUM_BUFFERS = 16;
   private static final int NUM_THREADS = 8;
@@ -66,7 +66,7 @@ static void setReadBufferManagerConfigs(int readAheadBlockSize) {
    * Returns the singleton instance of ReadBufferManagerV1.
    * @return the singleton instance of ReadBufferManagerV1
    */
-  static ReadBufferManagerV1 getBufferManager() {
+  public static ReadBufferManagerV1 getBufferManager() {
     if (bufferManager == null) {
       LOCK.lock();
       try {
@@ -88,7 +88,7 @@ static ReadBufferManagerV1 getBufferManager() {
   void init() {
     buffers = new byte[NUM_BUFFERS][];
     for (int i = 0; i < NUM_BUFFERS; i++) {
-      buffers[i] = new byte[getReadAheadBlockSize()];  // same buffers are reused. These byte arrays are never garbage collected
+      buffers[i] = new byte[getReadAheadBlockSize()];  // same buffers are reused. The byte array never goes back to GC
       getFreeList().add(i);
     }
     for (int i = 0; i < NUM_THREADS; i++) {