Associate transaction during synchronizations

wildfly · Mar 14, 2017 · 4cfb50b · 4cfb50b
1 parent d053c43
commit 4cfb50b
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 4 deletions.
diff --git a/src/main/java/org/wildfly/transaction/client/AbstractTransaction.java b/src/main/java/org/wildfly/transaction/client/AbstractTransaction.java
@@ -239,4 +239,28 @@ public <E extends Exception> void performAction(ExceptionRunnable<E> action) thr
     Object getOutflowLock() {
         return outflowLock;
     }
+
+    final class AssociatingSynchronization implements Synchronization {
+        private final Synchronization sync;
+
+        AssociatingSynchronization(final Synchronization sync) {
+            this.sync = sync;
+        }
+
+        public void beforeCompletion() {
+            try {
+                performConsumer(Synchronization::beforeCompletion, sync);
+            } catch (SystemException e) {
+                throw new SynchronizationException(e);
+            }
+        }
+
+        public void afterCompletion(final int status) {
+            try {
+                performConsumer(Synchronization::afterCompletion, sync, status);
+            } catch (SystemException e) {
+                throw new SynchronizationException(e);
+            }
+        }
+    }
 }
diff --git a/src/main/java/org/wildfly/transaction/client/LocalTransaction.java b/src/main/java/org/wildfly/transaction/client/LocalTransaction.java
@@ -139,7 +139,7 @@ public boolean delistResource(final XAResource xaRes, final int flag) throws Ill
 
     public void registerSynchronization(final Synchronization sync) throws RollbackException, IllegalStateException, SystemException {
         Assert.checkNotNullParam("sync", sync);
-        transaction.registerSynchronization(sync);
+        transaction.registerSynchronization(new AssociatingSynchronization(sync));
     }
 
     /**
@@ -162,7 +162,7 @@ public Xid getXid() {
 
     void registerInterposedSynchronization(final Synchronization sync) throws IllegalStateException {
         Assert.checkNotNullParam("sync", sync);
-        owner.getProvider().registerInterposedSynchronization(transaction, sync);
+        owner.getProvider().registerInterposedSynchronization(transaction, new AssociatingSynchronization(sync));
     }
 
     Object getResource(final Object key) throws NullPointerException {

diff --git a/src/main/java/org/wildfly/transaction/client/RemoteTransaction.java b/src/main/java/org/wildfly/transaction/client/RemoteTransaction.java
@@ -139,12 +139,12 @@ public boolean delistResource(final XAResource xaRes, final int flag) {
 
     public void registerSynchronization(final Synchronization sync) throws RollbackException, IllegalStateException, SystemException {
         Assert.checkNotNullParam("sync", sync);
-        stateRef.get().registerSynchronization(sync);
+        stateRef.get().registerSynchronization(new AssociatingSynchronization(sync));
     }
 
     void registerInterposedSynchronization(final Synchronization sync) throws IllegalStateException {
         Assert.checkNotNullParam("sync", sync);
-        stateRef.get().registerInterposedSynchronization(sync);
+        stateRef.get().registerInterposedSynchronization(new AssociatingSynchronization(sync));
     }
 
     public int hashCode() {