HADOOP-17377: ABFS: MsiTokenProvider doesn't retry HTTP 429/410 from the Instance Metadata Service

hadoop-tools/hadoop-azure/pom.xml

@@ -321,8 +321,23 @@
     <dependency>
       <groupId>org.mockito</groupId>
       <artifactId>mockito-core</artifactId>
+      <version>4.11.0</version>
       <scope>test</scope>
     </dependency>
+
+    <dependency>
+      <groupId>org.mockito</groupId>
+      <artifactId>mockito-inline</artifactId>
+      <version>4.11.0</version>
+      <scope>test</scope>
+      <exclusions>
+        <exclusion>
+          <groupId>org.mockito</groupId>
+          <artifactId>mockito-core</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
     <dependency>
       <groupId>org.apache.hadoop</groupId>
       <artifactId>hadoop-minikdc</artifactId>

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

@@ -213,7 +213,7 @@ public String getRequestId() {
       return this.requestId;
     }
 
-    protected HttpException(
+    public HttpException(
         final int httpErrorCode,
         final String requestId,
         final String message,
@@ -341,7 +341,7 @@ private static boolean isRecoverableFailure(IOException e) {
         || e instanceof FileNotFoundException);
   }
 
-  private static AzureADToken getTokenSingleCall(String authEndpoint,
+  public static AzureADToken getTokenSingleCall(String authEndpoint,
       String payload, Hashtable<String, String> headers, String httpMethod,
       boolean isMsi)
           throws IOException {

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

@@ -58,6 +58,13 @@ public class ExponentialRetryPolicy {
    */
   private static final double MAX_RANDOM_RATIO = 1.2;
 
+  /**
+   * Qualifies for retry based on
+   * https://learn.microsoft.com/en-us/azure/active-directory/
+   * managed-identities-azure-resources/how-to-use-vm-token#error-handling
+   */
+  private static final int HTTP_TOO_MANY_REQUESTS = 429;
+
   /**
    *  Holds the random number generator used to calculate randomized backoff intervals
    */
@@ -123,6 +130,9 @@ public ExponentialRetryPolicy(final int retryCount, final int minBackoff, final
    * and the current strategy. The valid http status code lies in the range of 1xx-5xx.
    * But an invalid status code might be set due to network or timeout kind of issues.
    * Such invalid status code also qualify for retry.
+   * HTTP status code 410 qualifies for retry based on
+   * https://docs.microsoft.com/en-in/azure/virtual-machines/linux/
+   * instance-metadata-service?tabs=windows#errors-and-debugging
    *
    * @param retryCount The current retry attempt count.
    * @param statusCode The status code of the response, or -1 for socket error.
@@ -132,6 +142,8 @@ public boolean shouldRetry(final int retryCount, final int statusCode) {
     return retryCount < this.retryCount
         && (statusCode < HTTP_CONTINUE
         || statusCode == HttpURLConnection.HTTP_CLIENT_TIMEOUT
+        || statusCode == HttpURLConnection.HTTP_GONE
+        || statusCode == HTTP_TOO_MANY_REQUESTS
         || (statusCode >= HttpURLConnection.HTTP_INTERNAL_ERROR
             && statusCode != HttpURLConnection.HTTP_NOT_IMPLEMENTED
             && statusCode != HttpURLConnection.HTTP_VERSION));

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

@@ -19,15 +19,22 @@
 package org.apache.hadoop.fs.azurebfs;
 
 import java.io.IOException;
+import java.lang.reflect.Field;
+import java.net.HttpURLConnection;
 import java.util.Date;
-
 import org.junit.Test;
+import org.mockito.MockedStatic;
+import org.mockito.Mockito;
 
 import org.apache.commons.lang3.StringUtils;
 import org.apache.hadoop.fs.azurebfs.oauth2.AccessTokenProvider;
+import org.apache.hadoop.fs.azurebfs.oauth2.AzureADAuthenticator;
 import org.apache.hadoop.fs.azurebfs.oauth2.AzureADToken;
 import org.apache.hadoop.fs.azurebfs.oauth2.MsiTokenProvider;
+import org.apache.hadoop.fs.azurebfs.services.ExponentialRetryPolicy;
 
+import static org.apache.hadoop.fs.azurebfs.constants.FileSystemConfigurations.DEFAULT_AZURE_OAUTH_TOKEN_FETCH_RETRY_MAX_ATTEMPTS;
+import static org.apache.hadoop.test.LambdaTestUtils.intercept;
 import static org.junit.Assume.assumeThat;
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.not;
@@ -40,13 +47,16 @@
 import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY;
 import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT;
 import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_ACCOUNT_OAUTH_MSI_TENANT;
+import static org.mockito.Mockito.times;
 
 /**
  * Test MsiTokenProvider.
  */
 public final class ITestAbfsMsiTokenProvider
     extends AbstractAbfsIntegrationTest {
 
+  private static final int HTTP_TOO_MANY_REQUESTS = 429;
+
   public ITestAbfsMsiTokenProvider() throws Exception {
     super();
   }
@@ -90,4 +100,109 @@ private String getTrimmedPasswordString(AbfsConfiguration conf, String key,
     return value.trim();
   }
 
+  /**
+   * Test to verify that token fetch is retried for throttling errors (too many requests 429).
+   * @throws Exception
+   */
+  @Test
+  public void testRetryForThrottling() throws Exception {
+    AbfsConfiguration conf = getConfiguration();
+
+    // Exception to be thrown with throttling error code 429.
+    AzureADAuthenticator.HttpException httpException
+        = new AzureADAuthenticator.HttpException(HTTP_TOO_MANY_REQUESTS,
+        "abc", "abc", "abc", "abc", "abc");
+
+    String tenantGuid = "abcd";
+    String clientId = "abcd";
+    String authEndpoint = getTrimmedPasswordString(conf,
+        FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT,
+        DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT);
+    String authority = getTrimmedPasswordString(conf,
+        FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY,
+        DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY);
+
+    // Mock the getTokenSingleCall to throw exception so the retry logic comes into place.
+    try (MockedStatic<AzureADAuthenticator> adAuthenticator = Mockito.mockStatic(
+        AzureADAuthenticator.class, Mockito.CALLS_REAL_METHODS)) {
+      adAuthenticator.when(
+          () -> AzureADAuthenticator.getTokenSingleCall(Mockito.anyString(),
+              Mockito.anyString(), Mockito.any(), Mockito.anyString(),
+              Mockito.anyBoolean())).thenThrow(httpException);
+
+      // Mock the tokenFetchRetryPolicy to verify retries.
+      ExponentialRetryPolicy exponentialRetryPolicy = Mockito.spy(
+          conf.getOauthTokenFetchRetryPolicy());
+      Field tokenFetchRetryPolicy = AzureADAuthenticator.class.getDeclaredField(
+          "tokenFetchRetryPolicy");
+      tokenFetchRetryPolicy.setAccessible(true);
+      tokenFetchRetryPolicy.set(ExponentialRetryPolicy.class,
+          exponentialRetryPolicy);
+
+      AccessTokenProvider tokenProvider = new MsiTokenProvider(authEndpoint,
+          tenantGuid, clientId, authority);
+      AzureADToken token = null;
+      intercept(AzureADAuthenticator.HttpException.class,
+          tokenProvider::getToken);
+
+      // If the status code doesn't qualify for retry shouldRetry returns false and the loop ends.
+      // It being called multiple times verifies that the retry was done for the throttling status code 429.
+      Mockito.verify(exponentialRetryPolicy,
+              times(DEFAULT_AZURE_OAUTH_TOKEN_FETCH_RETRY_MAX_ATTEMPTS + 1))
+          .shouldRetry(Mockito.anyInt(), Mockito.anyInt());
+    }
+  }
+
+  /**
+   * Test to verify that token fetch is not retried for resource not found errors.
+   * @throws Exception
+   */
+  @Test
+  public void testNoRetryForResourceNotFound() throws Exception {
+    AbfsConfiguration conf = getConfiguration();
+
+    // Exception to be thrown 404 error code.
+    AzureADAuthenticator.HttpException httpException
+        = new AzureADAuthenticator.HttpException(HttpURLConnection.HTTP_NOT_FOUND,
+        "abc", "abc", "abc", "abc", "abc");
+
+    String tenantGuid = "abcd";
+    String clientId = "abcd";
+    String authEndpoint = getTrimmedPasswordString(conf,
+        FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT,
+        DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_ENDPOINT);
+    String authority = getTrimmedPasswordString(conf,
+        FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY,
+        DEFAULT_FS_AZURE_ACCOUNT_OAUTH_MSI_AUTHORITY);
+
+    // Mock the getTokenSingleCall to throw exception.
+    try (MockedStatic<AzureADAuthenticator> adAuthenticator = Mockito.mockStatic(
+        AzureADAuthenticator.class, Mockito.CALLS_REAL_METHODS)) {
+      adAuthenticator.when(
+          () -> AzureADAuthenticator.getTokenSingleCall(Mockito.anyString(),
+              Mockito.anyString(), Mockito.any(), Mockito.anyString(),
+              Mockito.anyBoolean())).thenThrow(httpException);
+
+      // Mock the tokenFetchRetryPolicy to verify no retries.
+      ExponentialRetryPolicy exponentialRetryPolicy = Mockito.spy(
+          conf.getOauthTokenFetchRetryPolicy());
+      Field tokenFetchRetryPolicy = AzureADAuthenticator.class.getDeclaredField(
+          "tokenFetchRetryPolicy");
+      tokenFetchRetryPolicy.setAccessible(true);
+      tokenFetchRetryPolicy.set(ExponentialRetryPolicy.class,
+          exponentialRetryPolicy);
+
+      AccessTokenProvider tokenProvider = new MsiTokenProvider(authEndpoint,
+          tenantGuid, clientId, authority);
+      AzureADToken token = null;
+      intercept(AzureADAuthenticator.HttpException.class,
+          tokenProvider::getToken);
+
+      // If the status code doesn't qualify for retry shouldRetry returns false and the loop ends.
+      // It being called only once verifies that retry doesn't come into place..
+      Mockito.verify(exponentialRetryPolicy,
+              times(1))
+          .shouldRetry(Mockito.anyInt(), Mockito.anyInt());
+    }
+  }
 }