killbill-memoizeit

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/janitor/CompletionTaskBase.java b/payment/src/main/java/org/killbill/billing/payment/core/janitor/CompletionTaskBase.java
index 864832a..7df99f1 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/janitor/CompletionTaskBase.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/janitor/CompletionTaskBase.java
@@ -109,16 +109,16 @@ abstract class CompletionTaskBase<T> implements Runnable {
         public <T>  T doIteration();
     }
 
-    protected <T> T doJanitorOperationWithAccountLock(final JanitorIterationCallback callback, final Long accountRecordId, final InternalTenantContext internalTenantContext) {
+    protected <T> T doJanitorOperationWithAccountLock(final JanitorIterationCallback callback, final InternalTenantContext internalTenantContext) {
         GlobalLock lock = null;
         try {
-            final Account account = accountInternalApi.getAccountByRecordId(accountRecordId, internalTenantContext);
+            final Account account = accountInternalApi.getAccountByRecordId(internalTenantContext.getAccountRecordId(), internalTenantContext);
             lock = locker.lockWithNumberOfTries(LockerType.ACCNT_INV_PAY.toString(), account.getExternalKey(), ProcessorBase.NB_LOCK_TRY);
             return callback.doIteration();
         } catch (AccountApiException e) {
-            log.warn(String.format("Janitor failed to retrieve account with recordId %s", accountRecordId), e);
+            log.warn(String.format("Janitor failed to retrieve account with recordId %s", internalTenantContext.getAccountRecordId()), e);
         } catch (LockFailedException e) {
-            log.warn(String.format("Janitor failed to grab account with recordId %s", accountRecordId), e);
+            log.warn(String.format("Janitor failed to lock account with recordId %s", internalTenantContext.getAccountRecordId()), e);
         } finally {
             if (lock != null) {
                 lock.release();

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentAttemptTask.java b/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentAttemptTask.java
index 746f036..a02fb1a 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentAttemptTask.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentAttemptTask.java
@@ -58,6 +58,12 @@ import com.google.common.collect.Iterables;
  */
 public class IncompletePaymentAttemptTask extends CompletionTaskBase<PaymentAttemptModelDao> {
 
+    //
+    // Each paymentAttempt *should* transition to a new state, so fetching a limited size will still allow us to progress (as opposed to fetching the same entries over and over)
+    // We also don't expect to see too many entries in the INIT state.
+    //
+    private final static long MAX_ATTEMPTS_PER_ITERATIONS = 1000L;
+
     private final PluginRoutingPaymentAutomatonRunner pluginControlledPaymentAutomatonRunner;
 
     @Inject
@@ -72,7 +78,7 @@ public class IncompletePaymentAttemptTask extends CompletionTaskBase<PaymentAtte
 
     @Override
     public Iterable<PaymentAttemptModelDao> getItemsForIteration() {
-        final Pagination<PaymentAttemptModelDao> incompleteAttempts = paymentDao.getPaymentAttemptsByStateAcrossTenants(retrySMHelper.getInitialState().getName(), getCreatedDateBefore(), 0L, Long.MAX_VALUE);
+        final Pagination<PaymentAttemptModelDao> incompleteAttempts = paymentDao.getPaymentAttemptsByStateAcrossTenants(retrySMHelper.getInitialState().getName(), getCreatedDateBefore(), 0L, MAX_ATTEMPTS_PER_ITERATIONS);
         if (incompleteAttempts.getTotalNbRecords() > 0) {
             log.info("Janitor AttemptCompletionTask start run: found {} incomplete attempts", incompleteAttempts.getTotalNbRecords());
         }

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentTransactionTask.java b/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentTransactionTask.java
index 3f33684..2779e21 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentTransactionTask.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/janitor/IncompletePaymentTransactionTask.java
@@ -23,14 +23,11 @@ import java.util.List;
 import javax.inject.Inject;
 
 import org.joda.time.DateTime;
-import org.killbill.billing.ErrorCode;
-import org.killbill.billing.account.api.AccountApiException;
 import org.killbill.billing.account.api.AccountInternalApi;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
-import org.killbill.billing.payment.api.PaymentApiException;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionStatus;
 import org.killbill.billing.payment.core.PaymentTransactionInfoPluginConverter;
@@ -60,6 +57,7 @@ import com.google.common.collect.Iterables;
 
 public class IncompletePaymentTransactionTask extends CompletionTaskBase<PaymentTransactionModelDao> {
 
+
     private static final ImmutableList<TransactionStatus> TRANSACTION_STATUSES_TO_CONSIDER = ImmutableList.<TransactionStatus>builder()
                                                                                                           .add(TransactionStatus.PENDING)
                                                                                                           .add(TransactionStatus.UNKNOWN)
@@ -74,7 +72,8 @@ public class IncompletePaymentTransactionTask extends CompletionTaskBase<Payment
 
     @Override
     public Iterable<PaymentTransactionModelDao> getItemsForIteration() {
-        final Pagination<PaymentTransactionModelDao> result = paymentDao.getByTransactionStatusAcrossTenants(TRANSACTION_STATUSES_TO_CONSIDER, getCreatedDateBefore(), getCreatedDateAfter(), 0L, Long.MAX_VALUE);
+        // Code will go away in the next commit -- allow to fix h2...
+        final Pagination<PaymentTransactionModelDao> result = paymentDao.getByTransactionStatusAcrossTenants(TRANSACTION_STATUSES_TO_CONSIDER, getCreatedDateBefore(), getCreatedDateAfter(), 0L, 1000L);
         if (result.getTotalNbRecords() > 0) {
             log.info("Janitor IncompletePaymentTransactionTask start run: found {} pending/unknown payments", result.getTotalNbRecords());
         }
@@ -88,19 +87,23 @@ public class IncompletePaymentTransactionTask extends CompletionTaskBase<Payment
         doJanitorOperationWithAccountLock(new JanitorIterationCallback() {
             @Override
             public Void doIteration() {
+
+                // State may have changed since we originally retrieved with no lock
+                final PaymentTransactionModelDao rehydratedPaymentTransaction = paymentDao.getPaymentTransaction(paymentTransaction.getId(), internalTenantContext);
+
                 final TenantContext tenantContext = internalCallContextFactory.createTenantContext(internalTenantContext);
-                final PaymentModelDao payment = paymentDao.getPayment(paymentTransaction.getPaymentId(), internalTenantContext);
+                final PaymentModelDao payment = paymentDao.getPayment(rehydratedPaymentTransaction.getPaymentId(), internalTenantContext);
 
                 final PaymentMethodModelDao paymentMethod = paymentDao.getPaymentMethod(payment.getPaymentMethodId(), internalTenantContext);
                 final PaymentPluginApi paymentPluginApi = getPaymentPluginApi(payment, paymentMethod.getPluginName());
 
                 final PaymentTransactionInfoPlugin undefinedPaymentTransaction = new DefaultNoOpPaymentInfoPlugin(payment.getId(),
-                                                                                                                  paymentTransaction.getId(),
-                                                                                                                  paymentTransaction.getTransactionType(),
-                                                                                                                  paymentTransaction.getAmount(),
-                                                                                                                  paymentTransaction.getCurrency(),
-                                                                                                                  paymentTransaction.getCreatedDate(),
-                                                                                                                  paymentTransaction.getCreatedDate(),
+                                                                                                                  rehydratedPaymentTransaction.getId(),
+                                                                                                                  rehydratedPaymentTransaction.getTransactionType(),
+                                                                                                                  rehydratedPaymentTransaction.getAmount(),
+                                                                                                                  rehydratedPaymentTransaction.getCurrency(),
+                                                                                                                  rehydratedPaymentTransaction.getCreatedDate(),
+                                                                                                                  rehydratedPaymentTransaction.getCreatedDate(),
                                                                                                                   PaymentPluginStatus.UNDEFINED,
                                                                                                                   null);
                 PaymentTransactionInfoPlugin paymentTransactionInfoPlugin;
@@ -109,7 +112,7 @@ public class IncompletePaymentTransactionTask extends CompletionTaskBase<Payment
                     paymentTransactionInfoPlugin = Iterables.tryFind(result, new Predicate<PaymentTransactionInfoPlugin>() {
                         @Override
                         public boolean apply(final PaymentTransactionInfoPlugin input) {
-                            return input.getKbTransactionPaymentId().equals(paymentTransaction.getId());
+                            return input.getKbTransactionPaymentId().equals(rehydratedPaymentTransaction.getId());
                         }
                     }).or(new Supplier<PaymentTransactionInfoPlugin>() {
                         @Override
@@ -120,30 +123,38 @@ public class IncompletePaymentTransactionTask extends CompletionTaskBase<Payment
                 } catch (final Exception e) {
                     paymentTransactionInfoPlugin = undefinedPaymentTransaction;
                 }
-                updatePaymentAndTransactionIfNeeded(payment, paymentTransaction, paymentTransactionInfoPlugin, internalTenantContext);
+                updatePaymentAndTransactionIfNeeded(payment, rehydratedPaymentTransaction, paymentTransactionInfoPlugin, internalTenantContext);
                 return null;
             }
-        }, paymentTransaction.getAccountRecordId(), internalTenantContext);
+        }, internalTenantContext);
 
     }
 
-    public Boolean updatePaymentAndTransactionIfNeededWithAccountLock(final PaymentModelDao payment, final PaymentTransactionModelDao paymentTransaction, final PaymentTransactionInfoPlugin paymentTransactionInfoPlugin, final InternalTenantContext internalTenantContext) {
-        return doJanitorOperationWithAccountLock(new JanitorIterationCallback() {
+    public boolean updatePaymentAndTransactionIfNeededWithAccountLock(final PaymentModelDao payment, final PaymentTransactionModelDao paymentTransaction, final PaymentTransactionInfoPlugin paymentTransactionInfoPlugin, final InternalTenantContext internalTenantContext) {
+        // Can happen in the GET case, see PaymentProcessor#toPayment
+        if (!TRANSACTION_STATUSES_TO_CONSIDER.contains(paymentTransaction.getTransactionStatus())) {
+            // Nothing to do
+            return false;
+        }
+        final Boolean result =  doJanitorOperationWithAccountLock(new JanitorIterationCallback() {
             @Override
             public Boolean doIteration() {
-                return updatePaymentAndTransactionIfNeeded(payment, paymentTransaction, paymentTransactionInfoPlugin, internalTenantContext);
+                return updatePaymentAndTransactionInternal(payment, paymentTransaction, paymentTransactionInfoPlugin, internalTenantContext);
             }
-        }, internalTenantContext.getAccountRecordId(), internalTenantContext);
+        }, internalTenantContext);
+        return result != null && result;
     }
 
     private boolean updatePaymentAndTransactionIfNeeded(final PaymentModelDao payment, final PaymentTransactionModelDao paymentTransaction, final PaymentTransactionInfoPlugin paymentTransactionInfoPlugin, final InternalTenantContext internalTenantContext) {
-        final CallContext callContext = createCallContext("IncompletePaymentTransactionTask", internalTenantContext);
-
-        // Can happen in the GET case, see PaymentProcessor#toPayment
         if (!TRANSACTION_STATUSES_TO_CONSIDER.contains(paymentTransaction.getTransactionStatus())) {
             // Nothing to do
             return false;
         }
+        return updatePaymentAndTransactionInternal(payment, paymentTransaction, paymentTransactionInfoPlugin, internalTenantContext);
+    }
+
+    private boolean updatePaymentAndTransactionInternal(final PaymentModelDao payment, final PaymentTransactionModelDao paymentTransaction, final PaymentTransactionInfoPlugin paymentTransactionInfoPlugin, final InternalTenantContext internalTenantContext) {
+        final CallContext callContext = createCallContext("IncompletePaymentTransactionTask", internalTenantContext);
 
         // First obtain the new transactionStatus,
         // Then compute the new paymentState; this one is mostly interesting in case of success (to compute the lastSuccessPaymentState below)
@@ -189,8 +200,10 @@ public class IncompletePaymentTransactionTask extends CompletionTaskBase<Payment
                                                            paymentTransaction.getId(), transactionStatus, processedAmount, processedCurrency, gatewayErrorCode, gatewayError, internalCallContext);
 
         return true;
+
     }
 
+
     // Keep the existing currentTransactionStatus if we can't obtain a better answer from the plugin; if not, return the newTransactionStatus
     private TransactionStatus computeNewTransactionStatusFromPaymentTransactionInfoPlugin(final PaymentTransactionInfoPlugin input, final TransactionStatus currentTransactionStatus) {
         final TransactionStatus newTransactionStatus = PaymentTransactionInfoPluginConverter.toTransactionStatus(input);

diff --git a/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java b/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
index 69b2651..0af54d5 100644
--- a/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
+++ b/payment/src/main/java/org/killbill/billing/payment/core/PaymentProcessor.java
@@ -413,11 +413,11 @@ public class PaymentProcessor extends ProcessorBase {
             PaymentTransactionModelDao newPaymentTransactionModelDao = curPaymentTransactionModelDao;
 
             final PaymentTransactionInfoPlugin paymentTransactionInfoPlugin = findPaymentTransactionInfoPlugin(newPaymentTransactionModelDao, pluginTransactions);
-            if (paymentTransactionInfoPlugin != null && curPaymentTransactionModelDao.getTransactionStatus() == TransactionStatus.UNKNOWN) {
+            if (paymentTransactionInfoPlugin != null) {
                 // Make sure to invoke the Janitor task in case the plugin fixes its state on the fly
                 // See https://github.com/killbill/killbill/issues/341
-                final Boolean hasChanged = incompletePaymentTransactionTask.updatePaymentAndTransactionIfNeededWithAccountLock(newPaymentModelDao, newPaymentTransactionModelDao, paymentTransactionInfoPlugin, internalTenantContext);
-                if (hasChanged != null && hasChanged) {
+                final boolean hasChanged = incompletePaymentTransactionTask.updatePaymentAndTransactionIfNeededWithAccountLock(newPaymentModelDao, newPaymentTransactionModelDao, paymentTransactionInfoPlugin, internalTenantContext);
+                if (hasChanged) {
                     newPaymentModelDao = paymentDao.getPayment(newPaymentModelDao.getId(), internalTenantContext);
                     newPaymentTransactionModelDao = paymentDao.getPaymentTransaction(newPaymentTransactionModelDao.getId(), internalTenantContext);
                 }

diff --git a/payment/src/test/java/org/killbill/billing/payment/dao/TestPaymentDao.java b/payment/src/test/java/org/killbill/billing/payment/dao/TestPaymentDao.java
index e156f5b..9afbb45 100644
--- a/payment/src/test/java/org/killbill/billing/payment/dao/TestPaymentDao.java
+++ b/payment/src/test/java/org/killbill/billing/payment/dao/TestPaymentDao.java
@@ -496,7 +496,7 @@ public class TestPaymentDao extends PaymentTestSuiteWithEmbeddedDB {
             paymentDao.insertPaymentWithFirstTransaction(paymentModelDao1, transaction1, context1);
         }
 
-        final Pagination<PaymentTransactionModelDao> result =  paymentDao.getByTransactionStatusAcrossTenants(ImmutableList.of(TransactionStatus.UNKNOWN), clock.getUTCNow(), createdDate1, 0L, Long.MAX_VALUE);
+        final Pagination<PaymentTransactionModelDao> result =  paymentDao.getByTransactionStatusAcrossTenants(ImmutableList.of(TransactionStatus.UNKNOWN), clock.getUTCNow(), createdDate1, 0L, new Long(NB_ENTRIES));
         Assert.assertEquals(result.getTotalNbRecords(), new Long(NB_ENTRIES));
 
         final Iterator<PaymentTransactionModelDao> iterator = result.iterator();