killbill-memoizeit

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
index 5bcc4c4..e599a4f 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/api/DefaultSubscriptionApi.java
@@ -18,9 +18,13 @@ package com.ning.billing.entitlement.api;
 
 import java.util.ArrayList;
 import java.util.Collections;
+import java.util.HashSet;
+import java.util.LinkedList;
 import java.util.List;
+import java.util.Set;
 import java.util.UUID;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.joda.time.DateTimeZone;
@@ -40,7 +44,9 @@ import com.ning.billing.util.svcapi.account.AccountInternalApi;
 import com.ning.billing.util.svcapi.subscription.SubscriptionBaseInternalApi;
 
 import com.google.common.base.Function;
+import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
+import com.google.common.collect.Iterables;
 import com.google.common.collect.LinkedListMultimap;
 import com.google.common.collect.ListMultimap;
 
@@ -196,16 +202,29 @@ public class DefaultSubscriptionApi implements SubscriptionApi {
         }));
 
         final Account account = accountApi.getAccountById(baseBundle.getAccountId(), internalTenantContext);
+
+        final InternalTenantContext internalTenantContextWithAccountRecordId = internalCallContextFactory.createInternalTenantContext(account.getId(), context);
+
         final DateTimeZone accountTimeZone = account.getTimeZone();
 
-        final List<BlockingState> allBlockingStates = new ArrayList<BlockingState>();
-        allBlockingStates.addAll(blockingStateDao.getBlockingAll(account.getId(), internalTenantContext));
-        allBlockingStates.addAll(blockingStateDao.getBlockingAll(bundleId, internalTenantContext));
-        for (Entitlement cur : entitlements) {
-            allBlockingStates.addAll(blockingStateDao.getBlockingAll(cur.getId(), internalTenantContext));
-        }
+        final List<BlockingState> allBlockingStatesPerAccountRecordId = blockingStateDao.getBlockingAllForAccountRecordId(internalTenantContextWithAccountRecordId);
+        final Set<UUID> allEntitlementIds = new HashSet<UUID>(Collections2.<Entitlement, UUID>transform(entitlements, new Function<Entitlement, UUID>() {
+            @Override
+            public UUID apply(final Entitlement input) {
+                return input.getId();
+            }
+        }));
+
+        final List<BlockingState> filteredBlockingStates = new LinkedList<BlockingState>(Collections2.filter(allBlockingStatesPerAccountRecordId, new Predicate<BlockingState>() {
+            @Override
+            public boolean apply(final BlockingState input) {
+                return input.getType() == BlockingStateType.ACCOUNT ||
+                       (input.getType() == BlockingStateType.SUBSCRIPTION_BUNDLE && input.getBlockedId().equals(bundleId)) ||
+                       (input.getType() == BlockingStateType.SUBSCRIPTION && allEntitlementIds.contains(input.getBlockedId()));
+            }
+        }));
 
-        final SubscriptionBundleTimeline timeline = new DefaultSubscriptionBundleTimeline(accountTimeZone, account.getId(), bundleId, baseBundle.getExternalKey(), entitlements, allBlockingStates);
+        final SubscriptionBundleTimeline timeline = new DefaultSubscriptionBundleTimeline(accountTimeZone, account.getId(), bundleId, baseBundle.getExternalKey(), entitlements, filteredBlockingStates);
         final DefaultSubscriptionBundle bundle = new DefaultSubscriptionBundle(bundleId, baseBundle.getAccountId(), baseBundle.getExternalKey(), subscriptions, timeline, baseBundle.getCreatedDate(), baseBundle.getUpdatedDate());
         return bundle;
     }

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
index 37468a7..ade0cfb 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateDao.java
@@ -74,6 +74,14 @@ public interface BlockingStateDao {
      */
     public List<BlockingState> getBlockingAll(UUID blockableId, InternalTenantContext context);
 
+
+    /**
+     * Return all events (past and future) across all services) for a given context (account_record_id)
+     * @param context
+     * @return
+     */
+    public List<BlockingState> getBlockingAllForAccountRecordId(InternalTenantContext context);
+
     /**
      * Sets a new state for a specific service.
      *

diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
index ac1b851..e4d9587 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/DefaultBlockingStateDao.java
@@ -128,6 +128,22 @@ public class DefaultBlockingStateDao implements BlockingStateDao {
     }
 
     @Override
+    public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
+        return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<List<BlockingState>>() {
+            @Override
+            public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
+                final List<BlockingStateModelDao> models = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class).getByAccountRecordId(context);
+                return new ArrayList<BlockingState>(Collections2.transform(models, new Function<BlockingStateModelDao, BlockingState>() {
+                    @Override
+                    public BlockingState apply(@Nullable final BlockingStateModelDao src) {
+                        return BlockingStateModelDao.toBlockingState(src);
+                    }
+                }));
+            }
+        });
+    }
+
+    @Override
     public void setBlockingState(final BlockingState state, final Clock clock, final InternalCallContext context) {
         transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<Void>() {
             @Override

diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java b/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
index 0d8377c..81fae38 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/dao/MockBlockingStateDao.java
@@ -104,6 +104,11 @@ public class MockBlockingStateDao implements BlockingStateDao {
     }
 
     @Override
+    public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
+        throw new UnsupportedOperationException();
+    }
+
+    @Override
     public synchronized void setBlockingState(final BlockingState state, final Clock clock, final InternalCallContext context) {
         if (blockingStates.get(state.getBlockedId()) == null) {
             blockingStates.put(state.getBlockedId(), new ArrayList<BlockingState>());