killbill-memoizeit

entitlement: set blocking states in bulk This makes sure

1/12/2016 2:02:56 PM

Details

diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
index dffcddb..29400c3 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlement.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
  * License.  You may obtain a copy of the License at:
  *
@@ -17,7 +19,9 @@
 package org.killbill.billing.entitlement.api;
 
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.UUID;
 
@@ -296,14 +300,19 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
             } else {
                 getSubscriptionBase().cancel(callContext);
             }
-        } catch (SubscriptionBaseApiException e) {
+        } catch (final SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
         final BlockingState newBlockingState = new DefaultBlockingState(getId(), BlockingStateType.SUBSCRIPTION, DefaultEntitlementApi.ENT_STATE_CANCELLED, EntitlementService.ENTITLEMENT_SERVICE_NAME, true, true, false, effectiveCancelDate);
-        entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(newBlockingState, contextWithValidAccountRecordId);
+        final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+        final Collection<BlockingState> addOnsBlockingStates = computeAddOnBlockingStates(effectiveCancelDate, notificationEvents, callContext, contextWithValidAccountRecordId);
 
-        blockAddOnsIfRequired(effectiveCancelDate, callContext, contextWithValidAccountRecordId);
+        // Record the new state first, then insert the notifications to avoid race conditions
+        setBlockingStates(newBlockingState, addOnsBlockingStates, contextWithValidAccountRecordId);
+        for (final NotificationEvent notificationEvent : notificationEvents) {
+            recordFutureNotification(effectiveCancelDate, notificationEvent, contextWithValidAccountRecordId);
+        }
 
         return entitlementApi.getEntitlementForId(getId(), callContext);
     }
@@ -353,7 +362,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         if (getSubscriptionBase().getFutureEndDate() != null) {
             try {
                 getSubscriptionBase().uncancel(callContext);
-            } catch (SubscriptionBaseApiException e) {
+            } catch (final SubscriptionBaseApiException e) {
                 throw new EntitlementApiException(e);
             }
         }
@@ -394,14 +403,19 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
                 try {
                     // Cancel subscription base first, to correctly compute the add-ons entitlements we need to cancel (see below)
                     getSubscriptionBase().cancelWithPolicy(billingPolicy, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
 
                 final BlockingState newBlockingState = new DefaultBlockingState(getId(), BlockingStateType.SUBSCRIPTION, DefaultEntitlementApi.ENT_STATE_CANCELLED, EntitlementService.ENTITLEMENT_SERVICE_NAME, true, true, false, effectiveDate);
-                entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(newBlockingState, contextWithValidAccountRecordId);
+                final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+                final Collection<BlockingState> addOnsBlockingStates = computeAddOnBlockingStates(effectiveDate, notificationEvents, callContext, contextWithValidAccountRecordId);
 
-                blockAddOnsIfRequired(effectiveDate, callContext, contextWithValidAccountRecordId);
+                // Record the new state first, then insert the notifications to avoid race conditions
+                setBlockingStates(newBlockingState, addOnsBlockingStates, contextWithValidAccountRecordId);
+                for (final NotificationEvent notificationEvent : notificationEvents) {
+                    recordFutureNotification(effectiveDate, notificationEvent, contextWithValidAccountRecordId);
+                }
 
                 return entitlementApi.getEntitlementForId(getId(), callContext);
             }
@@ -454,19 +468,25 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
                 final DateTime effectiveChangeDate;
                 try {
                     effectiveChangeDate = getSubscriptionBase().changePlan(productName, billingPeriod, priceList, overrides, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
 
                 final InternalCallContext context = internalCallContextFactory.createInternalCallContext(getAccountId(), callContext);
                 try {
                     checker.checkBlockedChange(getSubscriptionBase(), effectiveChangeDate, context);
-                } catch (BlockingApiException e) {
+                } catch (final BlockingApiException e) {
                     throw new EntitlementApiException(e, e.getCode(), e.getMessage());
                 }
 
-                blockAddOnsIfRequired(effectiveChangeDate, callContext, context);
+                final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+                final Iterable<BlockingState> addOnsBlockingStates = computeAddOnBlockingStates(effectiveChangeDate, notificationEvents, callContext, context);
 
+                // Record the new state first, then insert the notifications to avoid race conditions
+                setBlockingStates(addOnsBlockingStates, context);
+                for (final NotificationEvent notificationEvent : notificationEvents) {
+                    recordFutureNotification(effectiveChangeDate, notificationEvent, context);
+                }
                 return entitlementApi.getEntitlementForId(getId(), callContext);
             }
         };
@@ -502,18 +522,24 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
                 final DateTime effectiveChangeDate = dateHelper.fromLocalDateAndReferenceTime(updatedPluginContext.getEffectiveDate(), getSubscriptionBase().getStartDate(), context);
                 try {
                     getSubscriptionBase().changePlanWithDate(productName, billingPeriod, priceList, overrides, effectiveChangeDate, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
 
                 try {
                     checker.checkBlockedChange(getSubscriptionBase(), effectiveChangeDate, context);
-                } catch (BlockingApiException e) {
+                } catch (final BlockingApiException e) {
                     throw new EntitlementApiException(e, e.getCode(), e.getMessage());
                 }
 
+                final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+                final Iterable<BlockingState> addOnsBlockingStates = computeAddOnBlockingStates(effectiveChangeDate, notificationEvents, callContext, context);
 
-                blockAddOnsIfRequired(effectiveChangeDate, callContext, context);
+                // Record the new state first, then insert the notifications to avoid race conditions
+                setBlockingStates(addOnsBlockingStates, context);
+                for (final NotificationEvent notificationEvent : notificationEvents) {
+                    recordFutureNotification(effectiveChangeDate, notificationEvent, context);
+                }
 
                 return entitlementApi.getEntitlementForId(getId(), callContext);
             }
@@ -551,17 +577,24 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
                 final DateTime effectiveChangeDate;
                 try {
                     effectiveChangeDate = getSubscriptionBase().changePlanWithPolicy(productName, billingPeriod, priceList, overrides, actionPolicy, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
 
                 try {
                     checker.checkBlockedChange(getSubscriptionBase(), effectiveChangeDate, context);
-                } catch (BlockingApiException e) {
+                } catch (final BlockingApiException e) {
                     throw new EntitlementApiException(e, e.getCode(), e.getMessage());
                 }
 
-                blockAddOnsIfRequired(effectiveChangeDate, callContext, context);
+                final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+                final Iterable<BlockingState> addOnsBlockingStates = computeAddOnBlockingStates(effectiveChangeDate, notificationEvents, callContext, context);
+
+                // Record the new state first, then insert the notifications to avoid race conditions
+                setBlockingStates(addOnsBlockingStates, context);
+                for (final NotificationEvent notificationEvent : notificationEvents) {
+                    recordFutureNotification(effectiveChangeDate, notificationEvent, context);
+                }
 
                 return entitlementApi.getEntitlementForId(getId(), callContext);
             }
@@ -573,11 +606,11 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
         eventsStream = eventsStreamBuilder.refresh(eventsStream, context);
     }
 
-    public void blockAddOnsIfRequired(final DateTime effectiveDate, final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
+    public Collection<BlockingState> computeAddOnBlockingStates(final DateTime effectiveDate, final Collection<NotificationEvent> notificationEvents, final TenantContext context, final InternalCallContext internalCallContext) throws EntitlementApiException {
         // Optimization - bail early
         if (!ProductCategory.BASE.equals(getSubscriptionBase().getCategory())) {
             // Only base subscriptions have add-ons
-            return;
+            return ImmutableList.<BlockingState>of();
         }
 
         // Get the latest state from disk (we just got cancelled or changed plan)
@@ -594,14 +627,11 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
             // go through the DAO (e.g. change)
             final boolean isBaseEntitlementCancelled = eventsStream.isEntitlementCancelled();
             final NotificationEvent notificationEvent = new EntitlementNotificationKey(getId(), getBundleId(), isBaseEntitlementCancelled ? EntitlementNotificationKeyAction.CANCEL : EntitlementNotificationKeyAction.CHANGE, effectiveDate);
-            recordFutureNotification(effectiveDate, notificationEvent, internalCallContext);
-            return;
+            notificationEvents.add(notificationEvent);
+            return ImmutableList.<BlockingState>of();
         }
 
-        final Collection<BlockingState> addOnsBlockingStates = eventsStream.computeAddonsBlockingStatesForNextSubscriptionBaseEvent(effectiveDate);
-        for (final BlockingState addOnBlockingState : addOnsBlockingStates) {
-            entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(addOnBlockingState, internalCallContext);
-        }
+        return eventsStream.computeAddonsBlockingStatesForNextSubscriptionBaseEvent(effectiveDate);
     }
 
     private void recordFutureNotification(final DateTime effectiveDate,
@@ -611,18 +641,29 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
             final NotificationQueue subscriptionEventQueue = notificationQueueService.getNotificationQueue(DefaultEntitlementService.ENTITLEMENT_SERVICE_NAME,
                                                                                                            DefaultEntitlementService.NOTIFICATION_QUEUE_NAME);
             subscriptionEventQueue.recordFutureNotification(effectiveDate, notificationEvent, context.getUserToken(), context.getAccountRecordId(), context.getTenantRecordId());
-        } catch (NoSuchNotificationQueue e) {
+        } catch (final NoSuchNotificationQueue e) {
             throw new RuntimeException(e);
-        } catch (IOException e) {
+        } catch (final IOException e) {
             throw new RuntimeException(e);
         }
     }
 
+    private void setBlockingStates(final BlockingState entitlementBlockingState, final Collection<BlockingState> addOnsBlockingStates, final InternalCallContext internalCallContext) {
+        final Collection<BlockingState> states = new LinkedList<BlockingState>();
+        states.add(entitlementBlockingState);
+        states.addAll(addOnsBlockingStates);
+        setBlockingStates(states, internalCallContext);
+    }
+
+    private void setBlockingStates(final Iterable<BlockingState> blockingStates, final InternalCallContext internalCallContext) {
+        entitlementUtils.setBlockingStatesAndPostBlockingTransitionEvent(blockingStates, getBundleId(), internalCallContext);
+    }
+
     //
     // Unfortunately the permission checks for the entitlement api cannot *simply* rely on the KillBillShiroAopModule because some of the operations (CANCEL, CHANGE) are
     // done through objects that are not injected by Guice, and so the check needs to happen explicitly.
     //
-    private void checkForPermissions(final Permission permission, final CallContext callContext) throws EntitlementApiException {
+    private void checkForPermissions(final Permission permission, final TenantContext callContext) throws EntitlementApiException {
         //
         // If authentication had been done (CorsBasicHttpAuthenticationFilter) we verify the correct permissions exist.
         //
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
index e8983c2..0560b0e 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/DefaultEntitlementApi.java
@@ -1,7 +1,9 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
  * License.  You may obtain a copy of the License at:
  *
@@ -17,7 +19,9 @@
 package org.killbill.billing.entitlement.api;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import javax.inject.Inject;
@@ -59,8 +63,6 @@ import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.bus.api.PersistentBus;
 import org.killbill.clock.Clock;
 import org.killbill.notificationq.api.NotificationQueueService;
-import org.slf4j.Logger;
-import org.slf4j.LoggerFactory;
 
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;
@@ -70,15 +72,12 @@ import com.google.common.collect.Lists;
 
 public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements EntitlementApi {
 
-    private static final Logger log = LoggerFactory.getLogger(DefaultEntitlementApi.class);
-
     public static final String ENT_STATE_BLOCKED = "ENT_BLOCKED";
     public static final String ENT_STATE_CLEAR = "ENT_CLEAR";
     public static final String ENT_STATE_CANCELLED = "ENT_CANCELLED";
 
     private final SubscriptionBaseInternalApi subscriptionBaseInternalApi;
     private final SubscriptionBaseTransferApi subscriptionBaseTransferApi;
-    private final AccountInternalApi accountApi;
     private final Clock clock;
     private final InternalCallContextFactory internalCallContextFactory;
     private final BlockingChecker checker;
@@ -104,7 +103,6 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
         this.internalCallContextFactory = internalCallContextFactory;
         this.subscriptionBaseInternalApi = subscriptionInternalApi;
         this.subscriptionBaseTransferApi = subscriptionTransferApi;
-        this.accountApi = accountApi;
         this.clock = clock;
         this.checker = checker;
         this.blockingStateDao = blockingStateDao;
@@ -119,7 +117,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
     @Override
     public Entitlement createBaseEntitlement(final UUID accountId, final PlanPhaseSpecifier planPhaseSpecifier, final String externalKey, final List<PlanPhasePriceOverride> overrides, final LocalDate effectiveDate, final Iterable<PluginProperty> properties, final CallContext callContext) throws EntitlementApiException {
 
-        EntitlementSpecifier entitlementSpecifier = new DefaultEntitlementSpecifier(planPhaseSpecifier, overrides);
+        final EntitlementSpecifier entitlementSpecifier = new DefaultEntitlementSpecifier(planPhaseSpecifier, overrides);
         final List<EntitlementSpecifier> entitlementSpecifierList = new ArrayList<EntitlementSpecifier>();
         entitlementSpecifierList.add(entitlementSpecifier);
         final EntitlementContext pluginContext = new DefaultEntitlementContext(OperationType.CREATE_SUBSCRIPTION,
@@ -152,7 +150,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                     return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
             }
@@ -216,8 +214,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
 
-
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
 
@@ -229,7 +226,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
     @Override
     public Entitlement addEntitlement(final UUID bundleId, final PlanPhaseSpecifier planPhaseSpecifier, final List<PlanPhasePriceOverride> overrides, final LocalDate effectiveDate, final Iterable<PluginProperty> properties, final CallContext callContext) throws EntitlementApiException {
 
-        EntitlementSpecifier entitlementSpecifier = new DefaultEntitlementSpecifier(planPhaseSpecifier, overrides);
+        final EntitlementSpecifier entitlementSpecifier = new DefaultEntitlementSpecifier(planPhaseSpecifier, overrides);
         final List<EntitlementSpecifier> entitlementSpecifierList = new ArrayList<EntitlementSpecifier>();
         entitlementSpecifierList.add(entitlementSpecifier);
         final EntitlementContext pluginContext = new DefaultEntitlementContext(OperationType.CREATE_SUBSCRIPTION,
@@ -267,7 +264,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                     return new DefaultEntitlement(subscription.getId(), eventsStreamBuilder, entitlementApi, pluginExecution,
                                                   blockingStateDao, subscriptionBaseInternalApi, checker, notificationQueueService,
                                                   entitlementUtils, dateHelper, clock, securityApi, internalCallContextFactory, callContext);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
             }
@@ -285,7 +282,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
             final InternalTenantContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalTenantContext(bundle.getAccountId(), context);
             final DateTime requestedDate = dateHelper.fromLocalDateAndReferenceTime(effectiveDate, baseSubscription.getStartDate(), contextWithValidAccountRecordId);
             return subscriptionBaseInternalApi.getDryRunChangePlanStatus(baseSubscription.getId(), targetProductName, requestedDate, contextWithValidAccountRecordId);
-        } catch (SubscriptionBaseApiException e) {
+        } catch (final SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
     }
@@ -302,7 +299,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
         final UUID accountId;
         try {
             accountId = subscriptionBaseInternalApi.getBundleFromId(bundleId, internalContext).getAccountId();
-        } catch (SubscriptionBaseApiException e) {
+        } catch (final SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
 
@@ -360,15 +357,13 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
 
     }
 
-
     @Override
-    public void setBlockingState(final UUID bundleId, final String stateName, final String serviceName, final LocalDate effectiveDate, boolean blockBilling, boolean blockEntitlement, boolean blockChange, final Iterable<PluginProperty> properties, final CallContext context)
+    public void setBlockingState(final UUID bundleId, final String stateName, final String serviceName, final LocalDate effectiveDate, final boolean blockBilling, final boolean blockEntitlement, final boolean blockChange, final Iterable<PluginProperty> properties, final CallContext context)
             throws EntitlementApiException {
         final InternalCallContext contextWithValidAccountRecordId = internalCallContextFactory.createInternalCallContext(bundleId, ObjectType.BUNDLE, context);
         super.setBlockingState(bundleId, stateName, serviceName, effectiveDate, blockBilling, blockEntitlement, blockChange, properties, contextWithValidAccountRecordId);
     }
 
-
     @Override
     public Iterable<BlockingState> getBlockingStatesForServiceAndType(final UUID blockableId, final BlockingStateType blockingStateType, final String serviceName, final TenantContext tenantContext) {
         // Not implemented see #431
@@ -426,15 +421,17 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
 
                     // Block all associated subscriptions - TODO Do we want to block the bundle as well (this will add an extra STOP_ENTITLEMENT event in the bundle timeline stream)?
                     // Note that there is no un-transfer at the moment, so we effectively add a blocking state on disk for all subscriptions
+                    final Map<BlockingState, UUID> blockingStates = new HashMap<BlockingState, UUID>();
                     for (final SubscriptionBase subscriptionBase : subscriptionBaseInternalApi.getSubscriptionsForBundle(baseBundle.getId(), null, contextWithValidAccountRecordId)) {
                         final BlockingState blockingState = new DefaultBlockingState(subscriptionBase.getId(), BlockingStateType.SUBSCRIPTION, DefaultEntitlementApi.ENT_STATE_CANCELLED, EntitlementService.ENTITLEMENT_SERVICE_NAME, true, true, false, requestedDate);
-                        entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(blockingState, contextWithValidAccountRecordId);
+                        blockingStates.put(blockingState, subscriptionBase.getBundleId());
                     }
+                    entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(blockingStates, contextWithValidAccountRecordId);
 
                     return newBundle.getId();
-                } catch (SubscriptionBaseTransferApiException e) {
+                } catch (final SubscriptionBaseTransferApiException e) {
                     throw new EntitlementApiException(e);
-                } catch (SubscriptionBaseApiException e) {
+                } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
             }
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultEntitlementApiBase.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultEntitlementApiBase.java
index 9651c95..0173d7c 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultEntitlementApiBase.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultEntitlementApiBase.java
@@ -78,6 +78,8 @@ import org.killbill.notificationq.api.NotificationQueueService.NoSuchNotificatio
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.collect.ImmutableList;
+
 public class DefaultEntitlementApiBase {
 
     private static final Logger log = LoggerFactory.getLogger(DefaultEntitlementApiBase.class);
@@ -277,9 +279,9 @@ public class DefaultEntitlementApiBase {
             final SubscriptionBase baseSubscription = inputBaseSubscription == null ? subscriptionInternalApi.getBaseSubscription(bundleId, internalCallContext) : inputBaseSubscription;
             final DateTime effectiveDate = dateHelper.fromLocalDateAndReferenceTime(localEffectiveDate, baseSubscription.getStartDate(), internalCallContext);
             final BlockingState state = new DefaultBlockingState(bundleId, BlockingStateType.SUBSCRIPTION_BUNDLE, stateName, serviceName, blockChange, blockEntitlement, blockBilling, effectiveDate);
-            entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(state, internalCallContext);
+            entitlementUtils.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableList.<BlockingState>of(state), bundleId, internalCallContext);
             return state.getId();
-        } catch (SubscriptionBaseApiException e) {
+        } catch (final SubscriptionBaseApiException e) {
             throw new EntitlementApiException(e);
         }
     }
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/BlockingStateDao.java b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/BlockingStateDao.java
index 770f56d..3327f70 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/BlockingStateDao.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/BlockingStateDao.java
@@ -19,6 +19,7 @@
 package org.killbill.billing.entitlement.dao;
 
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
@@ -29,6 +30,8 @@ import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.entitlement.api.EntitlementApiException;
 import org.killbill.billing.util.entity.dao.EntityDao;
 
+import com.google.common.base.Optional;
+
 public interface BlockingStateDao extends EntityDao<BlockingStateModelDao, BlockingState, EntitlementApiException> {
 
     /**
@@ -61,13 +64,12 @@ public interface BlockingStateDao extends EntityDao<BlockingStateModelDao, Block
     public List<BlockingState> getBlockingAllForAccountRecordId(InternalTenantContext context);
 
     /**
-     * Sets a new state for a specific service and send an event if needed
+     * Set new blocking states
      *
-     * @param state   blocking state to set
-     * @param bundleId bundle id of the associated bundle if the blocking state type is SUBSCRIPTION
+     * @param states  blocking states to set (mapped to the associated bundle id if the blocking state type is SUBSCRIPTION)
      * @param context call context
      */
-    public void setBlockingStateAndPostBlockingTransitionEvent(BlockingState state, UUID bundleId, InternalCallContext context);
+    public void setBlockingStatesAndPostBlockingTransitionEvent(Map<BlockingState, Optional<UUID>> states, InternalCallContext context);
 
     /**
      * Unactive the blocking state
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
index 6db86ed..cd65fa1 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/DefaultBlockingStateDao.java
@@ -25,6 +25,7 @@ import java.util.Comparator;
 import java.util.Date;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
@@ -65,6 +66,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import com.google.common.base.Function;
+import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
@@ -179,64 +181,68 @@ public class DefaultBlockingStateDao extends EntityDaoBase<BlockingStateModelDao
     }
 
     @Override
-    public void setBlockingStateAndPostBlockingTransitionEvent(final BlockingState state, final UUID bundleId, final InternalCallContext context) {
+    public void setBlockingStatesAndPostBlockingTransitionEvent(final Map<BlockingState, Optional<UUID>> states, final InternalCallContext context) {
         transactionalSqlDao.execute(new EntitySqlDaoTransactionWrapper<Void>() {
             @Override
             public Void inTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory) throws Exception {
                 final DateTime upToDate = clock.getUTCNow();
                 final BlockingStateSqlDao sqlDao = entitySqlDaoWrapperFactory.become(BlockingStateSqlDao.class);
-                final BlockingAggregator previousState = getBlockedStatus(sqlDao, entitySqlDaoWrapperFactory.getHandle(), state.getBlockedId(), state.getType(), bundleId, upToDate, context);
-
-                final BlockingStateModelDao newBlockingStateModelDao = new BlockingStateModelDao(state, context);
-
-                // Get all blocking states for that blocked id and service
-                final List<BlockingStateModelDao> allForBlockedItAndService = sqlDao.getBlockingHistoryForService(state.getBlockedId(), state.getService(), context);
-
-                // Add the new one (we rely below on the fact that the ID for newBlockingStateModelDao is now set)
-                allForBlockedItAndService.add(newBlockingStateModelDao);
-
-                // Re-order what should be the final list (allForBlockedItAndService is ordered by record_id in the SQL and we just added a new state)
-                final List<BlockingStateModelDao> allForBlockedItAndServiceOrdered = BLOCKING_STATE_MODEL_DAO_ORDERING.immutableSortedCopy(allForBlockedItAndService);
-
-                // Go through the (ordered) stream of blocking states for that blocked id and service and check
-                // if there is one or more blocking states for the same state following each others.
-                // If there are, delete them, as they are not needed anymore. A picture being worth a thousand words,
-                // if the current stream is: t0 S1 t1 S2 t3 S3 and we want to insert S2 at t0 < t1' < t1,
-                // the final stream should be: t0 S1 t1' S2 t3 S3 (and not t0 S1 t1' S2 t1 S2 t3 S3)
-                // Note that we also take care of the use case t0 S1 t1 S2 t2 S2 t3 S3 to cleanup legacy systems, although
-                // it shouldn't happen anymore
-                final Collection<UUID> blockingStatesToRemove = new HashSet<UUID>();
-                BlockingStateModelDao prevBlockingStateModelDao = null;
-                for (final BlockingStateModelDao blockingStateModelDao : allForBlockedItAndServiceOrdered) {
-                    if (prevBlockingStateModelDao != null && prevBlockingStateModelDao.getState().equals(blockingStateModelDao.getState())) {
-                        blockingStatesToRemove.add(blockingStateModelDao.getId());
+
+                for (final BlockingState state : states.keySet()) {
+                    final UUID bundleId = states.get(state).orNull();
+                    final BlockingAggregator previousState = getBlockedStatus(sqlDao, entitySqlDaoWrapperFactory.getHandle(), state.getBlockedId(), state.getType(), bundleId, upToDate, context);
+
+                    final BlockingStateModelDao newBlockingStateModelDao = new BlockingStateModelDao(state, context);
+
+                    // Get all blocking states for that blocked id and service
+                    final List<BlockingStateModelDao> allForBlockedItAndService = sqlDao.getBlockingHistoryForService(state.getBlockedId(), state.getService(), context);
+
+                    // Add the new one (we rely below on the fact that the ID for newBlockingStateModelDao is now set)
+                    allForBlockedItAndService.add(newBlockingStateModelDao);
+
+                    // Re-order what should be the final list (allForBlockedItAndService is ordered by record_id in the SQL and we just added a new state)
+                    final List<BlockingStateModelDao> allForBlockedItAndServiceOrdered = BLOCKING_STATE_MODEL_DAO_ORDERING.immutableSortedCopy(allForBlockedItAndService);
+
+                    // Go through the (ordered) stream of blocking states for that blocked id and service and check
+                    // if there is one or more blocking states for the same state following each others.
+                    // If there are, delete them, as they are not needed anymore. A picture being worth a thousand words,
+                    // if the current stream is: t0 S1 t1 S2 t3 S3 and we want to insert S2 at t0 < t1' < t1,
+                    // the final stream should be: t0 S1 t1' S2 t3 S3 (and not t0 S1 t1' S2 t1 S2 t3 S3)
+                    // Note that we also take care of the use case t0 S1 t1 S2 t2 S2 t3 S3 to cleanup legacy systems, although
+                    // it shouldn't happen anymore
+                    final Collection<UUID> blockingStatesToRemove = new HashSet<UUID>();
+                    BlockingStateModelDao prevBlockingStateModelDao = null;
+                    for (final BlockingStateModelDao blockingStateModelDao : allForBlockedItAndServiceOrdered) {
+                        if (prevBlockingStateModelDao != null && prevBlockingStateModelDao.getState().equals(blockingStateModelDao.getState())) {
+                            blockingStatesToRemove.add(blockingStateModelDao.getId());
+                        }
+                        prevBlockingStateModelDao = blockingStateModelDao;
                     }
-                    prevBlockingStateModelDao = blockingStateModelDao;
-                }
 
-                // Delete unnecessary states (except newBlockingStateModelDao, which doesn't exist in the database)
-                for (final UUID blockedId : blockingStatesToRemove) {
-                    if (!newBlockingStateModelDao.getId().equals(blockedId)) {
-                        sqlDao.unactiveEvent(blockedId.toString(), context);
+                    // Delete unnecessary states (except newBlockingStateModelDao, which doesn't exist in the database)
+                    for (final UUID blockedId : blockingStatesToRemove) {
+                        if (!newBlockingStateModelDao.getId().equals(blockedId)) {
+                            sqlDao.unactiveEvent(blockedId.toString(), context);
+                        }
                     }
-                }
 
-                // Create the state, if needed
-                if (!blockingStatesToRemove.contains(newBlockingStateModelDao.getId())) {
-                    sqlDao.create(newBlockingStateModelDao, context);
-                }
+                    // Create the state, if needed
+                    if (!blockingStatesToRemove.contains(newBlockingStateModelDao.getId())) {
+                        sqlDao.create(newBlockingStateModelDao, context);
+                    }
 
-                final BlockingAggregator currentState = getBlockedStatus(sqlDao, entitySqlDaoWrapperFactory.getHandle(), state.getBlockedId(), state.getType(), bundleId, upToDate, context);
-                if (previousState != null && currentState != null) {
-                    recordBusOrFutureNotificationFromTransaction(entitySqlDaoWrapperFactory,
-                                                                 state.getId(),
-                                                                 state.getEffectiveDate(),
-                                                                 state.getBlockedId(),
-                                                                 state.getType(),
-                                                                 state.getService(),
-                                                                 previousState,
-                                                                 currentState,
-                                                                 context);
+                    final BlockingAggregator currentState = getBlockedStatus(sqlDao, entitySqlDaoWrapperFactory.getHandle(), state.getBlockedId(), state.getType(), bundleId, upToDate, context);
+                    if (previousState != null && currentState != null) {
+                        recordBusOrFutureNotificationFromTransaction(entitySqlDaoWrapperFactory,
+                                                                     state.getId(),
+                                                                     state.getEffectiveDate(),
+                                                                     state.getBlockedId(),
+                                                                     state.getType(),
+                                                                     state.getService(),
+                                                                     previousState,
+                                                                     currentState,
+                                                                     context);
+                    }
                 }
 
                 return null;
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/ProxyBlockingStateDao.java b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/ProxyBlockingStateDao.java
index 4e6a6b9..6a64cc7 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/ProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/ProxyBlockingStateDao.java
@@ -54,6 +54,7 @@ import org.skife.jdbi.v2.IDBI;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Ordering;
@@ -231,8 +232,8 @@ public class ProxyBlockingStateDao implements BlockingStateDao {
     }
 
     @Override
-    public void setBlockingStateAndPostBlockingTransitionEvent(final BlockingState state, final UUID bundleId, final InternalCallContext context) {
-        delegate.setBlockingStateAndPostBlockingTransitionEvent(state, bundleId, context);
+    public void setBlockingStatesAndPostBlockingTransitionEvent(final Map<BlockingState, Optional<UUID>> states, final InternalCallContext context) {
+        delegate.setBlockingStatesAndPostBlockingTransitionEvent(states, context);
     }
 
     @Override
@@ -241,7 +242,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
+    // See DefaultEntitlement#computeAddOnBlockingStates
     private List<BlockingState> addBlockingStatesNotOnDisk(final List<BlockingState> blockingStatesOnDisk,
                                                            final InternalTenantContext context) {
         final Collection<BlockingState> blockingStatesOnDiskCopy = new LinkedList<BlockingState>(blockingStatesOnDisk);
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/DefaultEntitlementService.java b/entitlement/src/main/java/org/killbill/billing/entitlement/DefaultEntitlementService.java
index f62b795..7e189c0 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/DefaultEntitlementService.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/DefaultEntitlementService.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2015 Groupon, Inc
- * Copyright 2014-2015 The Billing Project, LLC
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
  * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
@@ -18,10 +18,14 @@
 
 package org.killbill.billing.entitlement;
 
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collection;
 import java.util.UUID;
 
 import org.joda.time.DateTime;
 import org.killbill.billing.callcontext.InternalCallContext;
+import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.DefaultBlockingTransitionInternalEvent;
 import org.killbill.billing.entitlement.api.DefaultEntitlement;
 import org.killbill.billing.entitlement.api.Entitlement;
@@ -30,12 +34,14 @@ import org.killbill.billing.entitlement.dao.BlockingStateDao;
 import org.killbill.billing.entitlement.engine.core.BlockingTransitionNotificationKey;
 import org.killbill.billing.entitlement.engine.core.EntitlementNotificationKey;
 import org.killbill.billing.entitlement.engine.core.EntitlementNotificationKeyAction;
+import org.killbill.billing.entitlement.engine.core.EntitlementUtils;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.platform.api.LifecycleHandlerType;
 import org.killbill.billing.platform.api.LifecycleHandlerType.LifecycleLevel;
 import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.CallOrigin;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
+import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.billing.util.callcontext.UserType;
 import org.killbill.bus.api.BusEvent;
 import org.killbill.bus.api.PersistentBus;
@@ -62,6 +68,7 @@ public class DefaultEntitlementService implements EntitlementService {
     private final BlockingStateDao blockingStateDao;
     private final PersistentBus eventBus;
     private final NotificationQueueService notificationQueueService;
+    private final EntitlementUtils entitlementUtils;
     private final InternalCallContextFactory internalCallContextFactory;
 
     private NotificationQueue entitlementEventQueue;
@@ -71,11 +78,13 @@ public class DefaultEntitlementService implements EntitlementService {
                                      final BlockingStateDao blockingStateDao,
                                      final PersistentBus eventBus,
                                      final NotificationQueueService notificationQueueService,
+                                     final EntitlementUtils entitlementUtils,
                                      final InternalCallContextFactory internalCallContextFactory) {
         this.entitlementInternalApi = entitlementInternalApi;
         this.blockingStateDao = blockingStateDao;
         this.eventBus = eventBus;
         this.notificationQueueService = notificationQueueService;
+        this.entitlementUtils = entitlementUtils;
         this.internalCallContextFactory = internalCallContextFactory;
     }
 
@@ -131,7 +140,7 @@ public class DefaultEntitlementService implements EntitlementService {
         try {
             if (EntitlementNotificationKeyAction.CHANGE.equals(entitlementNotificationKeyAction) ||
                 EntitlementNotificationKeyAction.CANCEL.equals(entitlementNotificationKeyAction)) {
-                ((DefaultEntitlement) entitlement).blockAddOnsIfRequired(key.getEffectiveDate(), callContext, internalCallContext);
+                blockAddOnsIfRequired(key, (DefaultEntitlement) entitlement, callContext, internalCallContext);
             } else if (EntitlementNotificationKeyAction.PAUSE.equals(entitlementNotificationKeyAction)) {
                 entitlementInternalApi.pause(key.getBundleId(), key.getEffectiveDate().toLocalDate(), ImmutableList.<PluginProperty>of(), internalCallContext);
             } else if (EntitlementNotificationKeyAction.RESUME.equals(entitlementNotificationKeyAction)) {
@@ -142,6 +151,30 @@ public class DefaultEntitlementService implements EntitlementService {
         }
     }
 
+    private void blockAddOnsIfRequired(final EntitlementNotificationKey key, final DefaultEntitlement entitlement, final TenantContext callContext, final InternalCallContext internalCallContext) throws EntitlementApiException {
+        final Collection<NotificationEvent> notificationEvents = new ArrayList<NotificationEvent>();
+        final Collection<BlockingState> blockingStates = entitlement.computeAddOnBlockingStates(key.getEffectiveDate(), notificationEvents, callContext, internalCallContext);
+        // Record the new state first, then insert the notifications to avoid race conditions
+        entitlementUtils.setBlockingStatesAndPostBlockingTransitionEvent(blockingStates, entitlement.getBundleId(), internalCallContext);
+        for (final NotificationEvent notificationEvent : notificationEvents) {
+            recordFutureNotification(key.getEffectiveDate(), notificationEvent, internalCallContext);
+        }
+    }
+
+    private void recordFutureNotification(final DateTime effectiveDate,
+                                          final NotificationEvent notificationEvent,
+                                          final InternalCallContext context) {
+        try {
+            final NotificationQueue subscriptionEventQueue = notificationQueueService.getNotificationQueue(DefaultEntitlementService.ENTITLEMENT_SERVICE_NAME,
+                                                                                                           DefaultEntitlementService.NOTIFICATION_QUEUE_NAME);
+            subscriptionEventQueue.recordFutureNotification(effectiveDate, notificationEvent, context.getUserToken(), context.getAccountRecordId(), context.getTenantRecordId());
+        } catch (final NoSuchNotificationQueue e) {
+            throw new RuntimeException(e);
+        } catch (final IOException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private void processBlockingNotification(final BlockingTransitionNotificationKey key, final InternalCallContext internalCallContext) {
         // Check if the blocking state has been deleted since
         if (blockingStateDao.getById(key.getBlockingStateId(), internalCallContext) == null) {
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EntitlementUtils.java b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EntitlementUtils.java
index b0aee5b..a4d33ac 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EntitlementUtils.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EntitlementUtils.java
@@ -18,8 +18,10 @@
 
 package org.killbill.billing.entitlement.engine.core;
 
+import java.util.Map;
 import java.util.UUID;
 
+import javax.annotation.Nullable;
 import javax.inject.Inject;
 
 import org.killbill.billing.callcontext.InternalCallContext;
@@ -33,7 +35,9 @@ import org.killbill.billing.subscription.api.SubscriptionBaseInternalApi;
 import org.killbill.billing.subscription.api.user.SubscriptionBaseApiException;
 import org.killbill.notificationq.api.NotificationQueueService;
 
+import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 
 public class EntitlementUtils {
@@ -52,10 +56,23 @@ public class EntitlementUtils {
         this.notificationQueueService = notificationQueueService;
     }
 
-    /**
-     * @param state   new state to store
-     * @param context call context
-     */
+    public void setBlockingStatesAndPostBlockingTransitionEvent(final Iterable<BlockingState> blockingStates, @Nullable final UUID bundleId, final InternalCallContext internalCallContext) {
+        final ImmutableMap.Builder<BlockingState, Optional<UUID>> states = new ImmutableMap.Builder<BlockingState, Optional<UUID>>();
+        final Optional<UUID> bundleIdOptional = Optional.<UUID>fromNullable(bundleId);
+        for (final BlockingState blockingState : blockingStates) {
+            states.put(blockingState, bundleIdOptional);
+        }
+        dao.setBlockingStatesAndPostBlockingTransitionEvent(states.build(), internalCallContext);
+    }
+
+    public void setBlockingStateAndPostBlockingTransitionEvent(final Map<BlockingState, UUID> blockingStates, final InternalCallContext internalCallContext) {
+        final ImmutableMap.Builder<BlockingState, Optional<UUID>> states = new ImmutableMap.Builder<BlockingState, Optional<UUID>>();
+        for (final BlockingState blockingState : blockingStates.keySet()) {
+            states.put(blockingState, Optional.<UUID>fromNullable(blockingStates.get(blockingState)));
+        }
+        dao.setBlockingStatesAndPostBlockingTransitionEvent(states.build(), internalCallContext);
+    }
+
     public void setBlockingStateAndPostBlockingTransitionEvent(final BlockingState state, final InternalCallContext context) {
         UUID bundleId = null;
         if (state.getType() == BlockingStateType.SUBSCRIPTION) {
@@ -65,7 +82,7 @@ public class EntitlementUtils {
                 throw new RuntimeException(e);
             }
         }
-        dao.setBlockingStateAndPostBlockingTransitionEvent(state, bundleId, context);
+        dao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(state, Optional.<UUID>fromNullable(bundleId)), context);
     }
 
     /**
diff --git a/entitlement/src/test/java/org/killbill/billing/entitlement/block/TestBlockingChecker.java b/entitlement/src/test/java/org/killbill/billing/entitlement/block/TestBlockingChecker.java
index 37c0e6b..bb480fa 100644
--- a/entitlement/src/test/java/org/killbill/billing/entitlement/block/TestBlockingChecker.java
+++ b/entitlement/src/test/java/org/killbill/billing/entitlement/block/TestBlockingChecker.java
@@ -20,11 +20,6 @@ package org.killbill.billing.entitlement.block;
 
 import java.util.UUID;
 
-import org.mockito.Mockito;
-import org.testng.Assert;
-import org.testng.annotations.BeforeMethod;
-import org.testng.annotations.Test;
-
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.entitlement.EntitlementTestSuiteNoDB;
@@ -36,6 +31,13 @@ import org.killbill.billing.junction.DefaultBlockingState;
 import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.billing.subscription.api.user.SubscriptionBaseApiException;
 import org.killbill.billing.subscription.api.user.SubscriptionBaseBundle;
+import org.mockito.Mockito;
+import org.testng.Assert;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 
 public class TestBlockingChecker extends EntitlementTestSuiteNoDB {
 
@@ -70,20 +72,19 @@ public class TestBlockingChecker extends EntitlementTestSuiteNoDB {
         ((MockBlockingStateDao) blockingStateDao).clear();
     }
 
-
     private void setStateBundle(final boolean bC, final boolean bE, final boolean bB) {
-        final BlockingState bundleState = new DefaultBlockingState(bundle.getId(), BlockingStateType.SUBSCRIPTION_BUNDLE,"state", "test-service", bC, bE, bB, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(bundleState, null, internalCallContext);
+        final BlockingState bundleState = new DefaultBlockingState(bundle.getId(), BlockingStateType.SUBSCRIPTION_BUNDLE, "state", "test-service", bC, bE, bB, clock.getUTCNow());
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(bundleState, Optional.<UUID>absent()), internalCallContext);
     }
 
     private void setStateAccount(final boolean bC, final boolean bE, final boolean bB) {
         final BlockingState accountState = new DefaultBlockingState(account.getId(), BlockingStateType.ACCOUNT, "state", "test-service", bC, bE, bB, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(accountState, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(accountState, Optional.<UUID>absent()), internalCallContext);
     }
 
     private void setStateSubscription(final boolean bC, final boolean bE, final boolean bB) {
         final BlockingState subscriptionState = new DefaultBlockingState(subscription.getId(), BlockingStateType.SUBSCRIPTION, "state", "test-service", bC, bE, bB, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(subscriptionState, subscription.getBundleId(), internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(subscriptionState, Optional.<UUID>of(subscription.getBundleId())), internalCallContext);
     }
 
     @Test(groups = "fast")
diff --git a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/MockBlockingStateDao.java b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/MockBlockingStateDao.java
index 4fc146b..bc51a98 100644
--- a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/MockBlockingStateDao.java
+++ b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/MockBlockingStateDao.java
@@ -35,6 +35,7 @@ import org.killbill.billing.entitlement.api.EntitlementApiException;
 import org.killbill.billing.util.entity.dao.MockEntityDaoBase;
 
 import com.google.common.base.Objects;
+import com.google.common.base.Optional;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
@@ -84,16 +85,18 @@ public class MockBlockingStateDao extends MockEntityDaoBase<BlockingStateModelDa
     }
 
     @Override
-    public synchronized void setBlockingStateAndPostBlockingTransitionEvent(final BlockingState state, final UUID bundleId, final InternalCallContext context) {
-        if (blockingStates.get(state.getBlockedId()) == null) {
-            blockingStates.put(state.getBlockedId(), new ArrayList<BlockingState>());
-        }
-        blockingStates.get(state.getBlockedId()).add(state);
+    public synchronized void setBlockingStatesAndPostBlockingTransitionEvent(final Map<BlockingState, Optional<UUID>> states, final InternalCallContext context) {
+        for (final BlockingState state : states.keySet()) {
+            if (blockingStates.get(state.getBlockedId()) == null) {
+                blockingStates.put(state.getBlockedId(), new ArrayList<BlockingState>());
+            }
+            blockingStates.get(state.getBlockedId()).add(state);
 
-        if (blockingStatesPerAccountRecordId.get(context.getAccountRecordId()) == null) {
-            blockingStatesPerAccountRecordId.put(context.getAccountRecordId(), new ArrayList<BlockingState>());
+            if (blockingStatesPerAccountRecordId.get(context.getAccountRecordId()) == null) {
+                blockingStatesPerAccountRecordId.put(context.getAccountRecordId(), new ArrayList<BlockingState>());
+            }
+            blockingStatesPerAccountRecordId.get(context.getAccountRecordId()).add(state);
         }
-        blockingStatesPerAccountRecordId.get(context.getAccountRecordId()).add(state);
     }
 
     @Override
diff --git a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestBlockingDao.java b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestBlockingDao.java
index ba25a23..72c13ad 100644
--- a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestBlockingDao.java
+++ b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestBlockingDao.java
@@ -32,6 +32,9 @@ import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.junction.DefaultBlockingState;
 
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
+
 public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
 
     @BeforeMethod(groups = "slow")
@@ -55,13 +58,13 @@ public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
         clock.setDay(new LocalDate(2012, 4, 1));
 
         final BlockingState state1 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(state1, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(state1, Optional.<UUID>absent()), internalCallContext);
 
         clock.addDays(1);
 
         final String overdueStateName2 = "NoReallyThisCantGoOn";
         final BlockingState state2 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName2, service, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(state2, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(state2, Optional.<UUID>absent()), internalCallContext);
 
         Assert.assertEquals(blockingStateDao.getBlockingStateForService(uuid, BlockingStateType.ACCOUNT, service, internalCallContext).getStateName(), state2.getStateName());
 
@@ -83,14 +86,14 @@ public class TestBlockingDao extends EntitlementTestSuiteWithEmbeddedDB {
         final boolean blockBilling = false;
 
         final BlockingState state1 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName, service1, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(state1, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(state1, Optional.<UUID>absent()), internalCallContext);
         clock.setDeltaFromReality(1000 * 3600 * 24);
 
         final String service2 = "TEST2";
 
         final String overdueStateName2 = "NoReallyThisCantGoOn";
         final BlockingState state2 = new DefaultBlockingState(uuid, BlockingStateType.ACCOUNT, overdueStateName2, service2, blockChange, blockEntitlement, blockBilling, clock.getUTCNow());
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(state2, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(state2, Optional.<UUID>absent()), internalCallContext);
 
         final List<BlockingState> history2 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(history2.size(), 2);
diff --git a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestDefaultBlockingStateDao.java b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestDefaultBlockingStateDao.java
index 8935cdd..62a2919 100644
--- a/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestDefaultBlockingStateDao.java
+++ b/entitlement/src/test/java/org/killbill/billing/entitlement/dao/TestDefaultBlockingStateDao.java
@@ -40,7 +40,9 @@ import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.entitlement.api.Entitlement;
 import org.killbill.billing.junction.DefaultBlockingState;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 
 public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbeddedDB {
 
@@ -74,7 +76,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // 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.setBlockingStateAndPostBlockingTransitionEvent(blockingState, entitlement.getBundleId(), internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState, Optional.<UUID>of(entitlement.getBundleId())), internalCallContext);
 
         Assert.assertEquals(blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext).size(), 1);
     }
@@ -97,7 +99,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // Set a state for service A
         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.setBlockingStateAndPostBlockingTransitionEvent(blockingState1, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState1, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates1 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates1.size(), 1);
         Assert.assertEquals(blockingStates1.get(0).getBlockedId(), blockableId);
@@ -106,7 +108,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         Assert.assertEquals(blockingStates1.get(0).getEffectiveDate(), stateDateTime);
 
         // Set the same state again - no change
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(blockingState1, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState1, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates2 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates2.size(), 1);
         Assert.assertEquals(blockingStates2.get(0).getBlockedId(), blockableId);
@@ -116,7 +118,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.setBlockingStateAndPostBlockingTransitionEvent(blockingState2, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState2, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates3 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates3.size(), 2);
         Assert.assertEquals(blockingStates3.get(0).getBlockedId(), blockableId);
@@ -131,7 +133,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // Set the state for service A in the future - there should be no change (already effective)
         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.setBlockingStateAndPostBlockingTransitionEvent(blockingState3, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState3, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates4 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates4.size(), 2);
         Assert.assertEquals(blockingStates4.get(0).getBlockedId(), blockableId);
@@ -146,7 +148,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // Set the state for service A in the past - the new effective date should be respected
         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.setBlockingStateAndPostBlockingTransitionEvent(blockingState4, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState4, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates5 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates5.size(), 2);
         Assert.assertEquals(blockingStates5.get(0).getBlockedId(), blockableId);
@@ -161,7 +163,7 @@ public class TestDefaultBlockingStateDao extends EntitlementTestSuiteWithEmbedde
         // Set a new state for service A
         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.setBlockingStateAndPostBlockingTransitionEvent(blockingState5, null, internalCallContext);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState5, Optional.<UUID>absent()), internalCallContext);
         final List<BlockingState> blockingStates6 = blockingStateDao.getBlockingAllForAccountRecordId(internalCallContext);
         Assert.assertEquals(blockingStates6.size(), 3);
         Assert.assertEquals(blockingStates6.get(0).getBlockedId(), blockableId);
diff --git a/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBillingApi.java b/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBillingApi.java
index e6e0041..5b4748b 100644
--- a/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBillingApi.java
+++ b/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBillingApi.java
@@ -39,6 +39,7 @@ import org.killbill.billing.catalog.api.Plan;
 import org.killbill.billing.catalog.api.PlanPhase;
 import org.killbill.billing.catalog.api.PriceList;
 import org.killbill.billing.catalog.api.PriceListSet;
+import org.killbill.billing.entitlement.api.BlockingState;
 import org.killbill.billing.entitlement.api.BlockingStateType;
 import org.killbill.billing.entitlement.api.Entitlement.EntitlementState;
 import org.killbill.billing.entitlement.dao.MockBlockingStateDao;
@@ -61,7 +62,9 @@ import org.testng.Assert;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
 
+import com.google.common.base.Optional;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNull;
@@ -182,8 +185,9 @@ public class TestBillingApi extends JunctionTestSuiteNoDB {
 
         final Account account = createAccount(32);
 
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(new DefaultBlockingState(bunId, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, now.plusDays(1)), null, internalCallContext);
-        blockingStateDao.setBlockingStateAndPostBlockingTransitionEvent(new DefaultBlockingState(bunId, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, now.plusDays(2)), null, internalCallContext);
+        final BlockingState blockingState1 = new DefaultBlockingState(bunId, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, now.plusDays(1));
+        final BlockingState blockingState2 = new DefaultBlockingState(bunId, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, now.plusDays(2));
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState1, Optional.<UUID>absent(), blockingState2, Optional.<UUID>absent()), internalCallContext);
 
         final SortedSet<BillingEvent> events = billingInternalApi.getBillingEventsForAccountAndUpdateAccountBCD(account.getId(), null, internalCallContext);
 
diff --git a/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBlockingCalculator.java b/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBlockingCalculator.java
index 4477367..8c26b09 100644
--- a/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBlockingCalculator.java
+++ b/junction/src/test/java/org/killbill/billing/junction/plumbing/billing/TestBlockingCalculator.java
@@ -31,10 +31,6 @@ import java.util.UUID;
 import org.joda.time.DateTime;
 import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
-import org.mockito.Mockito;
-import org.testng.annotations.BeforeMethod;
-import org.testng.annotations.Test;
-
 import org.killbill.billing.account.api.Account;
 import org.killbill.billing.catalog.MockPlan;
 import org.killbill.billing.catalog.MockPlanPhase;
@@ -51,6 +47,12 @@ import org.killbill.billing.junction.JunctionTestSuiteNoDB;
 import org.killbill.billing.junction.plumbing.billing.BlockingCalculator.DisabledDuration;
 import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.billing.subscription.api.SubscriptionBaseTransitionType;
+import org.mockito.Mockito;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableMap;
 
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.assertNotNull;
@@ -115,11 +117,12 @@ public class TestBlockingCalculator extends JunctionTestSuiteNoDB {
         billingEvents.add(C);
         billingEvents.add(D);
 
-        final List<BlockingState> blockingStates = new ArrayList<BlockingState>();
-        blockingStates.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, now));
-        blockingStates.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, now.plusDays(2)));
+        final BlockingState blockingState1 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, now);
+        final BlockingState blockingState2 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, now.plusDays(2));
 
-        setBlockingStates(blockingStates);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState1, Optional.<UUID>absent(),
+                                                                                                                        blockingState2, Optional.<UUID>absent()),
+                                                                         internalCallContext);
 
         blockingCalculator.insertBlockingEvents(billingEvents, new HashSet<UUID>(), internalCallContext);
 
@@ -745,13 +748,16 @@ public class TestBlockingCalculator extends JunctionTestSuiteNoDB {
         billingEvents.add(phase);
         billingEvents.add(upgrade);
 
-        final List<BlockingState> blockingEvents = new ArrayList<BlockingState>();
-        blockingEvents.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, false, false, new LocalDate(2012, 7, 5).toDateTimeAtStartOfDay(DateTimeZone.UTC)));
-        blockingEvents.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, new LocalDate(2012, 7, 15).toDateTimeAtStartOfDay(DateTimeZone.UTC)));
-        blockingEvents.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, new LocalDate(2012, 7, 24).toDateTimeAtStartOfDay(DateTimeZone.UTC)));
-        blockingEvents.add(new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, new LocalDate(2012, 7, 25).toDateTimeAtStartOfDay(DateTimeZone.UTC)));
+        final BlockingState blockingState1 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, false, false, new LocalDate(2012, 7, 5).toDateTimeAtStartOfDay(DateTimeZone.UTC));
+        final BlockingState blockingState2 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, new LocalDate(2012, 7, 15).toDateTimeAtStartOfDay(DateTimeZone.UTC));
+        final BlockingState blockingState3 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, DISABLED_BUNDLE, "test", true, true, true, new LocalDate(2012, 7, 24).toDateTimeAtStartOfDay(DateTimeZone.UTC));
+        final BlockingState blockingState4 = new DefaultBlockingState(bundleId1, BlockingStateType.SUBSCRIPTION_BUNDLE, CLEAR_BUNDLE, "test", false, false, false, new LocalDate(2012, 7, 25).toDateTimeAtStartOfDay(DateTimeZone.UTC));
 
-        setBlockingStates(blockingEvents);
+        blockingStateDao.setBlockingStatesAndPostBlockingTransitionEvent(ImmutableMap.<BlockingState, Optional<UUID>>of(blockingState1, Optional.<UUID>absent(),
+                                                                                                                        blockingState2, Optional.<UUID>absent(),
+                                                                                                                        blockingState3, Optional.<UUID>absent(),
+                                                                                                                        blockingState4, Optional.<UUID>absent()),
+                                                                         internalCallContext);
 
         blockingCalculator.insertBlockingEvents(billingEvents, new HashSet<UUID>(), internalCallContext);
 
@@ -768,10 +774,4 @@ public class TestBlockingCalculator extends JunctionTestSuiteNoDB {
         assertEquals(events.get(4).getEffectiveDate(), new LocalDate(2012, 7, 25).toDateTimeAtStartOfDay(DateTimeZone.UTC));
         assertEquals(events.get(4).getTransitionType(), SubscriptionBaseTransitionType.CHANGE);
     }
-
-    private void setBlockingStates(final List<BlockingState> blockingStates) {
-        for (final BlockingState blockingState : blockingStates) {
-            blockingStateDao.setBlockingState(blockingState, internalCallContext);
-        }
-    }
 }