Merge branch 'master' into dependabot/gradle/awssdkVersion-2.22.3

scalar-labs · Dec 25, 2023 · 0b8b9f2 · 0b8b9f2
2 parents d0ea7bd + 6f431ce
commit 0b8b9f2
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 10 deletions.
diff --git a/...main/java/com/scalar/db/common/ActiveTransactionManagedDistributedTransactionManager.java b/...main/java/com/scalar/db/common/ActiveTransactionManagedDistributedTransactionManager.java
@@ -19,6 +19,8 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -32,6 +34,17 @@ public abstract class ActiveTransactionManagedDistributedTransactionManager
 
   private final ActiveExpiringMap<String, ActiveTransaction> activeTransactions;
 
+  private final AtomicReference<BiConsumer<String, DistributedTransaction>>
+      transactionExpirationHandler =
+          new AtomicReference<>(
+              (id, t) -> {
+                try {
+                  t.rollback();
+                } catch (Exception e) {
+                  logger.warn("Rollback failed. transaction ID: {}", id, e);
+                }
+              });
+
   public ActiveTransactionManagedDistributedTransactionManager(DatabaseConfig config) {
     super(config);
     activeTransactions =
@@ -40,14 +53,14 @@ public ActiveTransactionManagedDistributedTransactionManager(DatabaseConfig conf
             TRANSACTION_EXPIRATION_INTERVAL_MILLIS,
             (id, t) -> {
               logger.warn("The transaction is expired. transaction ID: {}", id);
-              try {
-                t.rollback();
-              } catch (Exception e) {
-                logger.warn("Rollback failed. transaction ID: {}", id, e);
-              }
+              transactionExpirationHandler.get().accept(id, t);
             });
   }
 
+  public void setTransactionExpirationHandler(BiConsumer<String, DistributedTransaction> handler) {
+    transactionExpirationHandler.set(handler);
+  }
+
   private void add(ActiveTransaction transaction) throws TransactionException {
     if (activeTransactions.putIfAbsent(transaction.getId(), transaction).isPresent()) {
       transaction.rollback();

diff --git a/...n/java/com/scalar/db/common/ActiveTransactionManagedTwoPhaseCommitTransactionManager.java b/...n/java/com/scalar/db/common/ActiveTransactionManagedTwoPhaseCommitTransactionManager.java
@@ -21,6 +21,8 @@
 import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.BiConsumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -34,6 +36,17 @@ public abstract class ActiveTransactionManagedTwoPhaseCommitTransactionManager
 
   private final ActiveExpiringMap<String, ActiveTransaction> activeTransactions;
 
+  private final AtomicReference<BiConsumer<String, TwoPhaseCommitTransaction>>
+      transactionExpirationHandler =
+          new AtomicReference<>(
+              (id, t) -> {
+                try {
+                  t.rollback();
+                } catch (Exception e) {
+                  logger.warn("Rollback failed. transaction ID: {}", id, e);
+                }
+              });
+
   public ActiveTransactionManagedTwoPhaseCommitTransactionManager(DatabaseConfig config) {
     super(config);
     activeTransactions =
@@ -42,14 +55,15 @@ public ActiveTransactionManagedTwoPhaseCommitTransactionManager(DatabaseConfig c
             TRANSACTION_EXPIRATION_INTERVAL_MILLIS,
             (id, t) -> {
               logger.warn("The transaction is expired. transaction ID: {}", id);
-              try {
-                t.rollback();
-              } catch (Exception e) {
-                logger.warn("Rollback failed. transaction ID: {}", id, e);
-              }
+              transactionExpirationHandler.get().accept(id, t);
             });
   }
 
+  public void setTransactionExpirationHandler(
+      BiConsumer<String, TwoPhaseCommitTransaction> handler) {
+    transactionExpirationHandler.set(handler);
+  }
+
   private void add(ActiveTransaction transaction) throws TransactionException {
     if (activeTransactions.putIfAbsent(transaction.getId(), transaction).isPresent()) {
       transaction.rollback();