HADOOP-16948. Add unit test for acquire lease failure

Author: billierinaldi

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

@@ -21,9 +21,11 @@
 import java.io.IOException;
 import java.util.concurrent.TimeUnit;
 
+import org.apache.hadoop.thirdparty.com.google.common.annotations.VisibleForTesting;
 import org.apache.hadoop.thirdparty.com.google.common.util.concurrent.FutureCallback;
 import org.apache.hadoop.thirdparty.com.google.common.util.concurrent.ListenableScheduledFuture;
 import org.apache.hadoop.thirdparty.org.checkerframework.checker.nullness.qual.Nullable;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -51,8 +53,10 @@
 public final class AbfsLease {
   private static final Logger LOG = LoggerFactory.getLogger(AbfsLease.class);
 
-  static final int LEASE_ACQUIRE_RETRY_INTERVAL = 10; // Retry interval for acquiring lease in secs
-  static final int LEASE_ACQUIRE_MAX_RETRIES = 7; // Number of retries for acquiring lease
+  // Number of retries for acquiring lease
+  static final int DEFAULT_LEASE_ACQUIRE_MAX_RETRIES = 7;
+  // Retry interval for acquiring lease in secs
+  static final int DEFAULT_LEASE_ACQUIRE_RETRY_INTERVAL = 10;
 
   private final AbfsClient client;
   private final String path;
@@ -61,6 +65,7 @@ public final class AbfsLease {
   private volatile boolean leaseFreed;
   private volatile String leaseID = null;
   private volatile Throwable exception = null;
+  private volatile int acquireRetryCount = 0;
   private volatile ListenableScheduledFuture<AbfsRestOperation> future = null;
 
   public static class LeaseException extends AzureBlobFileSystemException {
@@ -74,6 +79,12 @@ public LeaseException(String s) {
   }
 
   public AbfsLease(AbfsClient client, String path) throws AzureBlobFileSystemException {
+    this(client, path, DEFAULT_LEASE_ACQUIRE_MAX_RETRIES, DEFAULT_LEASE_ACQUIRE_RETRY_INTERVAL);
+  }
+
+  @VisibleForTesting
+  public AbfsLease(AbfsClient client, String path, int acquireMaxRetries,
+      int acquireRetryInterval) throws AzureBlobFileSystemException {
     this.leaseFreed = false;
     this.client = client;
     this.path = path;
@@ -84,8 +95,8 @@ public AbfsLease(AbfsClient client, String path) throws AzureBlobFileSystemExcep
 
     // Try to get the lease a specified number of times, else throw an error
     RetryPolicy retryPolicy = RetryPolicies.retryUpToMaximumCountWithFixedSleep(
-        LEASE_ACQUIRE_MAX_RETRIES, LEASE_ACQUIRE_RETRY_INTERVAL, TimeUnit.SECONDS);
-    acquireLease(retryPolicy, 0, 0);
+        acquireMaxRetries, acquireRetryInterval, TimeUnit.SECONDS);
+    acquireLease(retryPolicy, 0, acquireRetryInterval, 0);
 
     while (leaseID == null && exception == null) {
       try {
@@ -103,7 +114,7 @@ public AbfsLease(AbfsClient client, String path) throws AzureBlobFileSystemExcep
     LOG.debug("Acquired lease {} on {}", leaseID, path);
   }
 
-  private void acquireLease(RetryPolicy retryPolicy, int numRetries, long delay)
+  private void acquireLease(RetryPolicy retryPolicy, int numRetries, int retryInterval, long delay)
       throws LeaseException {
     LOG.debug("Attempting to acquire lease on {}, retry {}", path, numRetries);
     if (future != null && !future.isDone()) {
@@ -124,7 +135,8 @@ public void onFailure(Throwable throwable) {
           if (RetryPolicy.RetryAction.RetryDecision.RETRY
               == retryPolicy.shouldRetry(null, numRetries, 0, true).action) {
             LOG.debug("Failed to acquire lease on {}, retrying: {}", path, throwable);
-            acquireLease(retryPolicy, numRetries + 1, LEASE_ACQUIRE_RETRY_INTERVAL);
+            acquireRetryCount++;
+            acquireLease(retryPolicy, numRetries + 1, retryInterval, retryInterval);
           } else {
             exception = throwable;
           }
@@ -168,4 +180,9 @@ public boolean isFreed() {
   public String getLeaseID() {
     return leaseID;
   }
+
+  @VisibleForTesting
+  public int getAcquireRetryCount() {
+    return acquireRetryCount;
+  }
 }

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

@@ -26,10 +26,18 @@
 import org.apache.hadoop.conf.Configuration;
 import org.apache.hadoop.fs.FSDataOutputStream;
 import org.apache.hadoop.fs.Path;
+import org.apache.hadoop.fs.azurebfs.contracts.exceptions.AzureBlobFileSystemException;
+import org.apache.hadoop.fs.azurebfs.services.AbfsClient;
+import org.apache.hadoop.fs.azurebfs.services.AbfsLease;
 import org.apache.hadoop.fs.azurebfs.services.AbfsOutputStream;
 import org.apache.hadoop.test.GenericTestUtils;
 import org.apache.hadoop.test.LambdaTestUtils;
 
+import static org.mockito.ArgumentMatchers.anyInt;
+import static org.mockito.ArgumentMatchers.anyString;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.spy;
+
 import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_INFINITE_LEASE_KEY;
 import static org.apache.hadoop.fs.azurebfs.constants.ConfigurationKeys.FS_AZURE_LEASE_THREADS;
 import static org.apache.hadoop.fs.azurebfs.constants.TestConfigurationKeys.FS_AZURE_TEST_NAMESPACE_ENABLED_ACCOUNT;
@@ -293,4 +301,36 @@ public void testFileSystemClose() throws Exception {
       return "Expected exception on new append after closed FS";
     });
   }
+
+  @Test(timeout = TEST_EXECUTION_TIMEOUT)
+  public void testAcquireRetry() throws Exception {
+    final Path testFilePath = new Path(path(methodName.getMethodName()), TEST_FILE);
+    final AzureBlobFileSystem fs = getCustomFileSystem(testFilePath.getParent(), 1);
+    fs.mkdirs(testFilePath.getParent());
+    fs.createNewFile(testFilePath);
+
+    AbfsLease lease = new AbfsLease(fs.getAbfsClient(), testFilePath.toUri().getPath());
+    Assert.assertNotNull("Did not successfully lease file", lease.getLeaseID());
+    lease.free();
+    Assert.assertEquals("Unexpected acquire retry count", 0, lease.getAcquireRetryCount());
+
+    AbfsClient mockClient = spy(fs.getAbfsClient());
+
+    doThrow(new AbfsLease.LeaseException("failed to acquire 1"))
+        .doThrow(new AbfsLease.LeaseException("failed to acquire 2"))
+        .doCallRealMethod()
+        .when(mockClient).acquireLease(anyString(), anyInt());
+
+    lease = new AbfsLease(mockClient, testFilePath.toUri().getPath(), 5, 1);
+    Assert.assertNotNull("Acquire lease should have retried", lease.getLeaseID());
+    lease.free();
+    Assert.assertEquals("Unexpected acquire retry count", 2, lease.getAcquireRetryCount());
+
+    doThrow(new AbfsLease.LeaseException("failed to acquire"))
+        .when(mockClient).acquireLease(anyString(), anyInt());
+
+    LambdaTestUtils.intercept(AzureBlobFileSystemException.class, () -> {
+      new AbfsLease(mockClient, testFilePath.toUri().getPath(), 5, 1);
+    });
+  }
 }