killbill-memoizeit

Details

diff --git a/api/src/main/java/org/killbill/billing/entitlement/AccountEventsStreams.java b/api/src/main/java/org/killbill/billing/entitlement/AccountEventsStreams.java
index bd674fb..e374d62 100644
--- a/api/src/main/java/org/killbill/billing/entitlement/AccountEventsStreams.java
+++ b/api/src/main/java/org/killbill/billing/entitlement/AccountEventsStreams.java
@@ -1,7 +1,8 @@
 /*
- * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 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:
  *
@@ -21,6 +22,7 @@ import java.util.Map;
 import java.util.UUID;
 
 import org.killbill.billing.account.api.ImmutableAccountData;
+import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.billing.subscription.api.user.SubscriptionBaseBundle;
 
 // Wrapper object to save on DAO calls
@@ -31,6 +33,9 @@ public interface AccountEventsStreams {
     // Map bundle id -> bundle
     public Map<UUID, SubscriptionBaseBundle> getBundles();
 
+    // Map bundle id -> subscriptions
+    public Map<UUID, Collection<SubscriptionBase>> getSubscriptions();
+
     // Map bundle id -> events streams
     public Map<UUID, Collection<EventsStream>> getEventsStreams();
 }
diff --git a/api/src/main/java/org/killbill/billing/subscription/api/SubscriptionBaseWithAddOns.java b/api/src/main/java/org/killbill/billing/subscription/api/SubscriptionBaseWithAddOns.java
index 5753c1a..3fc0da9 100644
--- a/api/src/main/java/org/killbill/billing/subscription/api/SubscriptionBaseWithAddOns.java
+++ b/api/src/main/java/org/killbill/billing/subscription/api/SubscriptionBaseWithAddOns.java
@@ -1,5 +1,6 @@
 /*
- * Copyright 2016 The Billing Project, LLC
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 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
@@ -17,14 +18,13 @@
 package org.killbill.billing.subscription.api;
 
 import java.util.List;
-import java.util.UUID;
 
 import org.joda.time.DateTime;
-import org.killbill.billing.subscription.api.SubscriptionBase;
+import org.killbill.billing.subscription.api.user.SubscriptionBaseBundle;
 
 public interface SubscriptionBaseWithAddOns {
 
-    public UUID getBundleId();
+    public SubscriptionBaseBundle getBundle();
 
     public List<SubscriptionBase> getSubscriptionBaseList();
 
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 4497b98..79f82d7 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
@@ -106,18 +106,7 @@ public class DefaultEntitlement extends EntityBase implements Entitlement {
     // Refresh-able
     protected EventsStream eventsStream;
 
-    public DefaultEntitlement(final UUID accountId, final UUID entitlementId, final EventsStreamBuilder eventsStreamBuilder,
-                              final EntitlementApi entitlementApi, final EntitlementPluginExecution pluginExecution, final BlockingStateDao blockingStateDao,
-                              final SubscriptionBaseInternalApi subscriptionInternalApi, final BlockingChecker checker,
-                              final NotificationQueueService notificationQueueService, final EntitlementUtils entitlementUtils,
-                              final EntitlementDateHelper dateHelper, final Clock clock, final SecurityApi securityApi,
-                              final InternalCallContextFactory internalCallContextFactory, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
-        this(eventsStreamBuilder.buildForEntitlement(entitlementId, internalTenantContext), eventsStreamBuilder,
-             entitlementApi, pluginExecution, blockingStateDao, subscriptionInternalApi, checker, notificationQueueService,
-             entitlementUtils, dateHelper, clock, securityApi, internalTenantContext, internalCallContextFactory);
-    }
-
-    public DefaultEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase subscription, final List<SubscriptionBase> allSubscriptionsForBundle, final EventsStreamBuilder eventsStreamBuilder,
+    public DefaultEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase subscription, final Collection<SubscriptionBase> allSubscriptionsForBundle, final EventsStreamBuilder eventsStreamBuilder,
                               final EntitlementApi entitlementApi, final EntitlementPluginExecution pluginExecution, final BlockingStateDao blockingStateDao,
                               final SubscriptionBaseInternalApi subscriptionInternalApi, final BlockingChecker checker,
                               final NotificationQueueService notificationQueueService, final EntitlementUtils entitlementUtils,
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 4c55b8b..ede2566 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
@@ -36,6 +36,7 @@ import org.killbill.billing.account.api.AccountInternalApi;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.BillingActionPolicy;
+import org.killbill.billing.catalog.api.Catalog;
 import org.killbill.billing.catalog.api.PlanPhasePriceOverride;
 import org.killbill.billing.catalog.api.PlanPhaseSpecifier;
 import org.killbill.billing.catalog.api.ProductCategory;
@@ -70,9 +71,11 @@ import org.killbill.notificationq.api.NotificationQueueService;
 
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;
+import com.google.common.collect.HashMultimap;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
+import com.google.common.collect.Multimap;
 
 import static org.killbill.billing.entitlement.logging.EntitlementLoggingHelper.logCreateEntitlement;
 import static org.killbill.billing.entitlement.logging.EntitlementLoggingHelper.logCreateEntitlementsWithAOs;
@@ -272,11 +275,11 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                                                                                          EntitlementService.ENTITLEMENT_SERVICE_NAME,
                                                                                          false, false, false,
                                                                                          dateHelper.fromLocalDateAndReferenceTime(baseEntitlementWithAddOnsSpecifier.getEntitlementEffectiveDate(), now, contextWithValidAccountRecordId));
-                            blockingStateMap.put(blockingState, subscriptionsWithAddOns.get(i).getBundleId());
+                            blockingStateMap.put(blockingState, subscriptionsWithAddOns.get(i).getBundle().getId());
                         }
                     }
                     entitlementUtils.setBlockingStateAndPostBlockingTransitionEvent(blockingStateMap, contextWithValidAccountRecordId);
-                    return buildEntitlementList(accountId, subscriptionsWithAddOns, contextWithValidAccountRecordId);
+                    return buildEntitlementList(subscriptionsWithAddOns, contextWithValidAccountRecordId);
                 } catch (final SubscriptionBaseApiException e) {
                     throw new EntitlementApiException(e);
                 }
@@ -285,12 +288,16 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
         return pluginExecution.executeWithPlugin(createBaseEntitlementsWithAddOns, pluginContext);
     }
 
-    private List<Entitlement> buildEntitlementList(final UUID accountId, final Iterable<SubscriptionBaseWithAddOns> subscriptionsWithAddOns, final InternalCallContext callContext) throws EntitlementApiException {
+    private List<Entitlement> buildEntitlementList(final Iterable<SubscriptionBaseWithAddOns> subscriptionsWithAddOns, final InternalCallContext internalCallContext) throws EntitlementApiException {
+        // Need to refresh all latest state as some bundles might have existed already
+        final AccountEventsStreams accountEventsStreams = eventsStreamBuilder.buildForAccount(internalCallContext);
+
         final List<Entitlement> result = new ArrayList<Entitlement>();
         for (final SubscriptionBaseWithAddOns subscriptionWithAddOns : subscriptionsWithAddOns) {
             for (final SubscriptionBase subscriptionBase : subscriptionWithAddOns.getSubscriptionBaseList()) {
-                final Entitlement entitlement = new DefaultEntitlement(accountId,
-                                                                       subscriptionBase.getId(),
+                final Entitlement entitlement = new DefaultEntitlement(subscriptionWithAddOns.getBundle(),
+                                                                       subscriptionBase,
+                                                                       accountEventsStreams.getSubscriptions().get(subscriptionBase.getBundleId()),
                                                                        eventsStreamBuilder,
                                                                        entitlementApi,
                                                                        pluginExecution,
@@ -303,7 +310,7 @@ public class DefaultEntitlementApi extends DefaultEntitlementApiBase implements 
                                                                        clock,
                                                                        securityApi,
                                                                        internalCallContextFactory,
-                                                                       callContext);
+                                                                       internalCallContext);
                 result.add(entitlement);
             }
         }
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultAccountEventsStreams.java b/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultAccountEventsStreams.java
index d1b09a8..362f05e 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultAccountEventsStreams.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/api/svcs/DefaultAccountEventsStreams.java
@@ -1,7 +1,8 @@
 /*
- * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 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:
  *
@@ -24,6 +25,7 @@ import java.util.UUID;
 import org.killbill.billing.account.api.ImmutableAccountData;
 import org.killbill.billing.entitlement.AccountEventsStreams;
 import org.killbill.billing.entitlement.EventsStream;
+import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.billing.subscription.api.user.SubscriptionBaseBundle;
 
 import com.google.common.collect.ImmutableList;
@@ -33,12 +35,15 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
 
     private final ImmutableAccountData account;
     private final Map<UUID, Collection<EventsStream>> eventsStreams;
+    private final Map<UUID, Collection<SubscriptionBase>> subscriptionsPerBundle;
     private final Map<UUID, SubscriptionBaseBundle> bundles = new HashMap<UUID, SubscriptionBaseBundle>();
 
     public DefaultAccountEventsStreams(final ImmutableAccountData account,
                                        final Iterable<SubscriptionBaseBundle> bundles,
+                                       final Map<UUID, Collection<SubscriptionBase>> subscriptionsPerBundle,
                                        final Map<UUID, Collection<EventsStream>> eventsStreams) {
         this.account = account;
+        this.subscriptionsPerBundle = subscriptionsPerBundle;
         this.eventsStreams = eventsStreams;
         for (final SubscriptionBaseBundle baseBundle : bundles) {
             this.bundles.put(baseBundle.getId(), baseBundle);
@@ -46,7 +51,7 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
     }
 
     public DefaultAccountEventsStreams(final ImmutableAccountData account) {
-        this(account, ImmutableList.<SubscriptionBaseBundle>of(), ImmutableMap.<UUID, Collection<EventsStream>>of());
+        this(account, ImmutableList.<SubscriptionBaseBundle>of(), ImmutableMap.<UUID, Collection<SubscriptionBase>>of(), ImmutableMap.<UUID, Collection<EventsStream>>of());
     }
 
     @Override
@@ -60,6 +65,11 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
     }
 
     @Override
+    public Map<UUID, Collection<SubscriptionBase>> getSubscriptions() {
+        return subscriptionsPerBundle;
+    }
+
+    @Override
     public Map<UUID, Collection<EventsStream>> getEventsStreams() {
         return eventsStreams;
     }
@@ -69,6 +79,7 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
         final StringBuilder sb = new StringBuilder("DefaultAccountEventsStreams{");
         sb.append("account=").append(account);
         sb.append(", eventsStreams=").append(eventsStreams);
+        sb.append(", subscriptionsPerBundle=").append(subscriptionsPerBundle);
         sb.append(", bundles=").append(bundles);
         sb.append('}');
         return sb.toString();
@@ -91,6 +102,9 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
         if (bundles != null ? !bundles.equals(that.bundles) : that.bundles != null) {
             return false;
         }
+        if (subscriptionsPerBundle != null ? !subscriptionsPerBundle.equals(that.subscriptionsPerBundle) : that.subscriptionsPerBundle != null) {
+            return false;
+        }
         if (eventsStreams != null ? !eventsStreams.equals(that.eventsStreams) : that.eventsStreams != null) {
             return false;
         }
@@ -102,6 +116,7 @@ public class DefaultAccountEventsStreams implements AccountEventsStreams {
     public int hashCode() {
         int result = account != null ? account.hashCode() : 0;
         result = 31 * result + (eventsStreams != null ? eventsStreams.hashCode() : 0);
+        result = 31 * result + (subscriptionsPerBundle != null ? subscriptionsPerBundle.hashCode() : 0);
         result = 31 * result + (bundles != null ? bundles.hashCode() : 0);
         return result;
     }
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
index fea8c12..7df7117 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/dao/OptimizedProxyBlockingStateDao.java
@@ -18,8 +18,8 @@
 
 package org.killbill.billing.entitlement.dao;
 
+import java.util.Collection;
 import java.util.LinkedList;
-import java.util.List;
 
 import javax.annotation.Nullable;
 
@@ -74,15 +74,15 @@ public class OptimizedProxyBlockingStateDao extends ProxyBlockingStateDao {
      * @return blocking states for that subscription
      * @throws EntitlementApiException
      */
-    public List<BlockingState> getBlockingHistory(final List<BlockingState> subscriptionBlockingStatesOnDisk,
-                                                  final List<BlockingState> allBlockingStatesOnDiskForAccount,
-                                                  final ImmutableAccountData account,
-                                                  final SubscriptionBaseBundle bundle,
-                                                  @Nullable final SubscriptionBase baseSubscription,
-                                                  final SubscriptionBase subscription,
-                                                  final List<SubscriptionBase> allSubscriptionsForBundle,
-                                                  final Catalog catalog,
-                                                  final InternalTenantContext context) throws EntitlementApiException {
+    public Collection<BlockingState> getBlockingHistory(final Collection<BlockingState> subscriptionBlockingStatesOnDisk,
+                                                        final Collection<BlockingState> allBlockingStatesOnDiskForAccount,
+                                                        final ImmutableAccountData account,
+                                                        final SubscriptionBaseBundle bundle,
+                                                        @Nullable final SubscriptionBase baseSubscription,
+                                                        final SubscriptionBase subscription,
+                                                        final Collection<SubscriptionBase> allSubscriptionsForBundle,
+                                                        final Catalog catalog,
+                                                        final InternalTenantContext context) throws EntitlementApiException {
         // 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
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/DefaultEventsStream.java b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/DefaultEventsStream.java
index 878f1a0..667b114 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/DefaultEventsStream.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/DefaultEventsStream.java
@@ -60,14 +60,14 @@ public class DefaultEventsStream implements EventsStream {
     private final ImmutableAccountData account;
     private final SubscriptionBaseBundle bundle;
     // All blocking states for the account, associated bundle or subscription
-    private final List<BlockingState> blockingStates;
+    private final Collection<BlockingState> blockingStates;
     private final BlockingChecker blockingChecker;
     // Base subscription for the bundle if it exists, null otherwise
     private final SubscriptionBase baseSubscription;
     // Subscription associated with this entitlement (equals to baseSubscription for base subscriptions)
     private final SubscriptionBase subscription;
     // All subscriptions for that bundle
-    private final List<SubscriptionBase> allSubscriptionsForBundle;
+    private final Collection<SubscriptionBase> allSubscriptionsForBundle;
     private final InternalTenantContext internalTenantContext;
     private final DateTime utcNow;
     private final LocalDate utcToday;
@@ -86,9 +86,9 @@ public class DefaultEventsStream implements EventsStream {
     private EntitlementState entitlementState;
 
     public DefaultEventsStream(final ImmutableAccountData account, final SubscriptionBaseBundle bundle,
-                               final List<BlockingState> blockingStates, final BlockingChecker blockingChecker,
+                               final Collection<BlockingState> blockingStates, final BlockingChecker blockingChecker,
                                @Nullable final SubscriptionBase baseSubscription, final SubscriptionBase subscription,
-                               final List<SubscriptionBase> allSubscriptionsForBundle,
+                               final Collection<SubscriptionBase> allSubscriptionsForBundle,
                                final int defaultBillCycleDayLocal,
                                final InternalTenantContext contextWithValidAccountRecordId, final DateTime utcNow) {
         sanityChecks(account, bundle, baseSubscription, subscription);
diff --git a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
index 7bbce78..63281e8 100644
--- a/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
+++ b/entitlement/src/main/java/org/killbill/billing/entitlement/engine/core/EventsStreamBuilder.java
@@ -194,15 +194,19 @@ public class EventsStreamBuilder {
 
         // Build the EventsStream objects
         final Map<UUID, Integer> bcdCache = new HashMap<UUID, Integer>();
-        final Map<UUID, Collection<EventsStream>> entitlementsPerBundle = new HashMap<UUID, Collection<EventsStream>>();
+        final Map<UUID, Collection<EventsStream>> eventsStreamPerBundle = new HashMap<UUID, Collection<EventsStream>>();
+        final Map<UUID, Collection<SubscriptionBase>> subscriptionsPerBundle = new HashMap<UUID, Collection<SubscriptionBase>>();
         for (final UUID bundleId : subscriptions.keySet()) {
             final SubscriptionBaseBundle bundle = bundlesPerId.get(bundleId);
             final List<SubscriptionBase> allSubscriptionsForBundle = subscriptions.get(bundleId);
             final SubscriptionBase baseSubscription = findBaseSubscription(allSubscriptionsForBundle);
             final List<BlockingState> bundleBlockingStates = MoreObjects.firstNonNull(blockingStatesPerBundle.get(bundleId), ImmutableList.<BlockingState>of());
 
-            if (entitlementsPerBundle.get(bundleId) == null) {
-                entitlementsPerBundle.put(bundleId, new LinkedList<EventsStream>());
+            if (eventsStreamPerBundle.get(bundleId) == null) {
+                eventsStreamPerBundle.put(bundleId, new LinkedList<EventsStream>());
+            }
+            if (subscriptionsPerBundle.get(bundleId) == null) {
+                subscriptionsPerBundle.put(bundleId, allSubscriptionsForBundle);
             }
 
             for (final SubscriptionBase subscription : allSubscriptionsForBundle) {
@@ -212,7 +216,7 @@ public class EventsStreamBuilder {
                 // for an add-on - which means going through the magic of ProxyBlockingStateDao, which will recursively
                 // create EventsStream objects. To avoid an infinite recursion, bypass ProxyBlockingStateDao when it's not
                 // needed, i.e. if this EventStream is for a standalone or a base subscription
-                final List<BlockingState> subscriptionBlockingStates;
+                final Collection<BlockingState> subscriptionBlockingStates;
                 if (baseSubscription == null || subscription.getId().equals(baseSubscription.getId())) {
                     subscriptionBlockingStates = subscriptionBlockingStatesOnDisk;
                 } else {
@@ -243,11 +247,11 @@ public class EventsStreamBuilder {
                                                                      bcdCache,
                                                                      catalog,
                                                                      internalTenantContext);
-                entitlementsPerBundle.get(bundleId).add(eventStream);
+                eventsStreamPerBundle.get(bundleId).add(eventStream);
             }
         }
 
-        return new DefaultAccountEventsStreams(account, bundles, entitlementsPerBundle);
+        return new DefaultAccountEventsStreams(account, bundles, subscriptionsPerBundle, eventsStreamPerBundle);
     }
 
     public EventsStream buildForEntitlement(final UUID entitlementId, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
@@ -265,7 +269,7 @@ public class EventsStreamBuilder {
         return buildForEntitlement(bundle, subscription, subscriptionsForBundle, internalTenantContext);
     }
 
-    public EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase subscription, final List<SubscriptionBase> allSubscriptionsForBundle, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
+    public EventsStream buildForEntitlement(final SubscriptionBaseBundle bundle, final SubscriptionBase subscription, final Collection<SubscriptionBase> allSubscriptionsForBundle, final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         final SubscriptionBase baseSubscription = findBaseSubscription(allSubscriptionsForBundle);
 
         final ImmutableAccountData account;
@@ -285,23 +289,23 @@ public class EventsStreamBuilder {
     }
 
     // Special signature for OptimizedProxyBlockingStateDao to save some DAO calls
-    public EventsStream buildForEntitlement(final List<BlockingState> blockingStatesForAccount,
+    public EventsStream buildForEntitlement(final Collection<BlockingState> blockingStatesForAccount,
                                             final ImmutableAccountData account,
                                             final SubscriptionBaseBundle bundle,
                                             final SubscriptionBase baseSubscription,
-                                            final List<SubscriptionBase> allSubscriptionsForBundle,
+                                            final Collection<SubscriptionBase> allSubscriptionsForBundle,
                                             final Catalog catalog,
                                             final InternalTenantContext internalTenantContext) throws EntitlementApiException {
         final Map<UUID, Integer> bcdCache = new HashMap<UUID, Integer>();
         return buildForEntitlement(blockingStatesForAccount, account, bundle, baseSubscription, baseSubscription, allSubscriptionsForBundle, bcdCache, catalog, internalTenantContext);
     }
 
-    private EventsStream buildForEntitlement(final List<BlockingState> blockingStatesForAccount,
+    private EventsStream buildForEntitlement(final Collection<BlockingState> blockingStatesForAccount,
                                              final ImmutableAccountData account,
                                              final SubscriptionBaseBundle bundle,
                                              @Nullable final SubscriptionBase baseSubscription,
                                              final SubscriptionBase subscription,
-                                             final List<SubscriptionBase> allSubscriptionsForBundle,
+                                             final Collection<SubscriptionBase> allSubscriptionsForBundle,
                                              final Map<UUID, Integer> bcdCache,
                                              final Catalog catalog,
                                              final InternalTenantContext internalTenantContext) throws EntitlementApiException {
@@ -363,8 +367,8 @@ public class EventsStreamBuilder {
                                              final SubscriptionBaseBundle bundle,
                                              @Nullable final SubscriptionBase baseSubscription,
                                              final SubscriptionBase subscription,
-                                             final List<SubscriptionBase> allSubscriptionsForBundle,
-                                             final List<BlockingState> blockingStates,
+                                             final Collection<SubscriptionBase> allSubscriptionsForBundle,
+                                             final Collection<BlockingState> blockingStates,
                                              final Map<UUID, Integer> bcdCache,
                                              final Catalog catalog,
                                              final InternalTenantContext internalTenantContext) throws EntitlementApiException {
diff --git a/subscription/src/main/java/org/killbill/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java b/subscription/src/main/java/org/killbill/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
index 2406202..7165973 100644
--- a/subscription/src/main/java/org/killbill/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
+++ b/subscription/src/main/java/org/killbill/billing/subscription/api/svcs/DefaultSubscriptionInternalApi.java
@@ -306,18 +306,17 @@ public class DefaultSubscriptionInternalApi extends SubscriptionApiBase implemen
                     }
                 }
 
-                final SubscriptionAndAddOnsSpecifier subscriptionAndAddOnsSpecifier = new SubscriptionAndAddOnsSpecifier(
-                        bundle.getId(),
-                        effectiveDate,
-                        verifyAndBuildSubscriptionSpecifiers(bundle.getId(),
-                                                             bundle.getExternalKey(),
-                                                             reorderedSpecifiers,
-                                                             entitlementWithAddOnsSpecifier.isMigrated(),
-                                                             context,
-                                                             now,
-                                                             effectiveDate,
-                                                             catalog,
-                                                             callContext)
+                final SubscriptionAndAddOnsSpecifier subscriptionAndAddOnsSpecifier = new SubscriptionAndAddOnsSpecifier(bundle,
+                                                                                                                         effectiveDate,
+                                                                                                                         verifyAndBuildSubscriptionSpecifiers(bundle.getId(),
+                                                                                                                                                              bundle.getExternalKey(),
+                                                                                                                                                              reorderedSpecifiers,
+                                                                                                                                                              entitlementWithAddOnsSpecifier.isMigrated(),
+                                                                                                                                                              context,
+                                                                                                                                                              now,
+                                                                                                                                                              effectiveDate,
+                                                                                                                                                              catalog,
+                                                                                                                                                              callContext)
                 );
                 subscriptionAndAddOns.add(subscriptionAndAddOnsSpecifier);
             }
diff --git a/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java b/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
index 807a2ad..1e1fc4e 100644
--- a/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
+++ b/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseApiService.java
@@ -135,7 +135,7 @@ public class DefaultSubscriptionBaseApiService implements SubscriptionBaseApiSer
                 createEvents(subscriptionAndAddOns.getSubscriptionSpecifiers(), context, eventsMap, subscriptionBaseList, fullCatalog);
                 subscriptionBaseAndAddOnsList.add(subscriptionBaseList);
 
-                final SubscriptionBaseWithAddOns subscriptionBaseWithAddOns = new DefaultSubscriptionBaseWithAddOns(subscriptionAndAddOns.getBundleId(),
+                final SubscriptionBaseWithAddOns subscriptionBaseWithAddOns = new DefaultSubscriptionBaseWithAddOns(subscriptionAndAddOns.getBundle(),
                                                                                                                     subscriptionBaseList,
                                                                                                                     subscriptionAndAddOns.getEffectiveDate());
                 allSubscriptions.add(subscriptionBaseWithAddOns);
diff --git a/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseWithAddOns.java b/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseWithAddOns.java
index 0050afd..c2cbbe7 100644
--- a/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseWithAddOns.java
+++ b/subscription/src/main/java/org/killbill/billing/subscription/api/user/DefaultSubscriptionBaseWithAddOns.java
@@ -1,5 +1,6 @@
 /*
- * Copyright 2016 The Billing Project, LLC
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 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
@@ -17,7 +18,6 @@
 package org.killbill.billing.subscription.api.user;
 
 import java.util.List;
-import java.util.UUID;
 
 import org.joda.time.DateTime;
 import org.killbill.billing.subscription.api.SubscriptionBase;
@@ -29,19 +29,19 @@ public class DefaultSubscriptionBaseWithAddOns implements SubscriptionBaseWithAd
 
     private static final Logger log = LoggerFactory.getLogger(DefaultSubscriptionBaseWithAddOns.class);
 
-    private final UUID bundleId;
+    private final SubscriptionBaseBundle bundle;
     private final List<SubscriptionBase> subscriptionBaseList;
     private final DateTime effectiveDate;
 
-    public DefaultSubscriptionBaseWithAddOns(final UUID bundleId, final List<SubscriptionBase> subscriptionBaseList, final DateTime effectiveDate) {
-        this.bundleId = bundleId;
+    public DefaultSubscriptionBaseWithAddOns(final SubscriptionBaseBundle bundle, final List<SubscriptionBase> subscriptionBaseList, final DateTime effectiveDate) {
+        this.bundle = bundle;
         this.subscriptionBaseList = subscriptionBaseList;
         this.effectiveDate = effectiveDate;
     }
 
     @Override
-    public UUID getBundleId() {
-        return bundleId;
+    public SubscriptionBaseBundle getBundle() {
+        return bundle;
     }
 
     @Override
diff --git a/subscription/src/main/java/org/killbill/billing/subscription/api/user/SubscriptionAndAddOnsSpecifier.java b/subscription/src/main/java/org/killbill/billing/subscription/api/user/SubscriptionAndAddOnsSpecifier.java
index 549335e..3b78ac5 100644
--- a/subscription/src/main/java/org/killbill/billing/subscription/api/user/SubscriptionAndAddOnsSpecifier.java
+++ b/subscription/src/main/java/org/killbill/billing/subscription/api/user/SubscriptionAndAddOnsSpecifier.java
@@ -1,5 +1,6 @@
 /*
- * Copyright 2016 The Billing Project, LLC
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 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
@@ -17,28 +18,25 @@
 package org.killbill.billing.subscription.api.user;
 
 import java.util.List;
-import java.util.UUID;
 
 import org.joda.time.DateTime;
 
 public class SubscriptionAndAddOnsSpecifier {
 
-    private UUID bundleId;
-    private DateTime effectiveDate;
+    private final SubscriptionBaseBundle bundle;
+    private final DateTime effectiveDate;
     private List<SubscriptionSpecifier> subscriptionSpecifiers;
 
-    public SubscriptionAndAddOnsSpecifier() {
-    }
-
-    public SubscriptionAndAddOnsSpecifier(final UUID bundleId, final DateTime effectiveDate,
+    public SubscriptionAndAddOnsSpecifier(final SubscriptionBaseBundle bundle,
+                                          final DateTime effectiveDate,
                                           final List<SubscriptionSpecifier> subscriptionSpecifiers) {
-        this.bundleId = bundleId;
+        this.bundle = bundle;
         this.effectiveDate = effectiveDate;
         this.subscriptionSpecifiers = subscriptionSpecifiers;
     }
 
-    public UUID getBundleId() {
-        return bundleId;
+    public SubscriptionBaseBundle getBundle() {
+        return bundle;
     }
 
     public DateTime getEffectiveDate() {