[GOBBLIN-2096] Retry Transient SQL Exception for MALA (#3982)

* Add retries to RetryableSQLException for MALA * Throw exception if exceed max retry
apache · Jun 20, 2024 · d9f83e6 · d9f83e6
1 parent 7503385
commit d9f83e6
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 5 deletions.
diff --git a/...n/java/org/apache/gobblin/service/modules/orchestration/MysqlMultiActiveLeaseArbiter.java b/...n/java/org/apache/gobblin/service/modules/orchestration/MysqlMultiActiveLeaseArbiter.java
@@ -23,6 +23,7 @@
 import java.sql.ResultSet;
 import java.sql.SQLException;
 import java.sql.SQLIntegrityConstraintViolationException;
+import java.sql.SQLTransientException;
 import java.sql.Timestamp;
 import java.util.Calendar;
 import java.util.Optional;
@@ -41,6 +42,7 @@
 import org.apache.gobblin.service.ServiceConfigKeys;
 import org.apache.gobblin.util.ConfigUtils;
 import org.apache.gobblin.util.DBStatementExecutor;
+import org.apache.gobblin.util.ExponentialBackoff;
 
 
 /**
@@ -171,6 +173,9 @@ public class MysqlMultiActiveLeaseArbiter implements MultiActiveLeaseArbiter {
   // Complete lease acquisition if values have not changed since lease was acquired
   protected static final String CONDITIONALLY_COMPLETE_LEASE_STATEMENT = "UPDATE %s SET "
       + "event_timestamp=event_timestamp, lease_acquisition_timestamp = NULL " + WHERE_CLAUSE_TO_MATCH_ROW;
+  protected static final int MAX_RETRIES = 3;
+  protected static final long MIN_INITIAL_DELAY_MILLIS = 20L;
+  protected static final long DELAY_FOR_RETRY_RANGE_MILLIS = 200L;
   private static final ThreadLocal<Calendar> UTC_CAL =
       ThreadLocal.withInitial(() -> Calendar.getInstance(TimeZone.getTimeZone("UTC")));
 
@@ -254,7 +259,10 @@ public LeaseAttemptStatus tryAcquireLease(DagActionStore.DagActionLeaseObject da
         log.debug("tryAcquireLease for [{}, is; {}, eventTimestamp: {}] - CASE 1: no existing row for this dag action,"
             + " then go ahead and insert", dagActionLeaseObject.getDagAction(),
             dagActionLeaseObject.isReminder() ? "reminder" : "original", dagActionLeaseObject.getEventTimeMillis());
-        int numRowsUpdated = attemptLeaseIfNewRow(dagActionLeaseObject.getDagAction());
+        int numRowsUpdated = attemptLeaseIfNewRow(dagActionLeaseObject.getDagAction(),
+            ExponentialBackoff.builder().maxRetries(MAX_RETRIES)
+                .initialDelay(MIN_INITIAL_DELAY_MILLIS + (long) Math.random() * DELAY_FOR_RETRY_RANGE_MILLIS)
+                .build());
        return evaluateStatusAfterLeaseAttempt(numRowsUpdated, dagActionLeaseObject.getDagAction(),
            Optional.empty(), dagActionLeaseObject.isReminder(), adoptConsensusFlowExecutionId);
       }
@@ -420,15 +428,25 @@ protected GetEventInfoResult createGetInfoResult(ResultSet resultSet) throws IOE
    * near past for it.
    * @return int corresponding to number of rows updated by INSERT statement to acquire lease
    */
-  protected int attemptLeaseIfNewRow(DagActionStore.DagAction dagAction) throws IOException {
+  protected int attemptLeaseIfNewRow(DagActionStore.DagAction dagAction, ExponentialBackoff exponentialBackoff) throws IOException {
     String formattedAcquireLeaseNewRowStatement =
         String.format(ACQUIRE_LEASE_IF_NEW_ROW_STATEMENT, this.leaseArbiterTableName);
     return dbStatementExecutor.withPreparedStatement(formattedAcquireLeaseNewRowStatement,
         insertStatement -> {
           completeInsertPreparedStatement(insertStatement, dagAction);
           try {
             return insertStatement.executeUpdate();
-          } catch (SQLIntegrityConstraintViolationException e) {
+          } catch (SQLTransientException e) {
+            try {
+              if (exponentialBackoff.awaitNextRetryIfAvailable()) {
+                return attemptLeaseIfNewRow(dagAction, exponentialBackoff);
+              }
+            } catch (InterruptedException e2) {
+              throw new IOException(e2);
+            }
+            throw e; 
+          }
+          catch (SQLIntegrityConstraintViolationException e) {
             if (!e.getMessage().contains("Duplicate entry")) {
               throw e;
             }

diff --git a/...va/org/apache/gobblin/service/modules/orchestration/MysqlMultiActiveLeaseArbiterTest.java b/...va/org/apache/gobblin/service/modules/orchestration/MysqlMultiActiveLeaseArbiterTest.java
@@ -23,6 +23,7 @@
 import java.util.Optional;
 import java.util.UUID;
 
+import org.apache.gobblin.util.ExponentialBackoff;
 import org.testng.Assert;
 import org.testng.annotations.AfterClass;
 import org.testng.annotations.BeforeClass;
@@ -201,9 +202,13 @@ Tests attemptLeaseIfNewRow() method to ensure a new row is inserted if no row ma
   @Test
   public void testAcquireLeaseIfNewRow() throws IOException {
     // Inserting the first time should update 1 row
-    Assert.assertEquals(mysqlMultiActiveLeaseArbiter.attemptLeaseIfNewRow(resumeDagAction), 1);
+    Assert.assertEquals(mysqlMultiActiveLeaseArbiter.attemptLeaseIfNewRow(resumeDagAction,
+        ExponentialBackoff.builder().maxRetries(MysqlMultiActiveLeaseArbiter.MAX_RETRIES)
+            .initialDelay(MysqlMultiActiveLeaseArbiter.MIN_INITIAL_DELAY_MILLIS).build()), 1);
     // Inserting the second time should not update any rows
-    Assert.assertEquals(mysqlMultiActiveLeaseArbiter.attemptLeaseIfNewRow(resumeDagAction), 0);
+    Assert.assertEquals(mysqlMultiActiveLeaseArbiter.attemptLeaseIfNewRow(resumeDagAction,
+        ExponentialBackoff.builder().maxRetries(MysqlMultiActiveLeaseArbiter.MAX_RETRIES)
+            .initialDelay(MysqlMultiActiveLeaseArbiter.MIN_INITIAL_DELAY_MILLIS).build()), 0);
   }
 
     /*