killbill-uncached

entitlement: further BlockingStateDao optimizations Save

12/3/2013 11:47:41 AM

Details

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 aaa286b..3dd9cdd 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
@@ -51,27 +51,6 @@ public interface BlockingStateDao extends EntityDao<BlockingStateModelDao, Block
     public List<BlockingState> getBlockingState(UUID blockableId, BlockingStateType blockingStateType, InternalTenantContext context);
 
     /**
-     * Returns the state history for that specific service
-     *
-     * @param blockableId       id of the blockable object
-     * @param blockingStateType blockable object type
-     * @param serviceName       name of the service
-     * @param context           call context
-     * @return list of blocking states for that blockable object and service
-     */
-    public List<BlockingState> getBlockingHistoryForService(UUID blockableId, BlockingStateType blockingStateType, String serviceName, InternalTenantContext context);
-
-    /**
-     * Return all the events (past and future) across all services
-     *
-     * @param blockableId       id of the blockable object
-     * @param blockingStateType blockable object type
-     * @param context           call context
-     * @return list of blocking states for that blockable object
-     */
-    public List<BlockingState> getBlockingAll(UUID blockableId, BlockingStateType blockingStateType, InternalTenantContext context);
-
-    /**
      * Return all events (past and future) across all services) for a given callcontext (account_record_id)
      *
      * @param context call context
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateSqlDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateSqlDao.java
index e95ebe9..c69758b 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateSqlDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/BlockingStateSqlDao.java
@@ -57,23 +57,16 @@ public interface BlockingStateSqlDao extends EntitySqlDao<BlockingStateModelDao,
                                                                  @Bind("effectiveDate") Date effectiveDate,
                                                                  @BindBean final InternalTenantContext context);
 
-
     @SqlQuery
     public abstract List<BlockingStateModelDao> getBlockingHistoryForService(@Bind("blockableId") UUID blockableId,
                                                                              @Bind("service") String serviceName,
                                                                              @BindBean final InternalTenantContext context);
 
-    @SqlQuery
-    public abstract List<BlockingStateModelDao> getBlockingAll(@Bind("blockableId") UUID blockableId,
-                                                               @BindBean final InternalTenantContext context);
-
-
     @SqlUpdate
     @Audited(ChangeType.UPDATE)
     public void unactiveEvent(@Bind("id") String id,
                               @BindBean final InternalCallContext context);
 
-
     public class BlockingHistorySqlMapper extends MapperBase implements ResultSetMapper<BlockingStateModelDao> {
 
         @Override
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 1292cbf..7f21843 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
@@ -111,44 +111,6 @@ public class DefaultBlockingStateDao extends EntityDaoBase<BlockingStateModelDao
     }
 
     @Override
-    public List<BlockingState> getBlockingHistoryForService(final UUID blockableId, final BlockingStateType blockingStateType, final String serviceName, final InternalTenantContext context) {
-        return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<List<BlockingState>>() {
-            @Override
-            public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
-                final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final List<BlockingStateModelDao> models = sqlDao.getBlockingHistoryForService(blockableId, serviceName, context);
-                final Collection<BlockingStateModelDao> modelsFiltered = filterBlockingStates(models, blockingStateType);
-                return new ArrayList<BlockingState>(Collections2.transform(modelsFiltered,
-                                                                           new Function<BlockingStateModelDao, BlockingState>() {
-                                                                               @Override
-                                                                               public BlockingState apply(@Nullable final BlockingStateModelDao src) {
-                                                                                   return BlockingStateModelDao.toBlockingState(src);
-                                                                               }
-                                                                           }));
-            }
-        });
-    }
-
-    @Override
-    public List<BlockingState> getBlockingAll(final UUID blockableId, final BlockingStateType blockingStateType, final InternalTenantContext context) {
-        return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<List<BlockingState>>() {
-            @Override
-            public List<BlockingState> inTransaction(final EntitySqlDaoWrapperFactory<EntitySqlDao> entitySqlDaoWrapperFactory) throws Exception {
-                final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final List<BlockingStateModelDao> models = sqlDao.getBlockingAll(blockableId, context);
-                final Collection<BlockingStateModelDao> modelsFiltered = filterBlockingStates(models, blockingStateType);
-                return new ArrayList<BlockingState>(Collections2.transform(modelsFiltered,
-                                                                           new Function<BlockingStateModelDao, BlockingState>() {
-                                                                               @Override
-                                                                               public BlockingState apply(@Nullable final BlockingStateModelDao src) {
-                                                                                   return BlockingStateModelDao.toBlockingState(src);
-                                                                               }
-                                                                           }));
-            }
-        });
-    }
-
-    @Override
     public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
         return transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<List<BlockingState>>() {
             @Override
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
index d810635..4f71fba 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
@@ -23,6 +23,7 @@ import javax.annotation.Nullable;
 
 import org.skife.jdbi.v2.IDBI;
 
+import com.ning.billing.account.api.Account;
 import com.ning.billing.callcontext.InternalTenantContext;
 import com.ning.billing.catalog.api.ProductCategory;
 import com.ning.billing.clock.Clock;
@@ -47,8 +48,31 @@ public class OptimizedProxyBlockingStateDao extends ProxyBlockingStateDao {
         super(eventsStreamBuilder, subscriptionBaseInternalApi, dbi, clock, cacheControllerDispatcher, nonEntityDao);
     }
 
-    // Special signature for EventsStreamBuilder to save some DAO calls
-    public List<BlockingState> getBlockingHistory(final List<BlockingState> blockingStatesOnDisk,
+    /**
+     * Retrieve blocking states for a given subscription
+     * <p/>
+     * If the specified subscription is not an add-on, we already have the blocking states
+     * (they are all on disk) - we simply return them and there is nothing to do.
+     * Otherwise, for add-ons, we will need to compute the blocking states not on disk.
+     * <p/>
+     * This is a special method for EventsStreamBuilder to save some DAO calls.
+     *
+     * @param subscriptionBlockingStatesOnDisk
+     *                                  blocking states on disk for that subscription
+     * @param allBlockingStatesOnDiskForAccount
+     *                                  all blocking states on disk for that account
+     * @param account                   account associated with the subscription
+     * @param bundle                    bundle associated with the subscription
+     * @param baseSubscription          base subscription (ProductCategory.BASE) associated with that bundle
+     * @param subscription              subscription for which to build blocking states
+     * @param allSubscriptionsForBundle all subscriptions associated with that bundle
+     * @param context                   call context
+     * @return blocking states for that subscription
+     * @throws EntitlementApiException
+     */
+    public List<BlockingState> getBlockingHistory(final List<BlockingState> subscriptionBlockingStatesOnDisk,
+                                                  final List<BlockingState> allBlockingStatesOnDiskForAccount,
+                                                  final Account account,
                                                   final SubscriptionBaseBundle bundle,
                                                   @Nullable final SubscriptionBase baseSubscription,
                                                   final SubscriptionBase subscription,
@@ -57,18 +81,20 @@ public class OptimizedProxyBlockingStateDao extends ProxyBlockingStateDao {
         // blockable id points to a subscription, but make sure it's an add-on
         if (!ProductCategory.ADD_ON.equals(subscription.getCategory())) {
             // blockable id points to a base or standalone subscription, there is nothing to do
-            return blockingStatesOnDisk;
+            return subscriptionBlockingStatesOnDisk;
         }
 
         // Find all base entitlements that we care about (for which we want to find future cancelled add-ons)
-        final Iterable<EventsStream> eventsStreams = ImmutableList.<EventsStream>of(eventsStreamBuilder.buildForEntitlement(bundle,
+        final Iterable<EventsStream> eventsStreams = ImmutableList.<EventsStream>of(eventsStreamBuilder.buildForEntitlement(allBlockingStatesOnDiskForAccount,
+                                                                                                                            account,
+                                                                                                                            bundle,
                                                                                                                             baseSubscription,
                                                                                                                             allSubscriptionsForBundle,
                                                                                                                             context));
 
         return addBlockingStatesNotOnDisk(subscription.getId(),
                                           BlockingStateType.SUBSCRIPTION,
-                                          new LinkedList<BlockingState>(blockingStatesOnDisk),
+                                          new LinkedList<BlockingState>(subscriptionBlockingStatesOnDisk),
                                           ImmutableList.<SubscriptionBase>of(baseSubscription),
                                           eventsStreams);
     }
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
index 7662f22..87ffd36 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/dao/ProxyBlockingStateDao.java
@@ -46,13 +46,11 @@ import com.ning.billing.entitlement.api.EntitlementApiException;
 import com.ning.billing.entitlement.engine.core.EventsStreamBuilder;
 import com.ning.billing.subscription.api.SubscriptionBase;
 import com.ning.billing.subscription.api.SubscriptionBaseInternalApi;
-import com.ning.billing.subscription.api.user.SubscriptionBaseApiException;
 import com.ning.billing.util.cache.CacheControllerDispatcher;
 import com.ning.billing.util.dao.NonEntityDao;
 import com.ning.billing.util.entity.Pagination;
 
 import com.google.common.base.Predicate;
-import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
 
@@ -164,21 +162,9 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
     }
 
     @Override
-    public List<BlockingState> getBlockingHistoryForService(final UUID blockableId, final BlockingStateType blockingStateType, final String serviceName, final InternalTenantContext context) {
-        final List<BlockingState> statesOnDisk = delegate.getBlockingHistoryForService(blockableId, blockingStateType, serviceName, context);
-        return addBlockingStatesNotOnDisk(blockableId, blockingStateType, statesOnDisk, context);
-    }
-
-    @Override
-    public List<BlockingState> getBlockingAll(final UUID blockableId, final BlockingStateType blockingStateType, final InternalTenantContext context) {
-        final List<BlockingState> statesOnDisk = delegate.getBlockingAll(blockableId, blockingStateType, context);
-        return addBlockingStatesNotOnDisk(blockableId, blockingStateType, statesOnDisk, context);
-    }
-
-    @Override
     public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
         final List<BlockingState> statesOnDisk = delegate.getBlockingAllForAccountRecordId(context);
-        return addBlockingStatesNotOnDisk(null, null, statesOnDisk, context);
+        return addBlockingStatesNotOnDisk(statesOnDisk, context);
     }
 
     @Override
@@ -193,9 +179,7 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
 
     // Add blocking states for add-ons, which would be impacted by a future cancellation or change of their base plan
     // See DefaultEntitlement#blockAddOnsIfRequired
-    private List<BlockingState> addBlockingStatesNotOnDisk(@Nullable final UUID blockableId,
-                                                           @Nullable final BlockingStateType blockingStateType,
-                                                           final List<BlockingState> blockingStatesOnDisk,
+    private List<BlockingState> addBlockingStatesNotOnDisk(final List<BlockingState> blockingStatesOnDisk,
                                                            final InternalTenantContext context) {
         final Collection<BlockingState> blockingStatesOnDiskCopy = new LinkedList<BlockingState>(blockingStatesOnDisk);
 
@@ -203,56 +187,29 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
         final Iterable<SubscriptionBase> baseSubscriptionsToConsider;
         final Iterable<EventsStream> eventsStreams;
         try {
-            if (blockingStateType == null) {
-                // We're coming from getBlockingAllForAccountRecordId
-                final Map<UUID, List<SubscriptionBase>> subscriptions = subscriptionInternalApi.getSubscriptionsForAccount(context);
-                baseSubscriptionsToConsider = Iterables.<SubscriptionBase>filter(Iterables.<SubscriptionBase>concat(subscriptions.values()),
-                                                                                 new Predicate<SubscriptionBase>() {
-                                                                                     @Override
-                                                                                     public boolean apply(final SubscriptionBase input) {
-                                                                                         return ProductCategory.BASE.equals(input.getCategory()) &&
-                                                                                                !EntitlementState.CANCELLED.equals(input.getState());
-                                                                                     }
-                                                                                 });
-                eventsStreams = Iterables.<EventsStream>concat(eventsStreamBuilder.buildForAccount(subscriptions, context).getEventsStreams().values());
-            } else if (BlockingStateType.SUBSCRIPTION.equals(blockingStateType)) {
-                // We're coming from getBlockingHistoryForService / getBlockingAll
-                final SubscriptionBase subscription = subscriptionInternalApi.getSubscriptionFromId(blockableId, context);
-
-                // blockable id points to a subscription, but make sure it's an add-on
-                if (ProductCategory.ADD_ON.equals(subscription.getCategory())) {
-                    final SubscriptionBase baseSubscription = subscriptionInternalApi.getBaseSubscription(subscription.getBundleId(), context);
-                    baseSubscriptionsToConsider = ImmutableList.<SubscriptionBase>of(baseSubscription);
-                    eventsStreams = ImmutableList.<EventsStream>of(eventsStreamBuilder.buildForEntitlement(baseSubscription, context));
-                } else {
-                    // blockable id points to a base or standalone subscription, there is nothing to do
-                    // Simply return the sorted list
-                    return BLOCKING_STATE_ORDERING.immutableSortedCopy(blockingStatesOnDisk);
-                }
-            } else {
-                // blockable id points to an account or bundle, in which case there are no extra blocking states to add
-                // Simply return the sorted list
-                return BLOCKING_STATE_ORDERING.immutableSortedCopy(blockingStatesOnDisk);
-            }
-        } catch (SubscriptionBaseApiException e) {
-            log.error("Error retrieving subscriptions for account record id " + context.getAccountRecordId(), e);
-            throw new RuntimeException(e);
+            final Map<UUID, List<SubscriptionBase>> subscriptions = subscriptionInternalApi.getSubscriptionsForAccount(context);
+            baseSubscriptionsToConsider = Iterables.<SubscriptionBase>filter(Iterables.<SubscriptionBase>concat(subscriptions.values()),
+                                                                             new Predicate<SubscriptionBase>() {
+                                                                                 @Override
+                                                                                 public boolean apply(final SubscriptionBase input) {
+                                                                                     return ProductCategory.BASE.equals(input.getCategory());
+                                                                                 }
+                                                                             });
+            eventsStreams = Iterables.<EventsStream>concat(eventsStreamBuilder.buildForAccount(subscriptions, context).getEventsStreams().values());
         } catch (EntitlementApiException e) {
             log.error("Error computing blocking states for addons for account record id " + context.getAccountRecordId(), e);
             throw new RuntimeException(e);
         }
 
-        return addBlockingStatesNotOnDisk(blockableId, blockingStateType, blockingStatesOnDiskCopy, baseSubscriptionsToConsider, eventsStreams);
+        return addBlockingStatesNotOnDisk(null, null, blockingStatesOnDiskCopy, baseSubscriptionsToConsider, eventsStreams);
     }
 
+    // Special signature for OptimizedProxyBlockingStateDao
     protected List<BlockingState> addBlockingStatesNotOnDisk(@Nullable final UUID blockableId,
                                                              @Nullable final BlockingStateType blockingStateType,
                                                              final Collection<BlockingState> blockingStatesOnDiskCopy,
                                                              final Iterable<SubscriptionBase> baseSubscriptionsToConsider,
                                                              final Iterable<EventsStream> eventsStreams) {
-        // Retrieve the cancellation blocking state on disk, if it exists (will be used later)
-        final BlockingState cancellationBlockingStateOnDisk = findEntitlementCancellationBlockingState(blockableId, blockingStatesOnDiskCopy);
-
         // Compute the blocking states not on disk for all base subscriptions
         final DateTime now = clock.getUTCNow();
         for (final SubscriptionBase baseSubscription : baseSubscriptionsToConsider) {
@@ -264,20 +221,23 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
                                                                                }
                                                                            });
 
-            // First, check to see if the base entitlement is cancelled. If so, cancel the
+            // First, check to see if the base entitlement is cancelled
             final Collection<BlockingState> blockingStatesNotOnDisk = eventsStream.computeAddonsBlockingStatesForFutureSubscriptionBaseEvents();
 
             // Inject the extra blocking states into the stream if needed
             for (final BlockingState blockingState : blockingStatesNotOnDisk) {
                 // If this entitlement is actually already cancelled, add the cancellation event we computed
                 // only if it's prior to the blocking state on disk (e.g. add-on future cancelled but base plan cancelled earlier).
-                final boolean overrideCancellationBlockingStateOnDisk = cancellationBlockingStateOnDisk != null &&
-                                                                        isEntitlementCancellationBlockingState(blockingState) &&
-                                                                        blockingState.getEffectiveDate().isBefore(cancellationBlockingStateOnDisk.getEffectiveDate());
+                BlockingState cancellationBlockingStateOnDisk = null;
+                boolean overrideCancellationBlockingStateOnDisk = false;
+                if (isEntitlementCancellationBlockingState(blockingState)) {
+                    cancellationBlockingStateOnDisk = findEntitlementCancellationBlockingState(blockingState.getBlockedId(), blockingStatesOnDiskCopy);
+                    overrideCancellationBlockingStateOnDisk = cancellationBlockingStateOnDisk != null && blockingState.getEffectiveDate().isBefore(cancellationBlockingStateOnDisk.getEffectiveDate());
+                }
 
                 if ((
                             blockingStateType == null ||
-                            // In case we're coming from getBlockingHistoryForService / getBlockingAll, make sure we don't add
+                            // In case we're coming from OptimizedProxyBlockingStateDao, make sure we don't add
                             // blocking states for other add-ons on that base subscription
                             (BlockingStateType.SUBSCRIPTION.equals(blockingStateType) && blockingState.getBlockedId().equals(blockableId))
                     ) && (
diff --git a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
index be8a899..a051f1a 100644
--- a/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
+++ b/entitlement/src/main/java/com/ning/billing/entitlement/engine/core/EventsStreamBuilder.java
@@ -126,10 +126,6 @@ public class EventsStreamBuilder {
             throw new EntitlementApiException(e);
         }
 
-        return buildForAccount(account, subscriptions, internalTenantContext);
-    }
-
-    private AccountEventsStreams buildForAccount(final Account account, final Map<UUID, List<SubscriptionBase>> subscriptions, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         if (subscriptions.isEmpty()) {
             // Bail early
             return new DefaultAccountEventsStreams(account);
@@ -197,6 +193,8 @@ public class EventsStreamBuilder {
                     subscriptionBlockingStates = subscriptionBlockingStatesOnDisk;
                 } else {
                     subscriptionBlockingStates = blockingStateDao.getBlockingHistory(subscriptionBlockingStatesOnDisk,
+                                                                                     blockingStatesForAccount,
+                                                                                     account,
                                                                                      bundle,
                                                                                      baseSubscription,
                                                                                      subscription,
@@ -239,45 +237,36 @@ public class EventsStreamBuilder {
             throw new EntitlementApiException(e);
         }
 
-        return buildForEntitlement(bundle, baseSubscription, subscription, allSubscriptionsForBundle, internalTenantContext);
-    }
-
-    // Special signature for ProxyBlockingStateDao to save some DAO calls
-    public EventsStream buildForEntitlement(final SubscriptionBase subscription, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
-        final SubscriptionBaseBundle bundle;
+        final Account account;
         try {
-            bundle = subscriptionInternalApi.getBundleFromId(subscription.getBundleId(), internalTenantContext);
-        } catch (SubscriptionBaseApiException e) {
+            account = accountInternalApi.getAccountById(bundle.getAccountId(), internalTenantContext);
+        } catch (AccountApiException e) {
             throw new EntitlementApiException(e);
         }
 
-        final List<SubscriptionBase> allSubscriptionsForBundle = subscriptionInternalApi.getSubscriptionsForBundle(subscription.getBundleId(), internalTenantContext);
-        return buildForEntitlement(bundle, subscription, subscription, allSubscriptionsForBundle, internalTenantContext);
+        // Retrieve the blocking states
+        final List<BlockingState> blockingStatesForAccount = defaultBlockingStateDao.getBlockingAllForAccountRecordId(internalTenantContext);
+
+        return buildForEntitlement(blockingStatesForAccount, account, bundle, baseSubscription, subscription, allSubscriptionsForBundle, internalTenantContext);
     }
 
     // Special signature for OptimizedProxyBlockingStateDao to save some DAO calls
-    public EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle,
-                                            final SubscriptionBase subscription,
+    public EventsStream buildForEntitlement(final List<BlockingState> blockingStatesForAccount,
+                                            final Account account,
+                                            final SubscriptionBaseBundle bundle,
+                                            final SubscriptionBase baseSubscription,
                                             final List<SubscriptionBase> allSubscriptionsForBundle,
                                             final InternalTenantContext internalTenantContext) throws EntitlementApiException {
-        return buildForEntitlement(bundle, subscription, subscription, allSubscriptionsForBundle, internalTenantContext);
+        return buildForEntitlement(blockingStatesForAccount, account, bundle, baseSubscription, baseSubscription, allSubscriptionsForBundle, internalTenantContext);
     }
 
-    private EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle,
+    private EventsStream buildForEntitlement(final List<BlockingState> blockingStatesForAccount,
+                                             final Account account,
+                                             final SubscriptionBaseBundle bundle,
                                              @Nullable final SubscriptionBase baseSubscription,
                                              final SubscriptionBase subscription,
                                              final List<SubscriptionBase> allSubscriptionsForBundle,
                                              final InternalTenantContext internalTenantContext) throws EntitlementApiException {
-        final Account account;
-        try {
-            account = accountInternalApi.getAccountById(bundle.getAccountId(), internalTenantContext);
-        } catch (AccountApiException e) {
-            throw new EntitlementApiException(e);
-        }
-
-        // Retrieve the blocking states
-        final Collection<BlockingState> blockingStatesForAccount = defaultBlockingStateDao.getBlockingAllForAccountRecordId(internalTenantContext);
-
         // Optimization: build lookup tables for blocking states states
         final Collection<BlockingState> accountBlockingStates = new LinkedList<BlockingState>();
         final Map<UUID, List<BlockingState>> blockingStatesPerSubscription = new HashMap<UUID, List<BlockingState>>();
@@ -308,9 +297,13 @@ public class EventsStreamBuilder {
         // needed, i.e. if this EventStream is for a standalone or a base subscription
         final Collection<BlockingState> subscriptionBlockingStates;
         if (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) {
+            // Note: we come here during the recursion from OptimizedProxyBlockingStateDao#getBlockingHistory
+            // (called by blockingStateDao.getBlockingHistory below)
             subscriptionBlockingStates = subscriptionBlockingStatesOnDisk;
         } else {
             subscriptionBlockingStates = blockingStateDao.getBlockingHistory(ImmutableList.<BlockingState>copyOf(subscriptionBlockingStatesOnDisk),
+                                                                             blockingStatesForAccount,
+                                                                             account,
                                                                              bundle,
                                                                              baseSubscription,
                                                                              subscription,
diff --git a/entitlement/src/main/resources/com/ning/billing/entitlement/dao/BlockingStateSqlDao.sql.stg b/entitlement/src/main/resources/com/ning/billing/entitlement/dao/BlockingStateSqlDao.sql.stg
index 973434a..62d8de4 100644
--- a/entitlement/src/main/resources/com/ning/billing/entitlement/dao/BlockingStateSqlDao.sql.stg
+++ b/entitlement/src/main/resources/com/ning/billing/entitlement/dao/BlockingStateSqlDao.sql.stg
@@ -92,18 +92,6 @@ and is_active
 ;
 >>
 
-getBlockingAll() ::= <<
-select
-<allTableFields()>
-from
-<tableName()>
-where blockable_id = :blockableId
-and is_active
-<AND_CHECK_TENANT()>
-<defaultOrderBy()>
-;
->>
-
 unactiveEvent() ::= <<
 update
 <tableName()>
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 3e3c69d..5f1faf6 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
@@ -42,11 +42,11 @@ public class MockBlockingStateDao extends MockEntityDaoBase<BlockingStateModelDa
     private final Map<UUID, List<BlockingState>> blockingStates = new HashMap<UUID, List<BlockingState>>();
     private final Map<Long, List<BlockingState>> blockingStatesPerAccountRecordId = new HashMap<Long, List<BlockingState>>();
 
-    // TODO This mock class should also check that events are past or present except for getBlockingAll
+    // TODO This mock class should also check that events are past or present
 
     @Override
     public BlockingState getBlockingStateForService(final UUID blockableId, final BlockingStateType blockingStateType, final String serviceName, final InternalTenantContext context) {
-        final List<BlockingState> states = getBlockingAll(blockableId, blockingStateType, context);
+        final List<BlockingState> states = blockingStates.get(blockableId);
         if (states == null) {
             return null;
         }
@@ -77,30 +77,6 @@ public class MockBlockingStateDao extends MockEntityDaoBase<BlockingStateModelDa
     }
 
     @Override
-    public List<BlockingState> getBlockingHistoryForService(final UUID overdueableId, final BlockingStateType blockingStateType, final String serviceName, final InternalTenantContext context) {
-        final List<BlockingState> states = blockingStates.get(overdueableId);
-        if (states == null) {
-            return new ArrayList<BlockingState>();
-        }
-        final ImmutableList<BlockingState> filtered = ImmutableList.<BlockingState>copyOf(Collections2.filter(states, new Predicate<BlockingState>() {
-            @Override
-            public boolean apply(@Nullable final BlockingState input) {
-                return input.getService().equals(serviceName);
-            }
-        }));
-
-        // Note! The returned list cannot be immutable!
-        return states == null ? new ArrayList<BlockingState>() : new ArrayList<BlockingState>(filtered);
-    }
-
-    @Override
-    public List<BlockingState> getBlockingAll(final UUID blockableId, final BlockingStateType blockingStateType, final InternalTenantContext context) {
-        final List<BlockingState> states = blockingStates.get(blockableId);
-        // Note! The returned list cannot be immutable!
-        return states == null ? new ArrayList<BlockingState>() : states;
-    }
-
-    @Override
     public List<BlockingState> getBlockingAllForAccountRecordId(final InternalTenantContext context) {
         return Objects.firstNonNull(blockingStatesPerAccountRecordId.get(context.getAccountRecordId()), ImmutableList.<BlockingState>of());
     }
diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
index db37ccd..264e810 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestBlockingDao.java
@@ -21,8 +21,10 @@ import java.util.UUID;
 
 import org.joda.time.LocalDate;
 import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
+import com.ning.billing.account.api.Account;
 import com.ning.billing.entitlement.EntitlementTestSuiteWithEmbeddedDB;
 import com.ning.billing.entitlement.api.BlockingState;
 import com.ning.billing.entitlement.api.BlockingStateType;
@@ -30,9 +32,16 @@ import com.ning.billing.junction.DefaultBlockingState;
 
 public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
 
-    @Test(groups = "slow")
-    public void testDao() {
+    @BeforeMethod(groups = "slow")
+    public void setUp() throws Exception {
+        final Account account = accountApi.createAccount(getAccountData(7), callContext);
 
+        // Override the context with the right account record id
+        internalCallContext = internalCallContextFactory.createInternalCallContext(account.getId(), callContext);
+    }
+
+    @Test(groups = "slow", description = "Check BlockingStateDao with a single service")
+    public void testDaoWithOneService() {
         final UUID uuid = UUID.randomUUID();
         final String overdueStateName = "WayPassedItMan";
         final String service = "TEST";
@@ -54,15 +63,15 @@ public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
 
         Assert.assertEquals(blockingStateDao.getBlockingStateForService(uuid, BlockingStateType.ACCOUNT, service, internalCallContext).getStateName(), state2.getStateName());
 
-        final List<BlockingState> states = blockingStateDao.getBlockingHistoryForService(uuid, BlockingStateType.ACCOUNT, service, internalCallContext);
+        final List<BlockingState> states = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(states.size(), 2);
 
         Assert.assertEquals(states.get(0).getStateName(), overdueStateName);
         Assert.assertEquals(states.get(1).getStateName(), overdueStateName2);
     }
 
-    @Test(groups = "slow")
-    public void testDaoHistory() throws Exception {
+    @Test(groups = "slow", description = "Check BlockingStateDao with multiple services")
+    public void testDaoWithMultipleServices() throws Exception {
         final UUID uuid = UUID.randomUUID();
         final String overdueStateName = "WayPassedItMan";
         final String service1 = "TEST";
@@ -81,7 +90,7 @@ public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
         final BlockingState state2 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName2, service2, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
         blockingStateDao.setBlockingState(state2, clock, internalCallContext);
 
-        final List<BlockingState> history2 = blockingStateDao.getBlockingAll(uuid, BlockingStateType.ACCOUNT, internalCallContext);
+        final List<BlockingState> history2 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(history2.size(), 2);
         Assert.assertEquals(history2.get(0).getStateName(), overdueStateName);
         Assert.assertEquals(history2.get(1).getStateName(), overdueStateName2);
diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestDefaultBlockingStateDao.java b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestDefaultBlockingStateDao.java
index 9300fc7..06989c9 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestDefaultBlockingStateDao.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/dao/TestDefaultBlockingStateDao.java
@@ -22,6 +22,7 @@ import java.util.UUID;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
 import com.ning.billing.account.api.Account;
@@ -38,12 +39,21 @@ import com.ning.billing.junction.DefaultBlockingState;
 
 public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbeddedDB {
 
+    private Account account;
+
+    @BeforeMethod(groups = "slow")
+    public void setUp() throws Exception {
+        account = accountApi.createAccount(getAccountData(7), callContext);
+
+        // Override the context with the right account record id
+        internalCallContext = internalCallContextFactory.createInternalCallContext(account.getId(), callContext);
+    }
+
     @Test(groups = "slow", description = "Verify we don't insert extra add-on events")
     public void testUnnecessaryEventsAreNotAdded() throws Exception {
         // This is a simple smoke test at the dao level only to make sure we do sane
         // things in case there are no future add-on cancellation events to add in the stream.
         // See TestEntitlementUtils for a more comprehensive test
-        final Account account = accountApi.createAccount(getAccountData(7), callContext);
         final PlanPhaseSpecifier spec = new PlanPhaseSpecifier("Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, PriceListSet.DEFAULT_PRICELIST_NAME, null);
         testListener.pushExpectedEvent(NextEvent.CREATE);
         final Entitlement entitlement = entitlementApi.createBaseEntitlement(account.getId(), spec, account.getExternalKey(), clock.getUTCToday(), callContext);
@@ -54,14 +64,14 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final String service = "service";
 
         // Verify initial state
-        Assert.assertEquals(blockingStateDao.getBlockingAll(entitlement.getId(), type, internalCallContext).size(), 0);
+        Assert.assertEquals(blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext).size(), 0);
 
         // Set a state
         final DateTime stateDateTime = new DateTime(2013, 5, 6, 10, 11, 12, DateTimeZone.UTC);
         final BlockingState blockingState = new DefaultBlockingState(entitlement.getId(), type, state, service, false, false, false, stateDateTime);
         blockingStateDao.setBlockingState(blockingState, clock, internalCallContext);
 
-        Assert.assertEquals(blockingStateDao.getBlockingAll(entitlement.getId(), type, internalCallContext).size(), 1);
+        Assert.assertEquals(blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext).size(), 1);
     }
 
     // See https://github.com/killbill/killbill/issues/111
@@ -75,7 +85,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final String serviceB = "service-B";
 
         // Verify initial state
-        Assert.assertEquals(blockingStateDao.getBlockingAll(blockableId, type, internalCallContext).size(), 0);
+        Assert.assertEquals(blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext).size(), 0);
 
         // Note: the checkers below rely on record_id ordering, not effective date
 
@@ -83,7 +93,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final DateTime stateDateTime = new DateTime(2013, 5, 6, 10, 11, 12, DateTimeZone.UTC);
         final BlockingState blockingState1 = new DefaultBlockingState(blockableId, type, state, serviceA, false, false, false, stateDateTime);
         blockingStateDao.setBlockingState(blockingState1, clock, internalCallContext);
-        final List<BlockingState> blockingStates1 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates1 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates1.size(), 1);
         Assert.assertEquals(blockingStates1.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates1.get(0).getStateName(), state);
@@ -92,7 +102,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
 
         // Set the same state again - no change
         blockingStateDao.setBlockingState(blockingState1, clock, internalCallContext);
-        final List<BlockingState> blockingStates2 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates2 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates2.size(), 1);
         Assert.assertEquals(blockingStates2.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates2.get(0).getStateName(), state);
@@ -102,7 +112,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // Set the state for service B
         final BlockingState blockingState2 = new DefaultBlockingState(blockableId, type, state, serviceB, false, false, false, stateDateTime);
         blockingStateDao.setBlockingState(blockingState2, clock, internalCallContext);
-        final List<BlockingState> blockingStates3 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates3 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates3.size(), 2);
         Assert.assertEquals(blockingStates3.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates3.get(0).getStateName(), state);
@@ -117,7 +127,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final DateTime stateDateTime2 = new DateTime(2013, 6, 6, 10, 11, 12, DateTimeZone.UTC);
         final BlockingState blockingState3 = new DefaultBlockingState(blockableId, type, state, serviceA, false, false, false, stateDateTime2);
         blockingStateDao.setBlockingState(blockingState3, clock, internalCallContext);
-        final List<BlockingState> blockingStates4 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates4 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates4.size(), 2);
         Assert.assertEquals(blockingStates4.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates4.get(0).getStateName(), state);
@@ -132,7 +142,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final DateTime stateDateTime3 = new DateTime(2013, 2, 6, 10, 11, 12, DateTimeZone.UTC);
         final BlockingState blockingState4 = new DefaultBlockingState(blockableId, type, state, serviceA, false, false, false, stateDateTime3);
         blockingStateDao.setBlockingState(blockingState4, clock, internalCallContext);
-        final List<BlockingState> blockingStates5 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates5 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates5.size(), 2);
         Assert.assertEquals(blockingStates5.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates5.get(0).getStateName(), state);
@@ -147,7 +157,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         final DateTime state2DateTime = new DateTime(2013, 12, 6, 10, 11, 12, DateTimeZone.UTC);
         final BlockingState blockingState5 = new DefaultBlockingState(blockableId, type, state2, serviceA, false, false, false, state2DateTime);
         blockingStateDao.setBlockingState(blockingState5, clock, internalCallContext);
-        final List<BlockingState> blockingStates6 = blockingStateDao.getBlockingAll(blockableId, type, internalCallContext);
+        final List<BlockingState> blockingStates6 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates6.size(), 3);
         Assert.assertEquals(blockingStates6.get(0).getBlockedId(), blockableId);
         Assert.assertEquals(blockingStates6.get(0).getStateName(), state);
diff --git a/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java b/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
index b623448..034a147 100644
--- a/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
+++ b/entitlement/src/test/java/com/ning/billing/entitlement/engine/core/TestEntitlementUtils.java
@@ -18,6 +18,7 @@ package com.ning.billing.entitlement.engine.core;
 
 import java.util.Collection;
 import java.util.List;
+import java.util.UUID;
 
 import javax.annotation.Nullable;
 
@@ -49,6 +50,7 @@ import com.ning.billing.entitlement.dao.BlockingStateSqlDao;
 
 import com.google.common.base.Objects;
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 
 public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
@@ -279,7 +281,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
         // Verify the blocking states DAO adds events not on disk for the first add-on...
         checkBlockingStatesDAO(changedBaseEntitlement, addOnEntitlement, baseEffectiveCancellationOrChangeDate, false);
         // ...but not for the second one
-        final List<BlockingState> blockingStatesForSecondAddOn = blockingStateDao.getBlockingAll(secondAddOnEntitlement.getId(), BlockingStateType.SUBSCRIPTION, internalCallContext);
+        final List<BlockingState> blockingStatesForSecondAddOn = blockingStatesForBlockedId(secondAddOnEntitlement.getId());
         Assert.assertEquals(blockingStatesForSecondAddOn.size(), 0);
     }
 
@@ -417,7 +419,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
 
     // Test the DAO
     private void checkBlockingStatesDAO(final DefaultEntitlement baseEntitlement, final DefaultEntitlement addOnEntitlement, final LocalDate effectiveBaseCancellationDate, final LocalDate effectiveAddOnCancellationDate, final boolean isBaseCancelled) {
-        final List<BlockingState> blockingStatesForBaseEntitlement = blockingStateDao.getBlockingAll(baseEntitlement.getId(), BlockingStateType.SUBSCRIPTION, internalCallContext);
+        final List<BlockingState> blockingStatesForBaseEntitlement = blockingStatesForBlockedId(baseEntitlement.getId());
         Assert.assertEquals(blockingStatesForBaseEntitlement.size(), isBaseCancelled ? 1 : 0);
         if (isBaseCancelled) {
             Assert.assertEquals(blockingStatesForBaseEntitlement.get(0).getBlockedId(), baseEntitlement.getId());
@@ -427,7 +429,7 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
             Assert.assertEquals(blockingStatesForBaseEntitlement.get(0).getStateName(), DefaultEntitlementApi.ENT_STATE_CANCELLED);
         }
 
-        final List<BlockingState> blockingStatesForAddOn = blockingStateDao.getBlockingAll(addOnEntitlement.getId(), BlockingStateType.SUBSCRIPTION, internalCallContext);
+        final List<BlockingState> blockingStatesForAddOn = blockingStatesForBlockedId(addOnEntitlement.getId());
         Assert.assertEquals(blockingStatesForAddOn.size(), 1);
         Assert.assertEquals(blockingStatesForAddOn.get(0).getBlockedId(), addOnEntitlement.getId());
         Assert.assertEquals(blockingStatesForAddOn.get(0).getEffectiveDate().toLocalDate(), effectiveAddOnCancellationDate);
@@ -469,4 +471,14 @@ public class TestEntitlementUtils extends EntitlementTestSuiteWithEmbeddedDB {
                                                                        });
         return eventsStream.computeAddonsBlockingStatesForNextSubscriptionBaseEvent(effectiveDate);
     }
+
+    private List<BlockingState> blockingStatesForBlockedId(final UUID blockedId) {
+        return ImmutableList.<BlockingState>copyOf(Iterables.<BlockingState>filter(blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext),
+                                                                                   new Predicate<BlockingState>() {
+                                                                                       @Override
+                                                                                       public boolean apply(final BlockingState input) {
+                                                                                           return input.getBlockedId().equals(blockedId);
+                                                                                       }
+                                                                                   }));
+    }
 }