killbill-memoizeit

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlInitiated.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlInitiated.java
index 54783cb..4ead401 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlInitiated.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/DefaultControlInitiated.java
@@ -79,7 +79,7 @@ public class DefaultControlInitiated implements LeavingStateCallback {
                 //
                 // We don't serialize any properties at this stage to avoid serializing sensitive information.
                 // However, if after going through the control plugins, the attempt end up in RETRIED state,
-                // the properties will be serialized in the enteringState( callback (any plugin that sets a
+                // the properties will be serialized in the enteringState callback (any plugin that sets a
                 // retried date is responsible to correctly remove sensitive information such as CVV, ...)
                 //
                 final byte[] serializedProperties = PluginPropertySerializer.serialize(ImmutableList.<PluginProperty>of());

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/PaymentStateControlContext.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/PaymentStateControlContext.java
index d9af3a9..8b51735 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/control/PaymentStateControlContext.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/control/PaymentStateControlContext.java
@@ -72,17 +72,6 @@ public class PaymentStateControlContext extends PaymentStateContext {
         this.result = result;
     }
 
-    public void setAmount(final BigDecimal adjustedAmount) {
-        this.amount = adjustedAmount;
-    }
-
-    public void setCurrency(final Currency currency) {
-        this.currency = currency;
-    }
-
-    public void setProperties(final Iterable<PluginProperty> properties) {
-        this.properties = properties;
-    }
 
     public PaymentTransaction getCurrentTransaction() {
         if (result == null || result.getTransactions() == null) {
@@ -91,7 +80,7 @@ public class PaymentStateControlContext extends PaymentStateContext {
         return Iterables.tryFind(result.getTransactions(), new Predicate<PaymentTransaction>() {
             @Override
             public boolean apply(final PaymentTransaction input) {
-                return ((DefaultPaymentTransaction) input).getAttemptId().equals(attemptId);
+                return ((DefaultPaymentTransaction) input).getAttemptId().equals(getAttemptId());
             }
         }).orNull();
     }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentEnteringStateCallback.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentEnteringStateCallback.java
index 9860f6b..f5dd2f2 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentEnteringStateCallback.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentEnteringStateCallback.java
@@ -17,8 +17,6 @@
 
 package org.killbill.billing.payment.core.sm.payments;
 
-import javax.annotation.Nullable;
-
 import org.killbill.automaton.Operation;
 import org.killbill.automaton.OperationResult;
 import org.killbill.automaton.State;
@@ -52,10 +50,6 @@ public abstract class PaymentEnteringStateCallback implements EnteringStateCallb
     public void enteringState(final State newState, final Operation.OperationCallback operationCallback, final OperationResult operationResult, final LeavingStateCallback leavingStateCallback) {
         logger.debug("Entering state {} with result {}", newState.getName(), operationResult);
 
-        if (paymentStateContext.isSkipOperationForUnknownTransaction()) {
-            return;
-        }
-
         // If the transaction was not created -- for instance we had an exception in leavingState callback then we bail; if not, then update state:
         if (paymentStateContext.getPaymentTransactionModelDao() != null && paymentStateContext.getPaymentTransactionModelDao().getId() != null) {
             final PaymentTransactionInfoPlugin paymentInfoPlugin = paymentStateContext.getPaymentTransactionInfoPlugin();

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
index 4896409..7d31cbb 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/payments/PaymentOperation.java
@@ -69,10 +69,6 @@ public abstract class PaymentOperation extends OperationCallbackBase<PaymentTran
     @Override
     public OperationResult doOperationCallback() throws OperationException {
 
-        if (paymentStateContext.isSkipOperationForUnknownTransaction()) {
-            return OperationResult.SUCCESS;
-        }
-
         try {
             this.plugin = daoHelper.getPaymentProviderPlugin();
 

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateContext.java b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateContext.java
index cbe72e7..8975b22 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateContext.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/sm/PaymentStateContext.java
@@ -37,31 +37,40 @@ import com.google.common.collect.ImmutableList;
 
 public class PaymentStateContext {
 
-    // HACK
-    protected UUID paymentMethodId;
-    protected UUID attemptId;
-
-    // Stateful objects created by the callbacks and passed to the other following callbacks in the automaton
-    protected List<PaymentTransactionModelDao> onLeavingStateExistingTransactions;
-    protected PaymentTransactionModelDao paymentTransactionModelDao;
-    protected PaymentTransactionInfoPlugin paymentTransactionInfoPlugin;
-    protected BigDecimal amount;
-    protected String paymentExternalKey;
-    protected String paymentTransactionExternalKey;
-    protected Currency currency;
-    protected Iterable<PluginProperty> properties;
-    protected boolean skipOperationForUnknownTransaction;
-
-    // Can be updated later via paymentTransactionModelDao (e.g. for auth or purchase)
-    protected final UUID paymentId;
-    protected final UUID transactionId;
-    protected final Account account;
-    protected final TransactionType transactionType;
-    protected final boolean shouldLockAccountAndDispatch;
-    protected final InternalCallContext internalCallContext;
-    protected final CallContext callContext;
-    protected final boolean isApiPayment;
-    protected final OperationResult overridePluginOperationResult;
+
+    // The following fields (paymentId, transactionId, amount, currency) may take their value from the paymentTransactionModelDao *when they are not already set*
+    private PaymentTransactionModelDao paymentTransactionModelDao;
+    // Initialized in CTOR or only set through paymentTransactionModelDao
+    private UUID paymentId;
+    private UUID transactionId;
+
+    // Can be overriden by control plugin
+    private BigDecimal amount;
+    private Currency currency;
+    private UUID paymentMethodId;
+    private Iterable<PluginProperty> properties;
+
+    // Set in the doOperationCallback when coming back from payment plugin
+    private PaymentTransactionInfoPlugin paymentTransactionInfoPlugin;
+
+    // Set in the control layer in the leavingState callback
+    private String paymentExternalKey;
+    private String paymentTransactionExternalKey;
+
+    // Set in the control layer after creating the attempt in the enteringState callback
+    private UUID attemptId;
+
+    // This is purely a performance improvement to avoid fetching the existing transactions for that payment throughout the state machine
+    private List<PaymentTransactionModelDao> onLeavingStateExistingTransactions;
+
+    // Immutable
+    private final Account account;
+    private final TransactionType transactionType;
+    private final boolean shouldLockAccountAndDispatch;
+    private final OperationResult overridePluginOperationResult;
+    private final InternalCallContext internalCallContext;
+    private final CallContext callContext;
+    private final boolean isApiPayment;
 
     // Use to create new transactions only
     public PaymentStateContext(final boolean isApiPayment, @Nullable final UUID paymentId, @Nullable final String paymentTransactionExternalKey, final TransactionType transactionType,
@@ -95,7 +104,6 @@ public class PaymentStateContext {
         this.internalCallContext = internalCallContext;
         this.callContext = callContext;
         this.onLeavingStateExistingTransactions = ImmutableList.of();
-        this.skipOperationForUnknownTransaction = false;
     }
 
     public boolean isApiPayment() {
@@ -112,6 +120,18 @@ public class PaymentStateContext {
 
     public void setPaymentTransactionModelDao(final PaymentTransactionModelDao paymentTransactionModelDao) {
         this.paymentTransactionModelDao = paymentTransactionModelDao;
+        if (paymentId == null) {
+            this.paymentId = paymentTransactionModelDao.getPaymentId();
+        }
+        if (transactionId == null) {
+            this.transactionId = paymentTransactionModelDao.getId();
+        }
+        if (amount == null) {
+            this.amount = paymentTransactionModelDao.getAmount();
+        }
+        if (currency == null) {
+            this.currency = paymentTransactionModelDao.getCurrency();
+        }
     }
 
     public List<PaymentTransactionModelDao> getOnLeavingStateExistingTransactions() {
@@ -131,11 +151,11 @@ public class PaymentStateContext {
     }
 
     public UUID getPaymentId() {
-        return paymentId != null ? paymentId : (paymentTransactionModelDao != null ? paymentTransactionModelDao.getPaymentId() : null);
+        return paymentId;
     }
 
     public UUID getTransactionId() {
-        return transactionId != null ? transactionId : (paymentTransactionModelDao != null ? paymentTransactionModelDao.getId() : null);
+        return transactionId;
     }
 
     public String getPaymentExternalKey() {
@@ -171,13 +191,11 @@ public class PaymentStateContext {
     }
 
     public BigDecimal getAmount() {
-        // For a complete operation, if the amount isn't specified, take the original amount on the transaction
-        return amount == null && paymentTransactionModelDao != null ? paymentTransactionModelDao.getAmount() : amount;
+        return amount;
     }
 
     public Currency getCurrency() {
-        // For a complete operation, if the currency isn't specified, take the original currency on the transaction
-        return currency == null && paymentTransactionModelDao != null ? paymentTransactionModelDao.getCurrency() : currency;
+        return currency;
     }
 
     public TransactionType getTransactionType() {
@@ -204,11 +222,16 @@ public class PaymentStateContext {
         return callContext;
     }
 
-    public boolean isSkipOperationForUnknownTransaction() {
-        return skipOperationForUnknownTransaction;
+    public void setAmount(final BigDecimal adjustedAmount) {
+        this.amount = adjustedAmount;
+    }
+
+    public void setCurrency(final Currency currency) {
+        this.currency = currency;
     }
 
-    public void setSkipOperationForUnknownTransaction(final boolean skipOperationForUnknownTransaction) {
-        this.skipOperationForUnknownTransaction = skipOperationForUnknownTransaction;
+    public void setProperties(final Iterable<PluginProperty> properties) {
+        this.properties = properties;
     }
+
 }

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryAuthorizeOperationCallback.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryAuthorizeOperationCallback.java
index 1a8f812..2a3a677 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryAuthorizeOperationCallback.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/MockRetryAuthorizeOperationCallback.java
@@ -71,23 +71,23 @@ public class MockRetryAuthorizeOperationCallback extends AuthorizeControlOperati
         }
         final PaymentModelDao payment = new PaymentModelDao(clock.getUTCNow(),
                                                             clock.getUTCNow(),
-                                                            paymentStateContext.account.getId(),
-                                                            paymentStateContext.paymentMethodId,
-                                                            paymentStateContext.paymentExternalKey);
+                                                            paymentStateContext.getAccount().getId(),
+                                                            paymentStateContext.getPaymentMethodId(),
+                                                            paymentStateContext.getPaymentExternalKey());
 
         final PaymentTransactionModelDao transaction = new PaymentTransactionModelDao(clock.getUTCNow(),
                                                                                       clock.getUTCNow(),
                                                                                       paymentStateContext.getAttemptId(),
-                                                                                      paymentStateContext.paymentTransactionExternalKey,
-                                                                                      paymentStateContext.paymentId,
-                                                                                      paymentStateContext.transactionType,
+                                                                                      paymentStateContext.getPaymentTransactionExternalKey(),
+                                                                                      paymentStateContext.getPaymentId(),
+                                                                                      paymentStateContext.getTransactionType(),
                                                                                       clock.getUTCNow(),
                                                                                       TransactionStatus.SUCCESS,
-                                                                                      paymentStateContext.amount,
-                                                                                      paymentStateContext.currency,
+                                                                                      paymentStateContext.getAmount(),
+                                                                                      paymentStateContext.getCurrency(),
                                                                                       "",
                                                                                       "");
-        final PaymentModelDao paymentModelDao = paymentDao.insertPaymentWithFirstTransaction(payment, transaction, paymentStateContext.internalCallContext);
+        final PaymentModelDao paymentModelDao = paymentDao.insertPaymentWithFirstTransaction(payment, transaction, paymentStateContext.getInternalCallContext());
         final PaymentTransaction convertedTransaction = new DefaultPaymentTransaction(transaction.getId(),
                                                                                                   paymentStateContext.getAttemptId(),
                                                                                                   transaction.getTransactionExternalKey(),

diff --git a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
index 36800f1..1717dc4 100644
--- a/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
+++ b/payment/src/test/java/org/killbill/billing/payment/core/sm/TestPaymentLeavingStateCallback.java
@@ -134,8 +134,8 @@ public class TestPaymentLeavingStateCallback extends PaymentTestSuiteWithEmbedde
                                                       paymentStateContext.getPaymentMethodId(),
                                                       paymentStateContext.getAmount(),
                                                       paymentStateContext.getCurrency(),
-                                                      paymentStateContext.shouldLockAccountAndDispatch,
-                                                      paymentStateContext.overridePluginOperationResult,
+                                                      paymentStateContext.shouldLockAccountAndDispatch(),
+                                                      paymentStateContext.getOverridePluginOperationResult(),
                                                       paymentStateContext.getProperties(),
                                                       internalCallContextForOtherAccount,
                                                       callContext);