killbill-memoizeit

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationDryRunInvoice.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationDryRunInvoice.java
index fc68f6a..53bb8d3 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationDryRunInvoice.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationDryRunInvoice.java
@@ -166,33 +166,54 @@ public class TestIntegrationDryRunInvoice extends TestIntegrationBase {
         final DateTime secondSubscriptionCreationDate = new DateTime(2014, 12, 15, 0, 0, 0, 0, testTimeZone);
         clock.setTime(secondSubscriptionCreationDate);
 
-        // Create the monthly
-        DefaultEntitlement baseEntitlementMonthly = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey", "Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
-        DefaultSubscriptionBase subscriptionMonthly = subscriptionDataFromSubscription(baseEntitlementMonthly.getSubscriptionBase());
+        // Create the first monthly
+        DefaultEntitlement baseEntitlementMonthly1 = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey1", "Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        DefaultSubscriptionBase subscriptionMonthly1 = subscriptionDataFromSubscription(baseEntitlementMonthly1.getSubscriptionBase());
         invoiceChecker.checkInvoice(account.getId(), invoiceItemCount++, callContext, new ExpectedInvoiceItemCheck(secondSubscriptionCreationDate.toLocalDate(), null, InvoiceItemType.FIXED, new BigDecimal("0")));
         // No end date for the trial item (fixed price of zero), and CTD should be today (i.e. when the trial started)
-        invoiceChecker.checkChargedThroughDate(subscriptionMonthly.getId(), clock.getUTCToday(), callContext);
+        invoiceChecker.checkChargedThroughDate(subscriptionMonthly1.getId(), clock.getUTCToday(), callContext);
+
+
+        // Create the second monthly
+        DefaultEntitlement baseEntitlementMonthly2 = createBaseEntitlementAndCheckForCompletion(account.getId(), "bundleKey2", "Pistol", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        DefaultSubscriptionBase subscriptionMonthly2 = subscriptionDataFromSubscription(baseEntitlementMonthly2.getSubscriptionBase());
+        invoiceChecker.checkInvoice(account.getId(), invoiceItemCount++, callContext, new ExpectedInvoiceItemCheck(secondSubscriptionCreationDate.toLocalDate(), null, InvoiceItemType.FIXED, new BigDecimal("0")));
+        // No end date for the trial item (fixed price of zero), and CTD should be today (i.e. when the trial started)
+        invoiceChecker.checkChargedThroughDate(subscriptionMonthly2.getId(), clock.getUTCToday(), callContext);
+
 
         // Verify next dryRun invoice and then move the clock to that date to also verify real invoice is the same
         expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 1, 14), new LocalDate(2015, 2, 14), InvoiceItemType.RECURRING, new BigDecimal("249.95")));
+        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 1, 14), new LocalDate(2015, 2, 14), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
         dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, DRY_RUN_UPCOMING_INVOICE_ARG, callContext);
         invoiceChecker.checkInvoiceNoAudits(dryRunInvoice, callContext, expectedInvoices);
 
-        busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
+        busHandler.pushExpectedEvents(NextEvent.PHASE, NextEvent.PHASE, NextEvent.INVOICE,  NextEvent.NULL_INVOICE, NextEvent.PAYMENT, NextEvent.INVOICE_PAYMENT);
         // 2015-1-14
         clock.addDays(30);
         assertListenerStatus();
         invoiceChecker.checkInvoice(account.getId(), invoiceItemCount++, callContext, expectedInvoices);
         expectedInvoices.clear();
 
-        // We test first the next expected invoice for a specific subscription: We can see the targetDate is 2015-2-14 and not 2015-2-1, and also we only see the item for the monthly subscription
-        final DryRunArguments dryRunUpcomingInvoiceWithFilterArg = new TestDryRunArguments(DryRunType.UPCOMING_INVOICE, null, null, null, null, null, null, subscriptionMonthly.getId(), null, null, null);
+        // We test first the next expected invoice for a specific subscription: We can see the targetDate is 2015-2-14 and not 2015-2-1
+        final DryRunArguments dryRunUpcomingInvoiceWithFilterArg1 = new TestDryRunArguments(DryRunType.UPCOMING_INVOICE, null, null, null, null, null, null, subscriptionMonthly1.getId(), null, null, null);
         expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("249.95")));
-        dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, dryRunUpcomingInvoiceWithFilterArg, callContext);
+        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+        dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, dryRunUpcomingInvoiceWithFilterArg1, callContext);
         assertEquals(dryRunInvoice.getTargetDate(), new LocalDate(2015, 2, 14));
         invoiceChecker.checkInvoiceNoAudits(dryRunInvoice, callContext, expectedInvoices);
         expectedInvoices.clear();
 
+
+        final DryRunArguments dryRunUpcomingInvoiceWithFilterArg2 = new TestDryRunArguments(DryRunType.UPCOMING_INVOICE, null, null, null, null, null, null, subscriptionMonthly2.getId(), null, null, null);
+        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("249.95")));
+        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
+        dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, dryRunUpcomingInvoiceWithFilterArg2, callContext);
+        assertEquals(dryRunInvoice.getTargetDate(), new LocalDate(2015, 2, 14));
+        invoiceChecker.checkInvoiceNoAudits(dryRunInvoice, callContext, expectedInvoices);
+        expectedInvoices.clear();
+
+
         // Then we test first the next expected invoice at the account level
         expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 1), new LocalDate(2016, 2, 1), InvoiceItemType.RECURRING, new BigDecimal("2399.95")));
         dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, DRY_RUN_UPCOMING_INVOICE_ARG, callContext);
@@ -206,6 +227,7 @@ public class TestIntegrationDryRunInvoice extends TestIntegrationBase {
         expectedInvoices.clear();
 
         expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("249.95")));
+        expectedInvoices.add(new ExpectedInvoiceItemCheck(new LocalDate(2015, 2, 14), new LocalDate(2015, 3, 14), InvoiceItemType.RECURRING, new BigDecimal("29.95")));
         dryRunInvoice = invoiceUserApi.triggerInvoiceGeneration(account.getId(), null, DRY_RUN_UPCOMING_INVOICE_ARG, callContext);
         invoiceChecker.checkInvoiceNoAudits(dryRunInvoice, callContext, expectedInvoices);
     }

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
index a755f9c..2ed5d58 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationInvoiceWithRepairLogic.java
@@ -41,7 +41,6 @@ import org.killbill.billing.catalog.api.ProductCategory;
 import org.killbill.billing.entitlement.api.DefaultEntitlement;
 import org.killbill.billing.entitlement.api.Entitlement.EntitlementActionPolicy;
 import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceItemType;
 import org.killbill.billing.invoice.api.InvoiceStatus;
@@ -51,7 +50,6 @@ import org.killbill.billing.invoice.dao.InvoiceModelDao;
 import org.killbill.billing.payment.api.Payment;
 import org.killbill.billing.payment.api.PluginProperty;
 import org.killbill.billing.payment.api.TransactionStatus;
-import org.testng.annotations.AfterMethod;
 import org.testng.annotations.Test;
 
 import com.google.common.collect.ImmutableList;
@@ -722,7 +720,7 @@ public class TestIntegrationInvoiceWithRepairLogic extends TestIntegrationBase {
         newItems.add(recurring3);
         newItems.add(repair3);
         shellInvoice.addInvoiceItems(newItems);
-        invoiceDao.createInvoice(shellInvoice, new FutureAccountNotifications(new HashMap<UUID, List<SubscriptionNotification>>()), internalCallContext);
+        invoiceDao.createInvoice(shellInvoice, new FutureAccountNotifications(), internalCallContext);
 
 
         // Move ahead one month, verify nothing from previous data was generated

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java b/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
index 7e56ce5..2f0562a 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/api/user/DefaultInvoiceUserApi.java
@@ -38,8 +38,6 @@ import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.invoice.InvoiceDispatcher;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.api.DryRunArguments;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
index bd35813..4c3745f 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/dao/DefaultInvoiceDao.java
@@ -26,6 +26,7 @@ import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
@@ -40,7 +41,6 @@ import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.entity.EntityPersistenceException;
 import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.api.DefaultInvoicePaymentErrorEvent;
 import org.killbill.billing.invoice.api.DefaultInvoicePaymentInfoEvent;
 import org.killbill.billing.invoice.api.Invoice;
@@ -285,7 +285,7 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
     @Override
     public List<InvoiceItemModelDao> createInvoices(final List<InvoiceModelDao> invoices,
                                                     final InternalCallContext context) {
-        return createInvoices(invoices, new FutureAccountNotifications(ImmutableMap.<UUID, List<SubscriptionNotification>>of()), context);
+        return createInvoices(invoices, new FutureAccountNotifications(), context);
     }
 
     private List<InvoiceItemModelDao> createInvoices(final Iterable<InvoiceModelDao> invoices,
@@ -982,20 +982,17 @@ public class DefaultInvoiceDao extends EntityDaoBase<InvoiceModelDao, Invoice, I
     private void notifyOfFutureBillingEvents(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
                                              final FutureAccountNotifications callbackDateTimePerSubscriptions, final InternalCallContext internalCallContext) {
 
+        for (final LocalDate notificationDate : callbackDateTimePerSubscriptions.getNotificationsForTrigger().keySet()) {
+            final DateTime notificationDateTime = internalCallContext.toUTCDateTime(notificationDate);
+            final Set<UUID> subscriptionIds = callbackDateTimePerSubscriptions.getNotificationsForTrigger().get(notificationDate);
+            nextBillingDatePoster.insertNextBillingNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionIds, notificationDateTime, internalCallContext);
+        }
+
         final long dryRunNotificationTime = invoiceConfig.getDryRunNotificationSchedule(internalCallContext).getMillis();
-        final boolean isInvoiceNotificationEnabled = dryRunNotificationTime > 0;
-        for (final UUID subscriptionId : callbackDateTimePerSubscriptions.getNotifications().keySet()) {
-            final List<SubscriptionNotification> callbackDateTimeUTC = callbackDateTimePerSubscriptions.getNotifications().get(subscriptionId);
-            for (final SubscriptionNotification cur : callbackDateTimeUTC) {
-                if (isInvoiceNotificationEnabled) {
-                    final DateTime curDryRunNotificationTime = cur.getEffectiveDate().minus(dryRunNotificationTime);
-                    final DateTime effectiveCurDryRunNotificationTime = (curDryRunNotificationTime.isAfter(clock.getUTCNow())) ? curDryRunNotificationTime : clock.getUTCNow();
-                    nextBillingDatePoster.insertNextBillingDryRunNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, effectiveCurDryRunNotificationTime, cur.getEffectiveDate(), internalCallContext);
-                }
-                if (cur.isForInvoiceNotificationTrigger()) {
-                    nextBillingDatePoster.insertNextBillingNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionId, cur.getEffectiveDate(), internalCallContext);
-                }
-            }
+        for (final LocalDate notificationDate : callbackDateTimePerSubscriptions.getNotificationsForDryRun().keySet()) {
+            final DateTime notificationDateTime = internalCallContext.toUTCDateTime(notificationDate);
+            final Set<UUID> subscriptionIds = callbackDateTimePerSubscriptions.getNotificationsForTrigger().get(notificationDate);
+            nextBillingDatePoster.insertNextBillingDryRunNotificationFromTransaction(entitySqlDaoWrapperFactory, accountId, subscriptionIds, notificationDateTime, notificationDateTime.plusMillis((int) dryRunNotificationTime), internalCallContext);
         }
     }
 

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
index 49d2288..63bd632 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/InvoiceDispatcher.java
@@ -24,6 +24,7 @@ import java.util.Collection;
 import java.util.Collections;
 import java.util.Comparator;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.LinkedList;
 import java.util.List;
@@ -53,7 +54,6 @@ import org.killbill.billing.entitlement.api.SubscriptionEventType;
 import org.killbill.billing.events.BusInternalEvent;
 import org.killbill.billing.events.EffectiveSubscriptionInternalEvent;
 import org.killbill.billing.events.InvoiceNotificationInternalEvent;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.api.DefaultInvoiceService;
 import org.killbill.billing.invoice.api.DryRunArguments;
 import org.killbill.billing.invoice.api.DryRunType;
@@ -115,6 +115,7 @@ import com.google.common.base.Preconditions;
 import com.google.common.base.Predicate;
 import com.google.common.collect.Collections2;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.common.collect.Lists;
 import com.google.common.collect.Ordering;
@@ -522,45 +523,70 @@ public class InvoiceDispatcher {
     }
 
     private FutureAccountNotifications createNextFutureNotificationDate(final InvoiceWithMetadata invoiceWithMetadata, final InternalCallContext context) {
-        final Map<UUID, List<SubscriptionNotification>> result = new HashMap<UUID, List<SubscriptionNotification>>();
 
-        for (final UUID subscriptionId : invoiceWithMetadata.getPerSubscriptionFutureNotificationDates().keySet()) {
+        final Map<LocalDate, Set<UUID>> notificationListForTrigger = new HashMap<LocalDate, Set<UUID>>();
 
-            final List<SubscriptionNotification> perSubscriptionNotifications = new ArrayList<SubscriptionNotification>();
+        for (final UUID subscriptionId : invoiceWithMetadata.getPerSubscriptionFutureNotificationDates().keySet()) {
 
             final SubscriptionFutureNotificationDates subscriptionFutureNotificationDates = invoiceWithMetadata.getPerSubscriptionFutureNotificationDates().get(subscriptionId);
-            // Add next recurring date if any
+
             if (subscriptionFutureNotificationDates.getNextRecurringDate() != null) {
-                perSubscriptionNotifications.add(new SubscriptionNotification(context.toUTCDateTime(subscriptionFutureNotificationDates.getNextRecurringDate()), true));
+                Set<UUID> subscriptionsForDates = notificationListForTrigger.get(subscriptionFutureNotificationDates.getNextRecurringDate());
+                if (subscriptionsForDates == null) {
+                    subscriptionsForDates = new HashSet<UUID>();
+                    notificationListForTrigger.put(subscriptionFutureNotificationDates.getNextRecurringDate(), subscriptionsForDates);
+                }
+                subscriptionsForDates.add(subscriptionId);
             }
-            // Add next usage dates if any
+
+
             if (subscriptionFutureNotificationDates.getNextUsageDates() != null) {
                 for (final UsageDef usageDef : subscriptionFutureNotificationDates.getNextUsageDates().keySet()) {
+
                     final LocalDate nextNotificationDateForUsage = subscriptionFutureNotificationDates.getNextUsageDates().get(usageDef);
-                    final DateTime subscriptionUsageCallbackDate = nextNotificationDateForUsage != null ? context.toUTCDateTime(nextNotificationDateForUsage) : null;
-                    perSubscriptionNotifications.add(new SubscriptionNotification(subscriptionUsageCallbackDate, true));
+                    Set<UUID> subscriptionsForDates = notificationListForTrigger.get(nextNotificationDateForUsage);
+                    if (subscriptionsForDates == null) {
+                        subscriptionsForDates = new HashSet<UUID>();
+                        notificationListForTrigger.put(nextNotificationDateForUsage, subscriptionsForDates);
+                    }
+                    subscriptionsForDates.add(subscriptionId);
                 }
             }
-            if (!perSubscriptionNotifications.isEmpty()) {
-                result.put(subscriptionId, perSubscriptionNotifications);
-            }
         }
 
-        // If dryRunNotification is enabled we also need to fetch the upcoming PHASE dates (we add SubscriptionNotification with isForInvoiceNotificationTrigger = false)
-        final boolean isInvoiceNotificationEnabled = invoiceConfig.getDryRunNotificationSchedule(context).getMillis() > 0;
+        final long dryRunNotificationTime = invoiceConfig.getDryRunNotificationSchedule(context).getMillis();
+        final boolean isInvoiceNotificationEnabled = dryRunNotificationTime > 0;
+
+        final Map<LocalDate, Set<UUID>> notificationListForDryRun  = isInvoiceNotificationEnabled ? new HashMap<LocalDate, Set<UUID>>() : ImmutableMap.<LocalDate, Set<UUID>>of();
         if (isInvoiceNotificationEnabled) {
-            final Map<UUID, DateTime> upcomingPhasesForSubscriptions = subscriptionApi.getNextFutureEventForSubscriptions(ImmutableList.<SubscriptionBaseTransitionType>of(SubscriptionBaseTransitionType.PHASE), context);
-            for (final UUID cur : upcomingPhasesForSubscriptions.keySet()) {
-                final DateTime curDate = upcomingPhasesForSubscriptions.get(cur);
-                List<SubscriptionNotification> resultValue = result.get(cur);
-                if (resultValue == null) {
-                    resultValue = new ArrayList<SubscriptionNotification>();
+            for (final LocalDate curDate : notificationListForTrigger.keySet()) {
+                final LocalDate curDryRunDate = context.toLocalDate(context.toUTCDateTime(curDate).minus(dryRunNotificationTime));
+                Set<UUID> subscriptionsForDryRunDates = notificationListForDryRun.get(curDryRunDate);
+                if (subscriptionsForDryRunDates == null) {
+                    subscriptionsForDryRunDates = new HashSet<UUID>();
+                    notificationListForDryRun.put(curDryRunDate, subscriptionsForDryRunDates);
                 }
-                resultValue.add(new SubscriptionNotification(curDate, false));
-                result.put(cur, resultValue);
+                subscriptionsForDryRunDates.addAll(notificationListForDryRun.get(curDate));
             }
+
+
+            final Map<UUID, DateTime> upcomingPhasesForSubscriptions = isInvoiceNotificationEnabled ?
+                                                                       subscriptionApi.getNextFutureEventForSubscriptions(ImmutableList.<SubscriptionBaseTransitionType>of(SubscriptionBaseTransitionType.PHASE), context) :
+                                                                       ImmutableMap.<UUID, DateTime>of();
+
+
+             for (UUID curId : upcomingPhasesForSubscriptions.keySet()) {
+                 final LocalDate curDryRunDate = context.toLocalDate(upcomingPhasesForSubscriptions.get(curId).minus(dryRunNotificationTime));
+                 Set<UUID> subscriptionsForDryRunDates = notificationListForDryRun.get(curDryRunDate);
+                 if (subscriptionsForDryRunDates == null) {
+                     subscriptionsForDryRunDates = new HashSet<UUID>();
+                     notificationListForDryRun.put(curDryRunDate, subscriptionsForDryRunDates);
+                 }
+                 subscriptionsForDryRunDates.add(curId);
+             }
         }
-        return new FutureAccountNotifications(result);
+
+        return new FutureAccountNotifications(notificationListForTrigger, notificationListForDryRun);
     }
 
     private List<InvoiceItemModelDao> transformToInvoiceModelDao(final List<InvoiceItem> invoiceItems) {
@@ -692,33 +718,25 @@ public class InvoiceDispatcher {
 
     public static class FutureAccountNotifications {
 
-        private final Map<UUID, List<SubscriptionNotification>> notifications;
 
-        public FutureAccountNotifications(final Map<UUID, List<SubscriptionNotification>> notifications) {
-            this.notifications = notifications;
-        }
+        private final Map<LocalDate, Set<UUID>> notificationListForTrigger;
+        private final Map<LocalDate, Set<UUID>> notificationListForDryRun;
 
-        public Map<UUID, List<SubscriptionNotification>> getNotifications() {
-            return notifications;
+        public FutureAccountNotifications() {
+            this(ImmutableMap.<LocalDate, Set<UUID>>of(), ImmutableMap.<LocalDate, Set<UUID>>of());
         }
 
-        public static class SubscriptionNotification {
-
-            private final DateTime effectiveDate;
-            private final boolean isForNotificationTrigger;
-
-            public SubscriptionNotification(final DateTime effectiveDate, final boolean isForNotificationTrigger) {
-                this.effectiveDate = effectiveDate;
-                this.isForNotificationTrigger = isForNotificationTrigger;
-            }
+        public FutureAccountNotifications(final Map<LocalDate, Set<UUID>> notificationListForTrigger, final Map<LocalDate, Set<UUID>> notificationListForDryRun) {
+            this.notificationListForTrigger = notificationListForTrigger;
+            this.notificationListForDryRun = notificationListForDryRun;
+        }
 
-            public DateTime getEffectiveDate() {
-                return effectiveDate;
-            }
+        public Map<LocalDate, Set<UUID>> getNotificationsForTrigger() {
+            return notificationListForTrigger;
+        }
 
-            public boolean isForInvoiceNotificationTrigger() {
-                return isForNotificationTrigger;
-            }
+        public Map<LocalDate, Set<UUID>> getNotificationsForDryRun() {
+            return notificationListForDryRun;
         }
     }
 

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java b/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
index c11d8dc..fb269c5 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/notification/DefaultNextBillingDatePoster.java
@@ -33,6 +33,7 @@ import org.killbill.notificationq.api.NotificationQueueService.NoSuchNotificatio
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.base.Joiner;
 import com.google.common.collect.ImmutableList;
 import com.google.inject.Inject;
 
@@ -40,6 +41,8 @@ public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
 
     private static final Logger log = LoggerFactory.getLogger(DefaultNextBillingDatePoster.class);
 
+    private static Joiner JOINER = Joiner.on(",");
+
     private final NotificationQueueService notificationQueueService;
 
     @Inject
@@ -50,24 +53,24 @@ public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
     @Override
     public void insertNextBillingNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
                                                              final UUID accountId,
-                                                             final UUID subscriptionId,
+                                                             final Iterable<UUID> subscriptionIds,
                                                              final DateTime futureNotificationTime,
                                                              final InternalCallContext internalCallContext) {
-        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionId, Boolean.FALSE, futureNotificationTime, futureNotificationTime, internalCallContext);
+        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionIds, Boolean.FALSE, futureNotificationTime, futureNotificationTime, internalCallContext);
     }
 
     @Override
     public void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
                                                                    final UUID accountId,
-                                                                   final UUID subscriptionId,
+                                                                   final Iterable<UUID> subscriptionIds,
                                                                    final DateTime futureNotificationTime,
                                                                    final DateTime targetDate,
                                                                    final InternalCallContext internalCallContext) {
-        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionId, Boolean.TRUE, futureNotificationTime, targetDate, internalCallContext);
+        insertNextBillingFromTransactionInternal(entitySqlDaoWrapperFactory, subscriptionIds, Boolean.TRUE, futureNotificationTime, targetDate, internalCallContext);
     }
 
     private void insertNextBillingFromTransactionInternal(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory,
-                                                          final UUID subscriptionId,
+                                                          final Iterable<UUID> subscriptionIds,
                                                           final Boolean isDryRunForInvoiceNotification,
                                                           final DateTime futureNotificationTime,
                                                           final DateTime targetDate,
@@ -97,22 +100,22 @@ public class DefaultNextBillingDatePoster implements NextBillingDatePoster {
             }
 
             if (existingNotificationForEffectiveDate == null) {
-                log.info("Queuing next billing date notification at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
+                log.info("Queuing next billing date notification at {} for subscriptionId {}", futureNotificationTime.toString(), JOINER.join(subscriptionIds));
 
-                final NextBillingDateNotificationKey newNotificationEvent = new NextBillingDateNotificationKey(null, ImmutableList.<UUID>of(subscriptionId), targetDate, isDryRunForInvoiceNotification);
+                final NextBillingDateNotificationKey newNotificationEvent = new NextBillingDateNotificationKey(null, subscriptionIds, targetDate, isDryRunForInvoiceNotification);
                 nextBillingQueue.recordFutureNotificationFromTransaction(entitySqlDaoWrapperFactory.getHandle().getConnection(), futureNotificationTime,
                                                                          newNotificationEvent, internalCallContext.getUserToken(),
                                                                          internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
             } else {
-                log.info("Updating next billing date notification event at {} for subscriptionId {}", futureNotificationTime.toString(), subscriptionId.toString());
-                final NextBillingDateNotificationKey updateNotificationEvent = new NextBillingDateNotificationKey(existingNotificationForEffectiveDate.getEvent(), ImmutableList.of(subscriptionId));
-                nextBillingQueue.updateFutureNotification(existingNotificationForEffectiveDate.getRecordId(), updateNotificationEvent, internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
+                log.info("Updating next billing date notification event at {} for subscriptionId {}", futureNotificationTime.toString(), JOINER.join(subscriptionIds));
+                final NextBillingDateNotificationKey updateNotificationEvent = new NextBillingDateNotificationKey(existingNotificationForEffectiveDate.getEvent(), subscriptionIds);
+                nextBillingQueue.updateFutureNotificationFromTransaction(entitySqlDaoWrapperFactory.getHandle().getConnection(), existingNotificationForEffectiveDate.getRecordId(), updateNotificationEvent, internalCallContext.getAccountRecordId(), internalCallContext.getTenantRecordId());
             }
 
         } catch (final NoSuchNotificationQueue e) {
             log.error("Attempting to put items on a non-existent queue (NextBillingDateNotifier).", e);
         } catch (final IOException e) {
-            log.error("Failed to serialize notificationKey for subscriptionId {}", subscriptionId);
+            log.error("Failed to serialize notificationKey for subscriptionId {}", subscriptionIds);
         }
     }
 

diff --git a/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java b/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
index 96836bd..d2dce63 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/notification/NextBillingDatePoster.java
@@ -27,8 +27,8 @@ import org.killbill.billing.util.entity.dao.EntitySqlDaoWrapperFactory;
 public interface NextBillingDatePoster {
 
     void insertNextBillingNotificationFromTransaction(EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, UUID accountId,
-                                                      UUID subscriptionId, DateTime futureNotificationTime, InternalCallContext internalCallContext);
+                                                      Iterable<UUID> subscriptionId, DateTime futureNotificationTime, InternalCallContext internalCallContext);
 
     void insertNextBillingDryRunNotificationFromTransaction(final EntitySqlDaoWrapperFactory entitySqlDaoWrapperFactory, final UUID accountId,
-                                                            final UUID subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext);
+                                                            final Iterable<UUID> subscriptionId, final DateTime futureNotificationTime, final DateTime targetDate, final InternalCallContext internalCallContext);
 }

diff --git a/invoice/src/test/java/org/killbill/billing/invoice/proRations/InvoiceTestUtils.java b/invoice/src/test/java/org/killbill/billing/invoice/proRations/InvoiceTestUtils.java
index fc5f88b..96d1868 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/proRations/InvoiceTestUtils.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/proRations/InvoiceTestUtils.java
@@ -28,8 +28,6 @@ import org.killbill.billing.account.api.AccountApiException;
 import org.killbill.billing.callcontext.InternalCallContext;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.entity.EntityPersistenceException;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications;
-import org.killbill.billing.invoice.InvoiceDispatcher.FutureAccountNotifications.SubscriptionNotification;
 import org.killbill.billing.invoice.TestInvoiceHelper;
 import org.killbill.billing.invoice.api.Invoice;
 import org.killbill.billing.invoice.api.InvoiceApiException;
@@ -47,7 +45,6 @@ import org.mockito.Mockito;
 import org.testng.Assert;
 
 import com.google.common.collect.ImmutableList;
-import com.google.common.collect.ImmutableMap;
 
 public class InvoiceTestUtils {