HADOOP-18501: ABFS: Partial read should add to throttling data: DRAFT

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

@@ -1106,4 +1106,8 @@ public void setEnableAbfsListIterator(boolean enableAbfsListIterator) {
     this.enableAbfsListIterator = enableAbfsListIterator;
   }
 
+  @VisibleForTesting
+  void setReadAheadQueueDepth(int depth) {
+    readAheadQueueDepth = depth;
+  }
 }

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

@@ -112,5 +112,7 @@ public final class AbfsHttpConstants {
   public static final char CHAR_STAR = '*';
   public static final char CHAR_PLUS = '+';
 
+  public static final String CONNECTION_RESET = "Connection reset";
+
   private AbfsHttpConstants() {}
 }

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

@@ -23,6 +23,7 @@
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.fs.azurebfs.AbfsStatistic;
 import org.apache.hadoop.fs.azurebfs.constants.HttpHeaderConfigurations;
 
@@ -70,7 +71,7 @@ static void updateMetrics(AbfsRestOperationType operationType,
       return;
     }
 
-    int status = abfsHttpOperation.getStatusCode();
+    final int status = abfsHttpOperation.getStatusCode();
     long contentLength = 0;
     // If the socket is terminated prior to receiving a response, the HTTP
     // status may be 0 or -1.  A status less than 200 or greater than or equal
@@ -88,9 +89,24 @@ static void updateMetrics(AbfsRestOperationType operationType,
         break;
       case ReadFile:
         String range = abfsHttpOperation.getConnection().getRequestProperty(HttpHeaderConfigurations.RANGE);
-        contentLength = getContentLengthIfKnown(range);
-        if (contentLength > 0) {
-          singleton.readThrottler.addBytesTransferred(contentLength,
+        long contentLengthRequested = getContentLengthIfKnown(range);
+        long bytesToBeAddedInMetric = contentLengthRequested;
+
+        long contentLengthReceived = abfsHttpOperation.getBytesReceived();
+        if (status == HttpURLConnection.HTTP_PARTIAL
+            && contentLengthRequested > contentLengthReceived) {
+          /*
+           * In case of server response status == 206 (partial content) and the
+           * contentLength received is lesser than the requested length, we have
+           * to take it as a throttling case and hence we need to add the remaining
+           * bytes (contentLengthReceived - contentLength) in failure data-point.
+           * */
+          bytesToBeAddedInMetric = contentLengthRequested - contentLengthReceived;
+          isFailedOperation = true;
+        }
+        if (bytesToBeAddedInMetric > 0) {
+          singleton.readThrottler.addBytesTransferred(
+              bytesToBeAddedInMetric,
               isFailedOperation);
         }
         break;
@@ -140,4 +156,19 @@ private static long getContentLengthIfKnown(String range) {
     }
     return contentLength;
   }
+
+  @VisibleForTesting
+  void setReadThrottler(final AbfsClientThrottlingAnalyzer readThrottler) {
+    this.readThrottler = readThrottler;
+  }
+
+  @VisibleForTesting
+  AbfsClientThrottlingAnalyzer getReadThrottler() {
+    return readThrottler;
+  }
+
+  @VisibleForTesting
+  static AbfsClientThrottlingIntercept getSingleton() {
+    return singleton;
+  }
 }

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

@@ -28,6 +28,7 @@
 import javax.net.ssl.HttpsURLConnection;
 import javax.net.ssl.SSLSocketFactory;
 
+import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.fs.azurebfs.utils.UriUtils;
 import org.apache.hadoop.security.ssl.DelegatingSSLSocketFactory;
 
@@ -563,4 +564,14 @@ public String getResponseHeader(final String httpHeader) {
       return "";
     }
   }
+
+  @VisibleForTesting
+  protected void setStatusCode(int statusCode) {
+    this.statusCode = statusCode;
+  }
+
+  @VisibleForTesting
+  protected void setBytesReceived(long bytesReceived) {
+    this.bytesReceived = bytesReceived;
+  }
 }

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

@@ -270,7 +270,7 @@ public synchronized int read(final byte[] b, final int off, final int len) throw
       if (currentLen <= 0 || currentLen > b.length - currentOff) {
         break;
       }
-    } while (lastReadBytes > 0);
+    } while (lastReadBytes >= 0); //backend server can return 0B with 206 httpStatus.
     return totalReadBytes > 0 ? totalReadBytes : lastReadBytes;
   }
 
@@ -505,14 +505,14 @@ private int readInternal(final long position, final byte[] b, final int offset,
         }
         return receivedBytes;
       }
-
-      // got nothing from read-ahead, do our own read now
-      receivedBytes = readRemote(position, b, offset, length, new TracingContext(tracingContext));
-      return receivedBytes;
     } else {
-      LOG.debug("read ahead disabled, reading remote");
-      return readRemote(position, b, offset, length, new TracingContext(tracingContext));
+      LOG.debug("read ahead disabled");
     }
+    LOG.debug("reading remote");
+    final long requiredLen = Math.min(length, contentLength - position);
+    return readRemote(position, b, offset, (int) requiredLen,
+        new TracingContext(tracingContext));
+
   }
 
   int readRemote(long position, byte[] b, int offset, int length, TracingContext tracingContext) throws IOException {

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

@@ -21,13 +21,15 @@
 import java.io.IOException;
 import java.io.UncheckedIOException;
 import java.net.HttpURLConnection;
+import java.net.SocketException;
 import java.net.URL;
 import java.net.UnknownHostException;
 import java.util.List;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import org.apache.hadoop.classification.VisibleForTesting;
 import org.apache.hadoop.fs.azurebfs.AbfsStatistic;
 import org.apache.hadoop.fs.azurebfs.constants.AbfsHttpConstants;
 import org.apache.hadoop.fs.azurebfs.contracts.exceptions.AbfsRestOperationException;
@@ -37,6 +39,8 @@
 import org.apache.hadoop.fs.azurebfs.utils.TracingContext;
 import org.apache.hadoop.fs.statistics.impl.IOStatisticsBinding;
 
+import static org.apache.hadoop.fs.azurebfs.constants.AbfsHttpConstants.CONNECTION_RESET;
+
 /**
  * The AbfsRestOperation for Rest AbfsClient.
  */
@@ -92,6 +96,10 @@ public URL getUrl() {
     return url;
   }
 
+  public String getMethod() {
+    return method;
+  }
+
   public List<AbfsHttpHeader> getRequestHeaders() {
     return requestHeaders;
   }
@@ -161,6 +169,7 @@ String getSasToken() {
    * @param bufferLength The length of the data in the buffer.
    * @param sasToken A sasToken for optional re-use by AbfsInputStream/AbfsOutputStream.
    */
+  @VisibleForTesting
   AbfsRestOperation(AbfsRestOperationType operationType,
                     AbfsClient client,
                     String method,
@@ -244,7 +253,7 @@ private boolean executeHttpOperation(final int retryCount,
     AbfsHttpOperation httpOperation = null;
     try {
       // initialize the HTTP request and open the connection
-      httpOperation = new AbfsHttpOperation(url, method, requestHeaders);
+      httpOperation = getHttpOperation();
       incrementCounter(AbfsStatistic.CONNECTIONS_MADE, 1);
       tracingContext.constructHeader(httpOperation);
 
@@ -311,6 +320,19 @@ private boolean executeHttpOperation(final int retryCount,
         LOG.debug("HttpRequestFailure: {}, {}", httpOperation, ex);
       }
 
+      /*
+       * In case of Connection_reset with status == 206 (partial-content)from
+       * server for read operation, the partial-result has to be sent back to
+       * AbfsInputStream which would retry for the remaining bytes.
+       * */
+      if (ex instanceof SocketException && CONNECTION_RESET.equals(
+          ex.getMessage())
+          && httpOperation.getStatusCode() == HttpURLConnection.HTTP_PARTIAL
+          && operationType == AbfsRestOperationType.ReadFile) {
+        result = httpOperation;
+        return true;
+      }
+
       if (!client.getRetryPolicy().shouldRetry(retryCount, -1)) {
         throw new InvalidAbfsRestOperationException(ex);
       }
@@ -331,6 +353,11 @@ private boolean executeHttpOperation(final int retryCount,
     return true;
   }
 
+  @VisibleForTesting
+  AbfsHttpOperation getHttpOperation() throws IOException {
+    return new AbfsHttpOperation(url, method, requestHeaders);
+  }
+
   /**
    * Incrementing Abfs counters with a long value.
    *

hadoop-tools/hadoop-azure/src/test/java/org/apache/hadoop/fs/azurebfs/AbstractAbfsIntegrationTest.java

@@ -298,6 +298,10 @@ public AzureBlobFileSystem getFileSystem() throws IOException {
     return abfs;
   }
 
+  public void setClient(final AbfsClient abfsClient) throws IOException {
+    abfs.getAbfsStore().setClient(abfsClient);
+  }
+
   public AzureBlobFileSystem getFileSystem(Configuration configuration) throws Exception{
     final AzureBlobFileSystem fs = (AzureBlobFileSystem) FileSystem.get(configuration);
     return fs;