killbill-aplcache

Rework Payment exception handling for state machine operation

6/20/2014 10:46:17 PM

Changes

payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentPluginOperation.java 4(+0 -4)

payment/src/main/java/org/killbill/billing/payment/core/sm/retryPluginOperation.java 4(+0 -4)

Details

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/PaymentGatewayProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/PaymentGatewayProcessor.java
index 9529b94..250ec49 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/PaymentGatewayProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/PaymentGatewayProcessor.java
@@ -17,14 +17,16 @@
 
 package org.killbill.billing.payment.core;
 
+import java.util.UUID;
 import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
-import org.killbill.automaton.OperationException;
 import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.account.api.AccountInternalApi;
@@ -46,9 +48,11 @@ import org.killbill.billing.util.dao.NonEntityDao;
 import org.killbill.bus.api.PersistentBus;
 import org.killbill.clock.Clock;
 import org.killbill.commons.locker.GlobalLocker;
+import org.killbill.commons.locker.LockFailedException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Objects;
 import com.google.inject.name.Named;
 
 import static org.killbill.billing.payment.glue.PaymentModule.PLUGIN_EXECUTOR_NAMED;
@@ -79,50 +83,62 @@ public class PaymentGatewayProcessor extends ProcessorBase {
     }
 
     public HostedPaymentPageFormDescriptor buildFormDescriptor(final Account account, final Iterable<PluginProperty> customFields, final Iterable<PluginProperty> properties, final CallContext callContext, final InternalCallContext internalCallContext) throws PaymentApiException {
-        try {
-            return paymentPluginFormDispatcher.dispatchWithTimeout(new CallableWithAccountLock<HostedPaymentPageFormDescriptor>(locker,
-                                                                                                                                account.getExternalKey(),
-                                                                                                                                new WithAccountLockCallback<HostedPaymentPageFormDescriptor>() {
+        return dispatchWithExceptionHandling(account,
+                                             new CallableWithAccountLock<HostedPaymentPageFormDescriptor, PaymentApiException>(locker,
+                                                                                                                               account.getExternalKey(),
+                                                                                                                               new WithAccountLockCallback<HostedPaymentPageFormDescriptor, PaymentApiException>() {
 
-                                                                                                                                    @Override
-                                                                                                                                    public HostedPaymentPageFormDescriptor doOperation() throws PaymentApiException {
-                                                                                                                                        final PaymentPluginApi plugin = getPaymentProviderPlugin(account, internalCallContext);
+                                                                                                                                   @Override
+                                                                                                                                   public HostedPaymentPageFormDescriptor doOperation() throws PaymentApiException {
+                                                                                                                                       final PaymentPluginApi plugin = getPaymentProviderPlugin(account, internalCallContext);
 
-                                                                                                                                        try {
-                                                                                                                                            return plugin.buildFormDescriptor(account.getId(), customFields, properties, callContext);
-                                                                                                                                        } catch (final RuntimeException e) {
-                                                                                                                                            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
-                                                                                                                                        } catch (final PaymentPluginApiException e) {
-                                                                                                                                            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
-                                                                                                                                        }
-                                                                                                                                    }
-                                                                                                                                }
-            ));
-        } catch (final TimeoutException e) {
-            throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_TIMEOUT, account.getId(), null);
-        } catch (final RuntimeException e) {
-            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
-        } catch (OperationException e) {
-            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
-        }
+                                                                                                                                       try {
+                                                                                                                                           return plugin.buildFormDescriptor(account.getId(), customFields, properties, callContext);
+                                                                                                                                       } catch (final RuntimeException e) {
+                                                                                                                                           throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
+                                                                                                                                       } catch (final PaymentPluginApiException e) {
+                                                                                                                                           throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_EXCEPTION, e.getErrorMessage());
+                                                                                                                                       }
+                                                                                                                                   }
+                                                                                                                               }),
+                                             paymentPluginFormDispatcher);
     }
 
     public GatewayNotification processNotification(final String notification, final String pluginName, final Iterable<PluginProperty> properties, final CallContext callContext) throws PaymentApiException {
+        return dispatchWithExceptionHandling(null,
+                                      new Callable<GatewayNotification>() {
+                                          @Override
+                                          public GatewayNotification call() throws PaymentApiException {
+                                              final PaymentPluginApi plugin = getPaymentPluginApi(pluginName);
+                                              try {
+                                                  return plugin.processNotification(notification, properties, callContext);
+                                              } catch (PaymentPluginApiException e) {
+                                                  throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_EXCEPTION, e.getErrorMessage());
+                                              }
+                                          }
+                                      }, paymentPluginNotificationDispatcher);
+    }
+
+    private static <ReturnType> ReturnType dispatchWithExceptionHandling(@Nullable final Account account, final Callable<ReturnType> callable, PluginDispatcher<ReturnType> pluginFormDispatcher) throws PaymentApiException {
+        final UUID accountId = account != null ? account.getId() : null;
+        final String accountExternalKey = account != null ? account.getExternalKey() : "";
         try {
-            return paymentPluginNotificationDispatcher.dispatchWithTimeout(new Callable<GatewayNotification>() {
-                                                                               @Override
-                                                                               public GatewayNotification call() throws Exception {
-                                                                                   final PaymentPluginApi plugin = getPaymentPluginApi(pluginName);
-                                                                                   return plugin.processNotification(notification, properties, callContext);
-                                                                               }
-                                                                           }
-                                                                          );
+            return pluginFormDispatcher.dispatchWithTimeout(callable);
         } catch (final TimeoutException e) {
-            throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_TIMEOUT, null, null);
-        } catch (final RuntimeException e) {
-            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
-        } catch (OperationException e) {
-            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
+            throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_TIMEOUT, accountId, null);
+        } catch (final InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new PaymentApiException(ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
+        } catch (final ExecutionException e) {
+            if (e.getCause() instanceof PaymentApiException) {
+                throw (PaymentApiException) e.getCause();
+            } else if (e.getCause() instanceof LockFailedException) {
+                final String format = String.format("Failed to lock account %s", accountExternalKey);
+                log.error(String.format(format), e);
+                throw new PaymentApiException(ErrorCode.PAYMENT_INTERNAL_ERROR, format);
+            } else {
+                throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
+            }
         }
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/PaymentMethodProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/PaymentMethodProcessor.java
index fe62885..f1913f4 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/PaymentMethodProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/PaymentMethodProcessor.java
@@ -93,7 +93,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
                                  final Iterable<PluginProperty> properties, final CallContext callContext, final InternalCallContext context)
             throws PaymentApiException {
         try {
-            return new WithAccountLock<UUID>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<UUID>() {
+            return new WithAccountLock<UUID, PaymentApiException>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<UUID, PaymentApiException>() {
 
                 @Override
                 public UUID doOperation() throws PaymentApiException {
@@ -121,7 +121,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
                 }
             });
         } catch (Exception e) {
-            throw new PaymentApiException(ErrorCode.__UNKNOWN_ERROR_CODE, "");
+            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
         }
     }
 
@@ -304,7 +304,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
                                      final Iterable<PluginProperty> properties, final CallContext callContext, final InternalCallContext context)
             throws PaymentApiException {
         try {
-            new WithAccountLock<Void>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<Void>() {
+            new WithAccountLock<Void, PaymentApiException>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<Void, PaymentApiException>() {
 
                 @Override
                 public Void doOperation() throws PaymentApiException {
@@ -340,14 +340,14 @@ public class PaymentMethodProcessor extends ProcessorBase {
                 }
             });
         } catch (Exception e) {
-            throw new PaymentApiException(ErrorCode.__UNKNOWN_ERROR_CODE, "");
+            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
         }
     }
 
     public void setDefaultPaymentMethod(final Account account, final UUID paymentMethodId, final Iterable<PluginProperty> properties, final CallContext callContext, final InternalCallContext context)
             throws PaymentApiException {
         try {
-            new WithAccountLock<Void>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<Void>() {
+            new WithAccountLock<Void, PaymentApiException>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<Void, PaymentApiException>() {
 
                 @Override
                 public Void doOperation() throws PaymentApiException {
@@ -370,7 +370,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
                 }
             });
         } catch (Exception e) {
-            throw new PaymentApiException(ErrorCode.__UNKNOWN_ERROR_CODE, e);
+            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
         }
     }
 
@@ -411,7 +411,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
         }
 
         try {
-            return new WithAccountLock<List<PaymentMethod>>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<List<PaymentMethod>>() {
+            return new WithAccountLock<List<PaymentMethod>, PaymentApiException>().processAccountWithLock(locker, account.getExternalKey(), new WithAccountLockCallback<List<PaymentMethod>, PaymentApiException>() {
 
                 @Override
                 public List<PaymentMethod> doOperation() throws PaymentApiException {
@@ -464,7 +464,7 @@ public class PaymentMethodProcessor extends ProcessorBase {
                 }
             });
         } catch (Exception e) {
-            throw new PaymentApiException(ErrorCode.__UNKNOWN_ERROR_CODE, e);
+            throw new PaymentApiException(e, ErrorCode.PAYMENT_INTERNAL_ERROR);
         }
     }
 
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/ProcessorBase.java b/payment/src/main/java/org/killbill/billing/payment/core/ProcessorBase.java
index 480282a..eb0f9f3 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/ProcessorBase.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/ProcessorBase.java
@@ -26,6 +26,7 @@ import java.util.concurrent.ExecutorService;
 
 import javax.annotation.Nullable;
 
+import org.killbill.automaton.OperationException;
 import org.killbill.billing.ErrorCode;
 import org.killbill.billing.ObjectType;
 import org.killbill.billing.account.api.Account;
@@ -166,71 +167,43 @@ public abstract class ProcessorBase {
         }
     }
 
-    protected Invoice rebalanceAndGetInvoice(final UUID accountId, final UUID invoiceId, final InternalCallContext context) throws InvoiceApiException {
-        invoiceApi.consumeExistingCBAOnAccountWithUnpaidInvoices(accountId, context);
-        final Invoice invoice = invoiceApi.getInvoiceById(invoiceId, context);
-        return invoice;
-    }
-
     protected TenantContext buildTenantContext(final InternalTenantContext context) {
         return context.toTenantContext(nonEntityDao.retrieveIdFromObject(context.getTenantRecordId(), ObjectType.TENANT));
     }
 
-    protected CallContext buildCallContext(final InternalCallContext context) {
-        return context.toCallContext(nonEntityDao.retrieveIdFromObject(context.getTenantRecordId(), ObjectType.TENANT));
-    }
-
     // TODO Rename - there is no lock!
-    public interface WithAccountLockCallback<T> {
-        public T doOperation() throws Exception;
+    public interface WithAccountLockCallback<ReturnType, ExceptionType extends Exception> {
+        public ReturnType doOperation() throws ExceptionType;
     }
 
-    public static class CallableWithAccountLock<T> implements Callable<T> {
+    public static class CallableWithAccountLock<ReturnType, ExceptionType extends Exception> implements Callable<ReturnType> {
 
         private final GlobalLocker locker;
         private final String accountExternalKey;
-        private final WithAccountLockCallback<T> callback;
+        private final WithAccountLockCallback<ReturnType, ExceptionType> callback;
 
         public CallableWithAccountLock(final GlobalLocker locker,
                                        final String accountExternalKey,
-                                       final WithAccountLockCallback<T> callback) {
+                                       final WithAccountLockCallback<ReturnType, ExceptionType> callback) {
             this.locker = locker;
             this.accountExternalKey = accountExternalKey;
             this.callback = callback;
         }
 
         @Override
-        public T call() throws Exception {
-            return new WithAccountLock<T>().processAccountWithLock(locker, accountExternalKey, callback);
-        }
-    }
-
-    public static class CallableWithoutAccountLock<T> implements Callable<T> {
-
-        private final WithAccountLockCallback<T> callback;
-
-        public CallableWithoutAccountLock(final WithAccountLockCallback<T> callback) {
-            this.callback = callback;
-        }
-
-        @Override
-        public T call() throws Exception {
-            return callback.doOperation();
+        public ReturnType call() throws ExceptionType, LockFailedException {
+            return new WithAccountLock<ReturnType, ExceptionType>().processAccountWithLock(locker, accountExternalKey, callback);
         }
     }
 
-    public static class WithAccountLock<T> {
+    public static class WithAccountLock<T, ExceptionType extends Exception> {
 
-        public T processAccountWithLock(final GlobalLocker locker, final String accountExternalKey, final WithAccountLockCallback<T> callback)
-                throws Exception {
+        public T processAccountWithLock(final GlobalLocker locker, final String accountExternalKey, final WithAccountLockCallback<T,ExceptionType > callback)
+                throws ExceptionType, LockFailedException {
             GlobalLock lock = null;
             try {
                 lock = locker.lockWithNumberOfTries(LockerType.ACCOUNT_FOR_INVOICE_PAYMENTS.toString(), accountExternalKey, NB_LOCK_TRY);
                 return callback.doOperation();
-            } catch (final LockFailedException e) {
-                final String format = String.format("Failed to lock account %s", accountExternalKey);
-                log.error(String.format(format), e);
-                throw new PaymentApiException(ErrorCode.PAYMENT_INTERNAL_ERROR, format);
             } finally {
                 if (lock != null) {
                     lock.release();
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
index 22f3213..c20a0fa 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/AuthorizeOperation.java
@@ -37,7 +37,7 @@ public class AuthorizeOperation extends DirectPaymentOperation {
     }
 
     @Override
-    protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+    protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
         logger.debug("Starting AUTHORIZE for payment {} ({} {})", directPaymentStateContext.getDirectPaymentId(), directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency());
         return plugin.authorizePayment(directPaymentStateContext.getAccount().getId(),
                                        directPaymentStateContext.getDirectPaymentId(),
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
index 21888b3..444b1aa 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CaptureOperation.java
@@ -37,7 +37,7 @@ public class CaptureOperation extends DirectPaymentOperation {
     }
 
     @Override
-    protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+    protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
         logger.debug("Starting CAPTURE for payment {} ({} {})", directPaymentStateContext.getDirectPaymentId(), directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency());
         return plugin.capturePayment(directPaymentStateContext.getAccount().getId(),
                                      directPaymentStateContext.getDirectPaymentId(),
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
index 059f67d..5c2c0c2 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/CreditOperation.java
@@ -37,7 +37,7 @@ public class CreditOperation extends DirectPaymentOperation {
     }
 
     @Override
-    protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+    protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
         logger.debug("Starting CREDIT for payment {} ({} {})", directPaymentStateContext.getDirectPaymentId(), directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency());
         return plugin.creditPayment(directPaymentStateContext.getAccount().getId(),
                                     directPaymentStateContext.getDirectPaymentId(),
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
index d7bc194..4934bba 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentAutomatonRunner.java
@@ -258,7 +258,6 @@ public class DirectPaymentAutomatonRunner {
 
             initialState.runOperation(operation, operationCallback, enteringStateCallback, leavingStateCallback);
         } catch (final MissingEntryException e) {
-            // TODO ErrorCode State Machine?
             throw new PaymentApiException(e.getCause(), ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
         } catch (final OperationException e) {
             if (e.getCause() == null) {
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentOperation.java
index 4386de1..9187437 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/DirectPaymentOperation.java
@@ -17,6 +17,9 @@
 
 package org.killbill.billing.payment.core.sm;
 
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
+
 import org.killbill.automaton.Operation.OperationCallback;
 import org.killbill.automaton.OperationException;
 import org.killbill.automaton.OperationResult;
@@ -29,8 +32,10 @@ import org.killbill.billing.payment.plugin.api.PaymentPluginApiException;
 import org.killbill.billing.payment.plugin.api.PaymentTransactionInfoPlugin;
 import org.killbill.commons.locker.GlobalLocker;
 
+import com.google.common.base.Objects;
+
 // Encapsulates the payment specific logic
-public abstract class DirectPaymentOperation extends PluginOperation implements OperationCallback {
+public abstract class DirectPaymentOperation extends OperationCallbackBase implements OperationCallback {
 
     protected final PaymentPluginApi plugin;
 
@@ -41,7 +46,6 @@ public abstract class DirectPaymentOperation extends PluginOperation implements 
         this.plugin = daoHelper.getPaymentProviderPlugin();
     }
 
-    protected abstract PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException;
 
     @Override
     public OperationResult doOperationCallback() throws OperationException {
@@ -52,10 +56,37 @@ public abstract class DirectPaymentOperation extends PluginOperation implements 
         }
     }
 
+    @Override
+    protected OperationException rewrapExecutionException(final DirectPaymentStateContext directPaymentStateContext, final ExecutionException e) {
+        if (e.getCause() instanceof PaymentPluginApiException) {
+            final Throwable realException = Objects.firstNonNull(e.getCause(), e);
+            logger.warn("Unsuccessful plugin call for account {}", directPaymentStateContext.getAccount().getExternalKey(), realException);
+            return new OperationException(realException, OperationResult.FAILURE);
+        } else /* if (e instanceof RuntimeException) */ {
+            logger.warn("Plugin call threw an exception for account {}", directPaymentStateContext.getAccount().getExternalKey(), e);
+            return new OperationException(e, OperationResult.EXCEPTION);
+        }
+    }
+
+    @Override
+    protected OperationException wrapTimeoutException(final DirectPaymentStateContext directPaymentStateContext, final TimeoutException e) {
+        logger.error("Plugin call TIMEOUT for account {}: {}", directPaymentStateContext.getAccount().getExternalKey(), e.getMessage());
+        return new OperationException(e, OperationResult.EXCEPTION);
+    }
+
+    @Override
+    protected OperationException wrapInterruptedException(final DirectPaymentStateContext directPaymentStateContext, final InterruptedException e) {
+        logger.error("Plugin call was interrupted for account {}: {}", directPaymentStateContext.getAccount().getExternalKey(), e.getMessage());
+        return new OperationException(e, OperationResult.EXCEPTION);
+    }
+
+    @Override
+    protected abstract PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException;
+
     private OperationResult doOperationCallbackWithDispatchAndAccountLock() throws OperationException {
-        return dispatchWithTimeout(new WithAccountLockCallback<OperationResult>() {
+        return dispatchWithAccountLockAndTimeout(new WithAccountLockCallback<OperationResult, OperationException>() {
             @Override
-            public OperationResult doOperation() throws Exception {
+            public OperationResult doOperation() throws OperationException {
                 return doSimpleOperationCallback();
             }
         });
@@ -71,14 +102,13 @@ public abstract class DirectPaymentOperation extends PluginOperation implements 
 
     private OperationResult doOperation() throws PaymentApiException {
         try {
-            final PaymentTransactionInfoPlugin paymentInfoPlugin = doPluginOperation();
+            final PaymentTransactionInfoPlugin paymentInfoPlugin = doCallSpecificOperationCallback();
 
             directPaymentStateContext.setPaymentInfoPlugin(paymentInfoPlugin);
 
             return processPaymentInfoPlugin();
         } catch (final PaymentPluginApiException e) {
-            // We don't care about the ErrorCode since it will be unwrapped
-            throw new PaymentApiException(e, ErrorCode.__UNKNOWN_ERROR_CODE, "");
+            throw new PaymentApiException(ErrorCode.PAYMENT_PLUGIN_EXCEPTION, e.getErrorMessage());
         }
     }
 
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
index 6161f11..1257847 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RefundOperation.java
@@ -37,7 +37,7 @@ public class RefundOperation extends DirectPaymentOperation {
     }
 
     @Override
-    protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+    protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
         logger.debug("Starting REFUND for payment {} ({} {})", directPaymentStateContext.getDirectPaymentId(), directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency());
         return plugin.refundPayment(directPaymentStateContext.getAccount().getId(),
                                     directPaymentStateContext.getDirectPaymentId(),
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryAuthorizeOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryAuthorizeOperationCallback.java
index 3f37aba..5993143 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryAuthorizeOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryAuthorizeOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryAuthorizeOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createAuthorization(directPaymentStateContext.account, directPaymentStateContext.paymentMethodId, directPaymentStateContext.directPaymentId, directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency(), directPaymentStateContext.directPaymentExternalKey, directPaymentStateContext.directPaymentTransactionExternalKey, false, directPaymentStateContext.getProperties(), directPaymentStateContext.callContext, directPaymentStateContext.internalCallContext);
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCaptureOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCaptureOperationCallback.java
index 1b683ba..508a83e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCaptureOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCaptureOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryCaptureOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createCapture(directPaymentStateContext.account, directPaymentStateContext.directPaymentId,
                                                     directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency(),
                                                     directPaymentStateContext.directPaymentTransactionExternalKey,
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCreditOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCreditOperationCallback.java
index 2e5844c..6f25017 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCreditOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryCreditOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryCreditOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createCredit(directPaymentStateContext.account, directPaymentStateContext.paymentMethodId, directPaymentStateContext.directPaymentId, directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency(), directPaymentStateContext.directPaymentExternalKey, directPaymentStateContext.directPaymentTransactionExternalKey, false, directPaymentStateContext.getProperties(), directPaymentStateContext.callContext, directPaymentStateContext.internalCallContext);
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryOperationCallback.java
index 47c69dd..3330cd6 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryOperationCallback.java
@@ -18,6 +18,8 @@ package org.killbill.billing.payment.core.sm;
 
 import java.math.BigDecimal;
 import java.util.UUID;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeoutException;
 
 import javax.annotation.Nullable;
 
@@ -25,7 +27,6 @@ import org.joda.time.DateTime;
 import org.killbill.automaton.Operation.OperationCallback;
 import org.killbill.automaton.OperationException;
 import org.killbill.automaton.OperationResult;
-import org.killbill.billing.ErrorCode;
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.callcontext.DefaultCallContext;
 import org.killbill.billing.catalog.api.Currency;
@@ -49,7 +50,7 @@ import org.killbill.commons.locker.GlobalLocker;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-public abstract class RetryOperationCallback extends PluginOperation implements OperationCallback {
+public abstract class RetryOperationCallback extends OperationCallbackBase implements OperationCallback {
 
     protected final DirectPaymentProcessor directPaymentProcessor;
     private final OSGIServiceRegistration<PaymentControlPluginApi> paymentControlPluginRegistry;
@@ -62,37 +63,14 @@ public abstract class RetryOperationCallback extends PluginOperation implements 
         this.paymentControlPluginRegistry = retryPluginRegistry;
     }
 
-    private PriorPaymentControlResult getPluginResult(final String pluginName, final PaymentControlContext paymentControlContext) throws PaymentControlApiException {
-
-        final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
-        final PriorPaymentControlResult result = plugin.priorCall(paymentControlContext);
-        return result;
-    }
 
-    private DateTime getNextRetryDate(final String pluginName, final PaymentControlContext paymentControlContext) {
-        try {
-            final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
-            final FailureCallResult result = plugin.onFailureCall(paymentControlContext);
-            return result.getNextRetryDate();
-        } catch (PaymentControlApiException e) {
-            logger.warn("Plugin " + pluginName + " failed to return next retryDate for payment " + paymentControlContext.getPaymentExternalKey(), e);
-            return null;
-        }
-    }
-
-    private void onCompletion(final String pluginName, final PaymentControlContext paymentControlContext) {
-        final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
-        try {
-            plugin.onSuccessCall(paymentControlContext);
-        } catch (PaymentControlApiException e) {
-            logger.warn("Plugin " + pluginName + " failed to complete onCompletion call for " + paymentControlContext.getPaymentExternalKey(), e);
-        }
-    }
+    @Override
+    protected abstract DirectPayment doCallSpecificOperationCallback() throws PaymentApiException;
 
     @Override
     public OperationResult doOperationCallback() throws OperationException {
 
-        return dispatchWithTimeout(new WithAccountLockCallback<OperationResult>() {
+        return dispatchWithAccountLockAndTimeout(new WithAccountLockCallback<OperationResult, OperationException>() {
 
             @Override
             public OperationResult doOperation() throws OperationException {
@@ -110,18 +88,19 @@ public abstract class RetryOperationCallback extends PluginOperation implements 
                                                                                                      retryableDirectPaymentStateContext.isApiPayment(),
                                                                                                      directPaymentStateContext.callContext);
 
-                // Note that we are using OperationResult.EXCEPTION result to transition to final ABORTED state -- see RetryStates.xml
                 final PriorPaymentControlResult pluginResult;
                 try {
                     pluginResult = getPluginResult(retryableDirectPaymentStateContext.getPluginName(), paymentControlContext);
                     if (pluginResult.isAborted()) {
+                        // Transition to ABORTED
                         return OperationResult.EXCEPTION;
                     }
                 } catch (PaymentControlApiException e) {
+                    // Transition to ABORTED and throw PaymentControlApiException to caller.
                     throw new OperationException(e, OperationResult.EXCEPTION);
                 }
 
-                boolean success = false;
+                boolean success;
                 try {
                     // Adjust amount with value returned by plugin if necessary
                     if (directPaymentStateContext.getAmount() == null ||
@@ -129,7 +108,7 @@ public abstract class RetryOperationCallback extends PluginOperation implements 
                         ((RetryableDirectPaymentStateContext) directPaymentStateContext).setAmount(pluginResult.getAdjustedAmount());
                     }
 
-                    final DirectPayment result = doPluginOperation();
+                    final DirectPayment result = doCallSpecificOperationCallback();
                     ((RetryableDirectPaymentStateContext) directPaymentStateContext).setResult(result);
                     final DirectPaymentTransaction transaction = ((RetryableDirectPaymentStateContext) directPaymentStateContext).getCurrentTransaction();
 
@@ -150,38 +129,88 @@ public abstract class RetryOperationCallback extends PluginOperation implements 
                                                                                                                     directPaymentStateContext.callContext);
 
                         onCompletion(retryableDirectPaymentStateContext.getPluginName(), updatedPaymentControlContext);
+                        return OperationResult.SUCCESS;
                     } else {
-                        // Error code?
-                        throwAndupdateRetryDateOnFailureOrException(retryableDirectPaymentStateContext, paymentControlContext, new PaymentApiException(ErrorCode.PAYMENT_INTERNAL_ERROR, "Plugin ERROR"));
+                        // Return an ABORTED/FAILURE state based on the retry result.
+                        return getOperationResultAndSetContext(retryableDirectPaymentStateContext, paymentControlContext);
                     }
-
                 } catch (PaymentApiException e) {
-                    throwAndupdateRetryDateOnFailureOrException(retryableDirectPaymentStateContext, paymentControlContext, e);
-                } catch (OperationException e) {
-                    // We need this catch clause to make sure this is not caught by the next more generic clause Exception
+                    // Wrap PaymentApiException, and throw a new OperationException with an ABORTED/FAILURE state based on the retry result.
+                    throw new OperationException(e, getOperationResultAndSetContext(retryableDirectPaymentStateContext, paymentControlContext));
+                } catch (RuntimeException e) {
+                    // Attempts to set the retry date in context if needed.
+                    getOperationResultAndSetContext(retryableDirectPaymentStateContext, paymentControlContext);
                     throw e;
-                } catch (Exception e) {
-                    // STEPH Any other exception we abort the retry logic, unclear if this is the *right* approach..
-                    throw new OperationException(e, OperationResult.EXCEPTION);
-                }
-                return OperationResult.SUCCESS;
-            }
-
-            private void throwAndupdateRetryDateOnFailureOrException(final RetryableDirectPaymentStateContext retryableDirectPaymentStateContext, final PaymentControlContext paymentControlContext,
-                                                                     @Nullable final PaymentApiException e) throws OperationException {
-                final DateTime retryDate = getNextRetryDate(retryableDirectPaymentStateContext.getPluginName(), paymentControlContext);
-                if (retryDate == null) {
-                    // STEPH only throw if e is not null
-                    throw new OperationException(e, OperationResult.EXCEPTION);
-                } else {
-                    // STEPH only throw if e is not null
-                    ((RetryableDirectPaymentStateContext) directPaymentStateContext).setRetryDate(retryDate);
-                    throw new OperationException(e, OperationResult.FAILURE);
                 }
             }
         });
     }
 
+    @Override
+    protected OperationException rewrapExecutionException(final DirectPaymentStateContext directPaymentStateContext, final ExecutionException e) {
+        if (e.getCause() instanceof OperationException) {
+            return (OperationException) e.getCause();
+        } else /* most probably RuntimeException */ {
+            logger.warn("RetryOperationCallback failed for account {}", directPaymentStateContext.getAccount().getExternalKey(), e);
+            return new OperationException(e, getOperationResultOnException(directPaymentStateContext));
+        }
+    }
+
+    @Override
+    protected OperationException wrapTimeoutException(final DirectPaymentStateContext directPaymentStateContext, final TimeoutException e) {
+        logger.error("RetryOperationCallback call TIMEOUT for account {}: {}", directPaymentStateContext.getAccount().getExternalKey(), e.getMessage());
+        return new OperationException(e, getOperationResultOnException(directPaymentStateContext));
+    }
+
+    @Override
+    protected OperationException wrapInterruptedException(final DirectPaymentStateContext directPaymentStateContext, final InterruptedException e) {
+        logger.error("RetryOperationCallback call was interrupted for account {}: {}", directPaymentStateContext.getAccount().getExternalKey(), e.getMessage());
+        return new OperationException(e, getOperationResultOnException(directPaymentStateContext));
+    }
+
+    private OperationResult getOperationResultOnException(final DirectPaymentStateContext directPaymentStateContext) {
+        final RetryableDirectPaymentStateContext retryableDirectPaymentStateContext = (RetryableDirectPaymentStateContext) directPaymentStateContext;
+        final OperationResult operationResult = retryableDirectPaymentStateContext.getRetryDate() != null ? OperationResult.FAILURE : OperationResult.EXCEPTION;
+        return operationResult;
+    }
+
+    private PriorPaymentControlResult getPluginResult(final String pluginName, final PaymentControlContext paymentControlContext) throws PaymentControlApiException {
+
+        final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
+        final PriorPaymentControlResult result = plugin.priorCall(paymentControlContext);
+        return result;
+    }
+
+    private OperationResult getOperationResultAndSetContext(final RetryableDirectPaymentStateContext retryableDirectPaymentStateContext, final PaymentControlContext paymentControlContext) {
+        final DateTime retryDate = getNextRetryDate(retryableDirectPaymentStateContext.getPluginName(), paymentControlContext);
+        if (retryDate != null) {
+            ((RetryableDirectPaymentStateContext) directPaymentStateContext).setRetryDate(retryDate);
+            return OperationResult.FAILURE;
+        } else {
+            return OperationResult.EXCEPTION;
+        }
+    }
+
+    private DateTime getNextRetryDate(final String pluginName, final PaymentControlContext paymentControlContext) {
+        try {
+            final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
+            final FailureCallResult result = plugin.onFailureCall(paymentControlContext);
+            return result.getNextRetryDate();
+        } catch (PaymentControlApiException e) {
+            logger.warn("Plugin " + pluginName + " failed to return next retryDate for payment " + paymentControlContext.getPaymentExternalKey(), e);
+            return null;
+        }
+    }
+
+    private void onCompletion(final String pluginName, final PaymentControlContext paymentControlContext) {
+        final PaymentControlPluginApi plugin = paymentControlPluginRegistry.getServiceForName(pluginName);
+        try {
+            plugin.onSuccessCall(paymentControlContext);
+        } catch (PaymentControlApiException e) {
+            logger.warn("Plugin " + pluginName + " failed to complete onCompletion call for " + paymentControlContext.getPaymentExternalKey(), e);
+        }
+    }
+
     public class DefaultPaymentControlContext extends DefaultCallContext implements PaymentControlContext {
 
         private final Account account;
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryPurchaseOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryPurchaseOperationCallback.java
index 759eccd..6127cef 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryPurchaseOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryPurchaseOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryPurchaseOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createPurchase(directPaymentStateContext.account, directPaymentStateContext.paymentMethodId, directPaymentStateContext.directPaymentId, directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency(), directPaymentStateContext.directPaymentExternalKey, directPaymentStateContext.directPaymentTransactionExternalKey, false, directPaymentStateContext.getProperties(), directPaymentStateContext.callContext, directPaymentStateContext.internalCallContext);
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryRefundOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryRefundOperationCallback.java
index dc812ae..414f2f8 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryRefundOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryRefundOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryRefundOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createRefund(directPaymentStateContext.account, directPaymentStateContext.directPaymentId, directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency(), directPaymentStateContext.directPaymentTransactionExternalKey, false, directPaymentStateContext.getProperties(), directPaymentStateContext.callContext, directPaymentStateContext.internalCallContext);
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryVoidOperationCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryVoidOperationCallback.java
index b274f37..45ea703 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryVoidOperationCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/RetryVoidOperationCallback.java
@@ -32,7 +32,7 @@ public class RetryVoidOperationCallback extends RetryOperationCallback {
     }
 
     @Override
-    protected DirectPayment doPluginOperation() throws PaymentApiException {
+    protected DirectPayment doCallSpecificOperationCallback() throws PaymentApiException {
         return directPaymentProcessor.createVoid(directPaymentStateContext.account, directPaymentStateContext.directPaymentId, directPaymentStateContext.directPaymentTransactionExternalKey, false, directPaymentStateContext.getProperties(), directPaymentStateContext.callContext, directPaymentStateContext.internalCallContext);
     }
 }
diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
index 20ed426..af6bd61 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/VoidOperation.java
@@ -37,7 +37,7 @@ public class VoidOperation extends DirectPaymentOperation {
     }
 
     @Override
-    protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+    protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
         logger.debug("Starting VOID for payment {} ({} {})", directPaymentStateContext.getDirectPaymentId(), directPaymentStateContext.getAmount(), directPaymentStateContext.getCurrency());
         return plugin.voidPayment(directPaymentStateContext.getAccount().getId(),
                                   directPaymentStateContext.getDirectPaymentId(),
diff --git a/payment/src/main/java/org/killbill/billing/payment/dispatcher/PluginDispatcher.java b/payment/src/main/java/org/killbill/billing/payment/dispatcher/PluginDispatcher.java
index a7f81e6..519107e 100644
--- a/payment/src/main/java/org/killbill/billing/payment/dispatcher/PluginDispatcher.java
+++ b/payment/src/main/java/org/killbill/billing/payment/dispatcher/PluginDispatcher.java
@@ -25,17 +25,7 @@ import java.util.concurrent.Future;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 
-import org.killbill.automaton.OperationException;
-import org.killbill.billing.ErrorCode;
-import org.killbill.billing.payment.api.PaymentApiException;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
-
-import com.google.common.base.Objects;
-
-public class PluginDispatcher<T> {
-
-    private static final Logger log = LoggerFactory.getLogger(PluginDispatcher.class);
+public class PluginDispatcher<ReturnType> {
 
     private final TimeUnit DEEFAULT_PLUGIN_TIMEOUT_UNIT = TimeUnit.SECONDS;
 
@@ -48,29 +38,13 @@ public class PluginDispatcher<T> {
     }
 
     // TODO Once we switch fully to automata, should this throw PaymentPluginApiException instead?
-    public T dispatchWithTimeout(final Callable<T> task) throws PaymentApiException, OperationException, TimeoutException {
+    public ReturnType dispatchWithTimeout(final Callable<ReturnType> task) throws TimeoutException, ExecutionException, InterruptedException {
         return dispatchWithTimeout(task, timeoutSeconds, DEEFAULT_PLUGIN_TIMEOUT_UNIT);
     }
 
-    public T dispatchWithTimeout(final Callable<T> task, final long timeout, final TimeUnit unit)
-            throws PaymentApiException, TimeoutException, OperationException {
-
-        try {
-            final Future<T> future = executor.submit(task);
-            return future.get(timeout, unit);
-        } catch (final ExecutionException e) {
-            if (e.getCause() instanceof PaymentApiException) {
-                throw (PaymentApiException) e.getCause();
-            } else if (e.getCause() instanceof OperationException) {
-                throw (OperationException) e.getCause();
-            } else if (e.getCause() instanceof RuntimeException) {
-                throw (RuntimeException) e.getCause();
-            } else {
-                throw new PaymentApiException(Objects.firstNonNull(e.getCause(), e), ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
-            }
-        } catch (final InterruptedException e) {
-            Thread.currentThread().interrupt();
-            throw new PaymentApiException(ErrorCode.PAYMENT_INTERNAL_ERROR, e.getMessage());
-        }
+    public ReturnType dispatchWithTimeout(final Callable<ReturnType> task, final long timeout, final TimeUnit unit)
+            throws TimeoutException, ExecutionException, InterruptedException {
+        final Future<ReturnType> future = executor.submit(task);
+        return future.get(timeout, unit);
     }
 }
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
index 3014415..20fcb90 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestDirectPaymentOperation.java
@@ -138,7 +138,7 @@ public class TestDirectPaymentOperation extends PaymentTestSuiteNoDB {
         }
 
         @Override
-        protected PaymentTransactionInfoPlugin doPluginOperation() throws PaymentPluginApiException {
+        protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
             if (paymentInfoPlugin == null) {
                 throw new RuntimeException("Exception expected by test");
             } else {
diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
index 9c20ade..b86ceef 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPluginOperation.java
@@ -39,6 +39,8 @@ import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionType;
 import org.killbill.billing.payment.core.ProcessorBase.WithAccountLockCallback;
 import org.killbill.billing.payment.dispatcher.PluginDispatcher;
+import org.killbill.billing.payment.plugin.api.PaymentPluginApiException;
+import org.killbill.billing.payment.plugin.api.PaymentTransactionInfoPlugin;
 import org.killbill.commons.locker.GlobalLocker;
 import org.killbill.commons.locker.memory.MemoryGlobalLocker;
 import org.mockito.Mockito;
@@ -64,7 +66,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
         testLocking(true);
     }
 
-    // STEPH the test now fails because the logic has been changed; we don't check in the dispatchWithTimeout
+    // STEPH the test now fails because the logic has been changed; we don't check in the dispatchWithAccountLockAndTimeout
     // method to see whether account should be locked or not. Instead we either (dispatch AND lock) OR
     // ! (dispatch AND lock)
     @Test(groups = "fast", enabled=false)
@@ -75,10 +77,10 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
     @Test(groups = "fast")
     public void testOperationTimeout() throws Exception {
         final CallbackTest callback = new CallbackTest(Integer.MAX_VALUE);
-        final PluginOperation pluginOperation = getPluginOperation(false, 1);
+        final DirectPaymentOperation pluginOperation = getPluginOperation(false, 1);
 
         try {
-            pluginOperation.dispatchWithTimeout(callback);
+            pluginOperation.dispatchWithAccountLockAndTimeout(callback);
             Assert.fail();
         } catch (final OperationException e) {
             Assert.assertEquals(e.getOperationResult(), OperationResult.EXCEPTION);
@@ -89,10 +91,10 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
     @Test(groups = "fast")
     public void testOperationThrowsPaymentApiException() throws Exception {
         final CallbackTest callback = new CallbackTest(new PaymentApiException(ErrorCode.__UNKNOWN_ERROR_CODE));
-        final PluginOperation pluginOperation = getPluginOperation();
+        final DirectPaymentOperation pluginOperation = getPluginOperation();
 
         try {
-            pluginOperation.dispatchWithTimeout(callback);
+            pluginOperation.dispatchWithAccountLockAndTimeout(callback);
             Assert.fail();
         } catch (final OperationException e) {
             Assert.assertEquals(e.getOperationResult(), OperationResult.FAILURE);
@@ -103,10 +105,10 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
     @Test(groups = "fast")
     public void testOperationThrowsRuntimeException() throws Exception {
         final CallbackTest callback = new CallbackTest(new NullPointerException("Expected for the test"));
-        final PluginOperation pluginOperation = getPluginOperation();
+        final DirectPaymentOperation pluginOperation = getPluginOperation();
 
         try {
-            pluginOperation.dispatchWithTimeout(callback);
+            pluginOperation.dispatchWithAccountLockAndTimeout(callback);
             Assert.fail();
         } catch (final OperationException e) {
             Assert.assertEquals(e.getOperationResult(), OperationResult.EXCEPTION);
@@ -117,7 +119,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
     private void testLocking(final boolean withAccountLock) throws Exception {
         final Semaphore available = new Semaphore(1, true);
         final CallbackTest callback = new CallbackTest(available);
-        final PluginOperation pluginOperation = getPluginOperation(withAccountLock);
+        final DirectPaymentOperation pluginOperation = getPluginOperation(withAccountLock);
 
         // Take the only permit
         available.acquire();
@@ -156,7 +158,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
         Assert.assertEquals(callback.getRunCount(), withAccountLock ? 1 : 2);
     }
 
-    private void runPluginOperationInBackground(final PluginOperation pluginOperation, final CallbackTest callback, final boolean shouldFailBecauseOfLockFailure) throws Exception {
+    private void runPluginOperationInBackground(final DirectPaymentOperation pluginOperation, final CallbackTest callback, final boolean shouldFailBecauseOfLockFailure) throws Exception {
         final AtomicBoolean threadStarted = new AtomicBoolean(false);
         final Thread t1 = new Thread(new Runnable() {
             @Override
@@ -165,7 +167,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
 
                 if (shouldFailBecauseOfLockFailure) {
                     try {
-                        pluginOperation.dispatchWithTimeout(callback);
+                        pluginOperation.dispatchWithAccountLockAndTimeout(callback);
                         Assert.fail();
                     } catch (final OperationException e) {
                         Assert.assertTrue(e.getCause() instanceof PaymentApiException);
@@ -174,7 +176,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
                     }
                 } else {
                     try {
-                        pluginOperation.dispatchWithTimeout(callback);
+                        pluginOperation.dispatchWithAccountLockAndTimeout(callback);
                     } catch (final OperationException e) {
                         Assert.fail(e.getMessage());
                     }
@@ -188,15 +190,15 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
         Awaitility.await().untilTrue(threadStarted);
     }
 
-    private PluginOperation getPluginOperation() throws PaymentApiException {
+    private DirectPaymentOperation getPluginOperation() throws PaymentApiException {
         return getPluginOperation(false);
     }
 
-    private PluginOperation getPluginOperation(final boolean shouldLockAccount) throws PaymentApiException {
+    private DirectPaymentOperation getPluginOperation(final boolean shouldLockAccount) throws PaymentApiException {
         return getPluginOperation(shouldLockAccount, Integer.MAX_VALUE);
     }
 
-    private PluginOperation getPluginOperation(final boolean shouldLockAccount, final int timeoutSeconds) throws PaymentApiException {
+    private DirectPaymentOperation getPluginOperation(final boolean shouldLockAccount, final int timeoutSeconds) throws PaymentApiException {
         final PluginDispatcher<OperationResult> paymentPluginDispatcher = new PluginDispatcher<OperationResult>(timeoutSeconds, Executors.newCachedThreadPool());
 
         final DirectPaymentStateContext directPaymentStateContext = new DirectPaymentStateContext(UUID.randomUUID(),
@@ -214,7 +216,7 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
         return new PluginOperationTest(locker, paymentPluginDispatcher, directPaymentStateContext);
     }
 
-    private static final class CallbackTest implements WithAccountLockCallback<OperationResult> {
+    private static final class CallbackTest implements WithAccountLockCallback<OperationResult, PaymentApiException> {
 
         private final AtomicInteger runCount = new AtomicInteger(0);
 
@@ -282,14 +284,15 @@ public class TestPluginOperation extends PaymentTestSuiteNoDB {
         }
     }
 
-    private static final class PluginOperationTest extends PluginOperation {
+    private static final class PluginOperationTest extends DirectPaymentOperation {
 
-        protected PluginOperationTest(final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher, final DirectPaymentStateContext directPaymentStateContext) {
-            super(locker, paymentPluginDispatcher, directPaymentStateContext);
+        protected PluginOperationTest(final GlobalLocker locker, final PluginDispatcher<OperationResult> paymentPluginDispatcher, final DirectPaymentStateContext directPaymentStateContext) throws PaymentApiException {
+            // STEPH null ?
+            super(null, locker, paymentPluginDispatcher, directPaymentStateContext);
         }
 
         @Override
-        protected <PluginResult> PluginResult doPluginOperation() throws Exception {
+        protected PaymentTransactionInfoPlugin doCallSpecificOperationCallback() throws PaymentPluginApiException {
             return null;
         }
     }
diff --git a/payment/src/test/java/org/killbill/billing/payment/dispatcher/TestPluginDispatcher.java b/payment/src/test/java/org/killbill/billing/payment/dispatcher/TestPluginDispatcher.java
index 3a69e46..e3fc071 100644
--- a/payment/src/test/java/org/killbill/billing/payment/dispatcher/TestPluginDispatcher.java
+++ b/payment/src/test/java/org/killbill/billing/payment/dispatcher/TestPluginDispatcher.java
@@ -17,6 +17,7 @@
 package org.killbill.billing.payment.dispatcher;
 
 import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.Executors;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
@@ -46,10 +47,10 @@ public class TestPluginDispatcher extends PaymentTestSuiteNoDB {
             Assert.fail("Failed : should have had Timeout exception");
         } catch (final TimeoutException e) {
             gotIt = true;
-        } catch (final PaymentApiException e) {
+        } catch (InterruptedException e) {
+            Assert.fail("Failed : should have had Timeout exception");
+        } catch (ExecutionException e) {
             Assert.fail("Failed : should have had Timeout exception");
-        } catch (OperationException e) {
-            Assert.fail("Failed : should have had OperationException exception");
         }
         Assert.assertTrue(gotIt);
     }
@@ -67,10 +68,14 @@ public class TestPluginDispatcher extends PaymentTestSuiteNoDB {
             Assert.fail("Failed : should have had Timeout exception");
         } catch (final TimeoutException e) {
             Assert.fail("Failed : should have had PaymentApiException exception");
-        } catch (final PaymentApiException e) {
-            gotIt = true;
-        } catch (OperationException e) {
-            Assert.fail("Failed : should have had OperationException exception");
+        } catch (InterruptedException e) {
+            Assert.fail("Failed : should have had PaymentApiException exception");
+        } catch (ExecutionException e) {
+            if (e.getCause() instanceof PaymentApiException) {
+                gotIt = true;
+            } else {
+                Assert.fail("Failed : should have had PaymentApiException exception");
+            }
         }
         Assert.assertTrue(gotIt);
     }
@@ -88,12 +93,16 @@ public class TestPluginDispatcher extends PaymentTestSuiteNoDB {
             Assert.fail("Failed : should have had Timeout exception");
         } catch (final TimeoutException e) {
             Assert.fail("Failed : should have had RuntimeException exception");
-        } catch (final PaymentApiException e) {
-            Assert.fail("Failed : should have had RuntimeException exception");
         } catch (final RuntimeException e) {
-            gotIt = true;
-        } catch (OperationException e) {
-            Assert.fail("Failed : should have had OperationException exception");
+            Assert.fail("Failed : should have had RuntimeException (wrapped in an ExecutionException)");
+        } catch (InterruptedException e) {
+            e.printStackTrace();
+        } catch (ExecutionException e) {
+            if (e.getCause() instanceof RuntimeException) {
+                gotIt = true;
+            } else {
+                Assert.fail("Failed : should have had RuntimeException exception");
+            }
         }
         Assert.assertTrue(gotIt);
     }