killbill-memoizeit

Introduce path optimization to not fetch all usage records

4/23/2015 4:40:04 PM

Details

diff --git a/api/src/main/java/org/killbill/billing/usage/InternalUserApi.java b/api/src/main/java/org/killbill/billing/usage/InternalUserApi.java
index 089a208..2dac91f 100644
--- a/api/src/main/java/org/killbill/billing/usage/InternalUserApi.java
+++ b/api/src/main/java/org/killbill/billing/usage/InternalUserApi.java
@@ -25,5 +25,5 @@ import org.killbill.billing.callcontext.InternalTenantContext;
 
 public interface InternalUserApi {
 
-    public List<RawUsage> getRawUsageForAccount(final UUID accountId, final LocalDate stateDate, final LocalDate endDate, final InternalTenantContext tenantContext);
+    public List<RawUsage> getRawUsageForAccount(final LocalDate stateDate, final LocalDate endDate, final InternalTenantContext tenantContext);
 }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java b/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
index c298495..61c3063 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/generator/DefaultInvoiceGenerator.java
@@ -20,9 +20,11 @@ package org.killbill.billing.invoice.generator;
 
 import java.math.BigDecimal;
 import java.util.ArrayList;
-import java.util.Collections;
+import java.util.HashMap;
 import java.util.Iterator;
+import java.util.LinkedList;
 import java.util.List;
+import java.util.Map;
 import java.util.UUID;
 
 import javax.annotation.Nullable;
@@ -50,13 +52,11 @@ import org.killbill.billing.invoice.model.InvalidDateSequenceException;
 import org.killbill.billing.invoice.model.RecurringInvoiceItem;
 import org.killbill.billing.invoice.model.RecurringInvoiceItemData;
 import org.killbill.billing.invoice.tree.AccountItemTree;
+import org.killbill.billing.invoice.usage.RawUsageOptimizer;
 import org.killbill.billing.invoice.usage.SubscriptionConsumableInArrear;
 import org.killbill.billing.junction.BillingEvent;
 import org.killbill.billing.junction.BillingEventSet;
-import org.killbill.billing.usage.InternalUserApi;
 import org.killbill.billing.usage.RawUsage;
-import org.killbill.billing.util.callcontext.InternalCallContextFactory;
-import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.billing.util.config.InvoiceConfig;
 import org.killbill.billing.util.currency.KillBillMoney;
 import org.killbill.clock.Clock;
@@ -66,6 +66,7 @@ import org.slf4j.LoggerFactory;
 import com.google.common.base.Function;
 import com.google.common.base.Predicate;
 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.inject.Inject;
@@ -76,13 +77,13 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
 
     private final Clock clock;
     private final InvoiceConfig config;
-    private final InternalUserApi usageApi;
+    private final RawUsageOptimizer rawUsageOptimizer;
 
     @Inject
-    public DefaultInvoiceGenerator(final Clock clock, final InternalUserApi usageApi, final InvoiceConfig config) {
+    public DefaultInvoiceGenerator(final Clock clock, final InvoiceConfig config, final RawUsageOptimizer rawUsageOptimizer) {
         this.clock = clock;
         this.config = config;
-        this.usageApi = usageApi;
+        this.rawUsageOptimizer = rawUsageOptimizer;
     }
 
     /*
@@ -106,17 +107,18 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
         final List<InvoiceItem> inAdvanceItems = generateInAdvanceInvoiceItems(account.getId(), invoiceId, events, existingInvoices, adjustedTargetDate, targetCurrency);
         invoice.addInvoiceItems(inAdvanceItems);
 
-        final List<InvoiceItem> usageItems = generateUsageInvoiceItems(account, invoiceId, events, existingInvoices, targetDate, context);
+        final List<InvoiceItem> usageItems = generateUsageConsumableInArrearItems(account, invoiceId, events, existingInvoices, targetDate, context);
         invoice.addInvoiceItems(usageItems);
 
         return invoice.getInvoiceItems().size() != 0 ? invoice : null;
     }
 
-    // STEPH_USAGE Only deals with consumable in arrear usage billing.
-    private List<InvoiceItem> generateUsageInvoiceItems(final Account account,
-                                                        final UUID invoiceId, final BillingEventSet eventSet,
-                                                        @Nullable final List<Invoice> existingInvoices, final LocalDate targetDate,
-                                                        final InternalCallContext internalCallContext) throws InvoiceApiException {
+    private List<InvoiceItem> generateUsageConsumableInArrearItems(final Account account,
+                                                                   final UUID invoiceId, final BillingEventSet eventSet,
+                                                                   @Nullable final List<Invoice> existingInvoices, final LocalDate targetDate,
+                                                                   final InternalCallContext internalCallContext) throws InvoiceApiException {
+
+        final Map<UUID, List<InvoiceItem>> perSubscriptionConsumableInArrearUsageItems = extractPerSubscriptionExistingConsumableInArrearUsageItems(eventSet.getUsages(), existingInvoices);
         try {
             final List<InvoiceItem> items = Lists.newArrayList();
             final Iterator<BillingEvent> events = eventSet.iterator();
@@ -143,7 +145,7 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
                         }
                     }).orNull() != null;
                     if (foundUsage) {
-                        rawUsage = usageApi.getRawUsageForAccount(account.getId(), new LocalDate(event.getEffectiveDate(), account.getTimeZone()), targetDate, internalCallContext);
+                        rawUsage = rawUsageOptimizer.getConsumableInArrearUsage(new LocalDate(event.getEffectiveDate(), account.getTimeZone()), targetDate, Iterables.concat(perSubscriptionConsumableInArrearUsageItems.values()), eventSet.getUsages(), internalCallContext);
                         seenAnyUsageItems = true;
                     }
                 }
@@ -152,7 +154,8 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
                 final UUID subscriptionId = event.getSubscription().getId();
                 if (curSubscriptionId != null && !curSubscriptionId.equals(subscriptionId)) {
                     final SubscriptionConsumableInArrear subscriptionConsumableInArrear = new SubscriptionConsumableInArrear(invoiceId, curEvents, rawUsage, targetDate);
-                    items.addAll(subscriptionConsumableInArrear.computeMissingUsageInvoiceItems(extractUsageItemsForSubscription(curSubscriptionId, existingInvoices)));
+                    final List<InvoiceItem> consumableInUsageArrearItems = perSubscriptionConsumableInArrearUsageItems.get(curSubscriptionId);
+                    items.addAll(subscriptionConsumableInArrear.computeMissingUsageInvoiceItems(consumableInUsageArrearItems != null ? consumableInUsageArrearItems : ImmutableList.<InvoiceItem>of()));
                     curEvents = Lists.newArrayList();
                 }
                 curSubscriptionId = subscriptionId;
@@ -161,7 +164,8 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
             if (curSubscriptionId != null) {
 
                 final SubscriptionConsumableInArrear subscriptionConsumableInArrear = new SubscriptionConsumableInArrear(invoiceId, curEvents, rawUsage, targetDate);
-                items.addAll(subscriptionConsumableInArrear.computeMissingUsageInvoiceItems(extractUsageItemsForSubscription(curSubscriptionId, existingInvoices)));
+                final List<InvoiceItem> consumableInUsageArrearItems = perSubscriptionConsumableInArrearUsageItems.get(curSubscriptionId);
+                items.addAll(subscriptionConsumableInArrear.computeMissingUsageInvoiceItems(consumableInUsageArrearItems != null ? consumableInUsageArrearItems : ImmutableList.<InvoiceItem>of()));
             }
             return items;
 
@@ -170,25 +174,39 @@ public class DefaultInvoiceGenerator implements InvoiceGenerator {
         }
     }
 
-    private List<InvoiceItem> extractUsageItemsForSubscription(final UUID subscriptionId, @Nullable final List<Invoice> existingInvoices) {
+    private Map<UUID, List<InvoiceItem>> extractPerSubscriptionExistingConsumableInArrearUsageItems(final Map<String, Usage> knownUsage, @Nullable final List<Invoice> existingInvoices) {
 
-        if (existingInvoices == null) {
-            return Collections.emptyList();
+        if (existingInvoices == null || existingInvoices.isEmpty()) {
+            return ImmutableMap.of();
         }
 
-        final Iterable usageItems = Iterables.concat(Iterables.transform(existingInvoices, new Function<Invoice, Iterable<InvoiceItem>>() {
+        final Map<UUID, List<InvoiceItem>> result = new HashMap<UUID, List<InvoiceItem>>();
+        final Iterable<InvoiceItem> usageConsumableInArrearItems = Iterables.concat(Iterables.transform(existingInvoices, new Function<Invoice, Iterable<InvoiceItem>>() {
             @Override
             public Iterable<InvoiceItem> apply(final Invoice input) {
 
                 return Iterables.filter(input.getInvoiceItems(), new Predicate<InvoiceItem>() {
                     @Override
                     public boolean apply(final InvoiceItem input) {
-                        return input.getInvoiceItemType() == InvoiceItemType.USAGE && input.getSubscriptionId().equals(subscriptionId);
+                        if (input.getInvoiceItemType() == InvoiceItemType.USAGE) {
+                            final Usage usage = knownUsage.get(input.getUsageName());
+                            return usage.getUsageType() == UsageType.CONSUMABLE && usage.getBillingMode() == BillingMode.IN_ARREAR;
+                        }
+                        return false;
                     }
                 });
             }
         }));
-        return ImmutableList.<InvoiceItem>copyOf(usageItems);
+
+        for (InvoiceItem cur : usageConsumableInArrearItems) {
+            List<InvoiceItem> perSubscriptionUsageItems = result.get(cur.getSubscriptionId());
+            if (perSubscriptionUsageItems == null) {
+                perSubscriptionUsageItems = new LinkedList<InvoiceItem>();
+                result.put(cur.getSubscriptionId(), perSubscriptionUsageItems);
+            }
+            perSubscriptionUsageItems.add(cur);
+        }
+        return result;
     }
 
     private List<InvoiceItem> generateInAdvanceInvoiceItems(final UUID accountId, final UUID invoiceId, final BillingEventSet eventSet,
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/glue/DefaultInvoiceModule.java b/invoice/src/main/java/org/killbill/billing/invoice/glue/DefaultInvoiceModule.java
index 5929709..f106da5 100644
--- a/invoice/src/main/java/org/killbill/billing/invoice/glue/DefaultInvoiceModule.java
+++ b/invoice/src/main/java/org/killbill/billing/invoice/glue/DefaultInvoiceModule.java
@@ -48,6 +48,7 @@ import org.killbill.billing.invoice.notification.NextBillingDatePoster;
 import org.killbill.billing.invoice.notification.NullInvoiceNotifier;
 import org.killbill.billing.invoice.plugin.api.InvoicePluginApi;
 import org.killbill.billing.invoice.template.bundles.DefaultResourceBundleFactory;
+import org.killbill.billing.invoice.usage.RawUsageOptimizer;
 import org.killbill.billing.osgi.api.OSGIServiceRegistration;
 import org.killbill.billing.platform.api.KillbillConfigSource;
 import org.killbill.billing.util.config.InvoiceConfig;
@@ -156,7 +157,7 @@ public class DefaultInvoiceModule extends KillBillModule implements InvoiceModul
         installInvoicePaymentApi();
         installInvoiceMigrationApi();
         installResourceBundleFactory();
-
+        bind(RawUsageOptimizer.class).asEagerSingleton();;
         bind(InvoiceApiHelper.class).asEagerSingleton();
     }
 }
diff --git a/invoice/src/main/java/org/killbill/billing/invoice/usage/RawUsageOptimizer.java b/invoice/src/main/java/org/killbill/billing/invoice/usage/RawUsageOptimizer.java
new file mode 100644
index 0000000..effa298
--- /dev/null
+++ b/invoice/src/main/java/org/killbill/billing/invoice/usage/RawUsageOptimizer.java
@@ -0,0 +1,141 @@
+/*
+ * Copyright 2014-2015 Groupon, Inc
+ * Copyright 2014-2015 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
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.invoice.usage;
+
+import java.util.List;
+import java.util.ListIterator;
+import java.util.Map;
+import java.util.Set;
+
+import javax.annotation.Nullable;
+import javax.inject.Inject;
+
+import org.joda.time.LocalDate;
+import org.killbill.billing.callcontext.InternalCallContext;
+import org.killbill.billing.catalog.api.BillingPeriod;
+import org.killbill.billing.catalog.api.Usage;
+import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.invoice.model.UsageInvoiceItem;
+import org.killbill.billing.usage.InternalUserApi;
+import org.killbill.billing.usage.RawUsage;
+import org.killbill.billing.util.config.InvoiceConfig;
+import org.killbill.clock.Clock;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Function;
+import com.google.common.base.Predicate;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
+import com.google.common.collect.Iterables;
+import com.google.common.collect.Ordering;
+
+public class RawUsageOptimizer {
+
+    private static final Ordering<InvoiceItem> USAGE_ITEM_ORDERING = Ordering.natural()
+                                                                             .onResultOf(new Function<InvoiceItem, Comparable>() {
+                                                                                 @Override
+                                                                                 public Comparable apply(final InvoiceItem invoiceItem) {
+                                                                                     return invoiceItem.getEndDate();
+                                                                                 }
+                                                                             });
+
+    private static final Logger log = LoggerFactory.getLogger(RawUsageOptimizer.class);
+
+    private final InternalUserApi usageApi;
+    private final InvoiceConfig config;
+
+    @Inject
+    public RawUsageOptimizer(final InvoiceConfig config, final InternalUserApi usageApi) {
+        this.usageApi = usageApi;
+        this.config = config;
+    }
+
+    public List<RawUsage> getConsumableInArrearUsage(final LocalDate firstEventStartDate, final LocalDate targetDate, final Iterable<InvoiceItem> existingUsageItems, final Map<String, Usage> knownUsage, final InternalCallContext internalCallContext) {
+        final LocalDate targetStartDate = config.getMaxRawUsagePreviousPeriod() > 0 ? getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, existingUsageItems, knownUsage) : firstEventStartDate;
+        return usageApi.getRawUsageForAccount(targetStartDate, targetDate, internalCallContext);
+    }
+
+    @VisibleForTesting
+    LocalDate getOptimizedRawUsageStartDate(final LocalDate firstEventStartDate, final LocalDate targetDate, final Iterable<InvoiceItem> existingUsageItems, final Map<String, Usage> knownUsage) {
+
+        if (!existingUsageItems.iterator().hasNext()) {
+            return firstEventStartDate;
+
+        }
+        // Extract all usage billing period known in that catalog
+        final Set<BillingPeriod> knownUsageBillingPeriod = ImmutableSet.copyOf(Iterables.transform(knownUsage.values(), new Function<Usage, BillingPeriod>() {
+            @Nullable
+            @Override
+            public BillingPeriod apply(final Usage input) {
+                return input.getBillingPeriod();
+            }
+        }));
+
+        // Make sure all usage items are sorted by endDate
+        final List<InvoiceItem> sortedUsageItems = USAGE_ITEM_ORDERING.sortedCopy(existingUsageItems);
+
+        // Compute an array with one date per BillingPeriod:
+        // If BillingPeriod is never defined in the catalog (no need to look for items), we initialize its value such that
+        final LocalDate[] perBillingPeriodMostRecentConsumableInArrearItemEndDate = new LocalDate[BillingPeriod.values().length];
+        int idx = 0;
+        for (BillingPeriod bp : BillingPeriod.values()) {
+            final LocalDate makerDateThanCannotBeChosenAsTheMinOfAllDates = targetDate.plusMonths(config.getMaxRawUsagePreviousPeriod() * bp.getNumberOfMonths());
+            perBillingPeriodMostRecentConsumableInArrearItemEndDate[idx++] = (knownUsageBillingPeriod.contains(bp)) ? null : makerDateThanCannotBeChosenAsTheMinOfAllDates;
+        }
+
+        final ListIterator<InvoiceItem> iterator = sortedUsageItems.listIterator(sortedUsageItems.size());
+        while (iterator.hasPrevious()) {
+
+            final UsageInvoiceItem item = (UsageInvoiceItem) iterator.previous();
+            final Usage usage = knownUsage.get(item.getUsageName());
+
+            if (perBillingPeriodMostRecentConsumableInArrearItemEndDate[usage.getBillingPeriod().ordinal()] == null) {
+                perBillingPeriodMostRecentConsumableInArrearItemEndDate[usage.getBillingPeriod().ordinal()] = item.getEndDate();
+                if (!Iterables.any(ImmutableList.copyOf(perBillingPeriodMostRecentConsumableInArrearItemEndDate), new Predicate<LocalDate>() {
+                    @Override
+                    public boolean apply(@Nullable final LocalDate input) {
+                        return input == null;
+                    }
+                })) {
+                    break;
+                }
+            }
+        }
+
+        // Extract the min from all the dates
+        LocalDate targetStartDate = null;
+        idx = 0;
+        for (BillingPeriod bp : BillingPeriod.values()) {
+            if (bp != BillingPeriod.NO_BILLING_PERIOD) {
+                final LocalDate tmp = perBillingPeriodMostRecentConsumableInArrearItemEndDate[idx];
+                final LocalDate targetBillingPeriodDate = tmp != null ? tmp.minusMonths(config.getMaxRawUsagePreviousPeriod() * bp.getNumberOfMonths()) : null;
+                if (targetStartDate == null || (targetBillingPeriodDate != null && targetBillingPeriodDate.compareTo(targetStartDate) < 0)) {
+                    targetStartDate = targetBillingPeriodDate;
+                }
+            }
+            idx++;
+        }
+
+        final LocalDate result = targetStartDate.compareTo(firstEventStartDate) > 0 ? targetStartDate : firstEventStartDate;
+        log.info("RawUsageOptimizer rawEventStartDate = {}, firstEventStartDate = {}", result, firstEventStartDate);
+        return result;
+    }
+
+}
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/generator/TestDefaultInvoiceGenerator.java b/invoice/src/test/java/org/killbill/billing/invoice/generator/TestDefaultInvoiceGenerator.java
index e83ad6f..9094ddc 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/generator/TestDefaultInvoiceGenerator.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/generator/TestDefaultInvoiceGenerator.java
@@ -122,8 +122,13 @@ public class TestDefaultInvoiceGenerator extends InvoiceTestSuiteNoDB {
             public TimeSpan getDryRunNotificationSchedule() {
                 return new TimeSpan("0s");
             }
+
+            @Override
+            public int getMaxRawUsagePreviousPeriod() {
+                return -1;
+            }
         };
-        this.generator = new DefaultInvoiceGenerator(clock, null, invoiceConfig);
+        this.generator = new DefaultInvoiceGenerator(clock, invoiceConfig, null);
         this.account = new MockAccountBuilder().name(UUID.randomUUID().toString().substring(1, 8))
                                                .firstNameLength(6)
                                                .email(UUID.randomUUID().toString().substring(1, 8))
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/InvoiceTestSuiteNoDB.java b/invoice/src/test/java/org/killbill/billing/invoice/InvoiceTestSuiteNoDB.java
index 3db8940..1069bfc 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/InvoiceTestSuiteNoDB.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/InvoiceTestSuiteNoDB.java
@@ -29,6 +29,7 @@ import org.killbill.billing.invoice.api.formatters.ResourceBundleFactory;
 import org.killbill.billing.invoice.dao.InvoiceDao;
 import org.killbill.billing.invoice.generator.InvoiceGenerator;
 import org.killbill.billing.invoice.glue.TestInvoiceModuleNoDB;
+import org.killbill.billing.invoice.usage.RawUsageOptimizer;
 import org.killbill.billing.junction.BillingInternalApi;
 import org.killbill.billing.lifecycle.api.BusService;
 import org.killbill.billing.platform.api.KillbillConfigSource;
@@ -94,6 +95,8 @@ public abstract class InvoiceTestSuiteNoDB extends GuicyKillbillTestSuiteNoDB {
     protected UsageUserApi usageUserApi;
     @Inject
     protected ResourceBundleFactory resourceBundleFactory;
+    @Inject
+    protected RawUsageOptimizer rawUsageOptimizer;
 
     @Override
     protected KillbillConfigSource getConfigSource() {
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
index 5961ded..0830548 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestContiguousIntervalConsumableInArrear.java
@@ -28,6 +28,7 @@ import org.joda.time.LocalDate;
 import org.killbill.billing.catalog.DefaultTier;
 import org.killbill.billing.catalog.DefaultTieredBlock;
 import org.killbill.billing.catalog.DefaultUsage;
+import org.killbill.billing.catalog.api.BillingPeriod;
 import org.killbill.billing.catalog.api.CatalogApiException;
 import org.killbill.billing.catalog.api.Currency;
 import org.killbill.billing.catalog.api.Usage;
@@ -71,7 +72,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTier(block);
-        final DefaultUsage usage = createDefaultUsage(usageName, tier);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
 
         final LocalDate targetDate = startDate.plusDays(1);
         final ContiguousIntervalConsumableInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
@@ -109,7 +110,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultTieredBlock block2 = createDefaultTieredBlock("unit", 1000, 100, BigDecimal.ONE);
         final DefaultTier tier2 = createDefaultTier(block2);
-        final DefaultUsage usage = createDefaultUsage(usageName, tier1, tier2);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier1, tier2);
 
         final LocalDate targetDate = new LocalDate(2014, 03, 20);
         final ContiguousIntervalConsumableInArrear intervalConsumableInArrear = createContiguousIntervalConsumableInArrear(usage, ImmutableList.<RawUsage>of(), targetDate, false,
@@ -139,7 +140,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
 
         final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier = createDefaultTier(block);
-        final DefaultUsage usage = createDefaultUsage(usageName, tier);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
 
         final LocalDate targetDate = endDate;
 
@@ -199,7 +200,7 @@ public class TestContiguousIntervalConsumableInArrear extends TestUsageInArrearB
         final DefaultTier tier = createDefaultTier(tieredBlock1, tieredBlock2);
 
 
-        final DefaultUsage usage = createDefaultUsage(usageName, tier);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
 
 
         final LocalDate t0 = new LocalDate(2015, 03, BCD);
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java
new file mode 100644
index 0000000..0b20533
--- /dev/null
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestRawUsageOptimizer.java
@@ -0,0 +1,150 @@
+/*
+ * Copyright 2014-2015 Groupon, Inc
+ * Copyright 2014-2015 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
+ * License.  You may obtain a copy of the License at:
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.killbill.billing.invoice.usage;
+
+import java.math.BigDecimal;
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+import org.joda.time.LocalDate;
+import org.killbill.billing.catalog.DefaultTier;
+import org.killbill.billing.catalog.DefaultTieredBlock;
+import org.killbill.billing.catalog.DefaultUsage;
+import org.killbill.billing.catalog.api.BillingPeriod;
+import org.killbill.billing.catalog.api.Currency;
+import org.killbill.billing.catalog.api.Usage;
+import org.killbill.billing.invoice.api.InvoiceItem;
+import org.killbill.billing.invoice.model.UsageInvoiceItem;
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+public class TestRawUsageOptimizer extends TestUsageInArrearBase {
+
+    @BeforeClass(groups = "fast")
+    protected void beforeClass() throws Exception {
+        super.beforeClass();
+    }
+
+    @BeforeMethod(groups = "fast")
+    public void beforeMethod() {
+        super.beforeMethod();
+    }
+
+
+    @Test(groups = "fast")
+    public void testWithNoItems() {
+
+        final LocalDate firstEventStartDate = new LocalDate(2014, 03, 15);
+
+        final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
+
+        final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
+        final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
+        final DefaultTier tier = createDefaultTier(block);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
+        knownUsage.put(usageName, usage);
+
+        final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, firstEventStartDate.plusDays(1), invoiceItems, knownUsage);
+        Assert.assertEquals(result.compareTo(firstEventStartDate), 0);
+    }
+
+
+    @Test(groups = "fast")
+    public void testWithOneMonthlyUsageSectionTooFewItems() {
+
+        final LocalDate firstEventStartDate = new LocalDate(2014, 03, 15);
+
+        final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
+        for (int i = 0; i < 1; i++) {
+            invoiceItems.add(createUsageItem(firstEventStartDate.plusMonths(i)));
+        }
+        final LocalDate targetDate = invoiceItems.get(invoiceItems.size() - 1).getEndDate();
+
+        final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
+        final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
+        final DefaultTier tier = createDefaultTier(block);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
+        knownUsage.put(usageName, usage);
+
+        final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage);
+        // The largest endDate for ii is 2014-04-15, and by default org.killbill.invoice.readMaxRawUsagePreviousPeriod == 2 => targetDate =>  2014-02-15,
+        // so we default to firstEventStartDate = 2014-03-15
+        Assert.assertEquals(result.compareTo(firstEventStartDate), 0);
+    }
+
+    @Test(groups = "fast")
+    public void testWithOneMonthlyUsageSectionAndEnoughUsageItems() {
+
+        final LocalDate firstEventStartDate = new LocalDate(2014, 03, 15);
+
+        final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
+        for (int i = 0; i < 5; i++) {
+            invoiceItems.add(createUsageItem(firstEventStartDate.plusMonths(i)));
+        }
+        final LocalDate targetDate = invoiceItems.get(invoiceItems.size() - 1).getEndDate();
+
+        final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
+        final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
+        final DefaultTier tier = createDefaultTier(block);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
+        knownUsage.put(usageName, usage);
+
+        final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage);
+        // The largest endDate for ii is 2014-08-15, and by default org.killbill.invoice.readMaxRawUsagePreviousPeriod == 2 => targetDate =>  2014-06-15
+        Assert.assertEquals(result.compareTo(new LocalDate(2014, 06, 15)), 0, "112 got " + result);
+    }
+
+
+    @Test(groups = "fast")
+    public void testWithOneMonthlyAndOneNonActiveAnnualUsageSectionAndEnoughUsageItems() {
+
+        final LocalDate firstEventStartDate = new LocalDate(2014, 03, 15);
+
+        final List<InvoiceItem> invoiceItems = new ArrayList<InvoiceItem>();
+        for (int i = 0; i < 5; i++) {
+            invoiceItems.add(createUsageItem(firstEventStartDate.plusMonths(i)));
+        }
+        final LocalDate targetDate = invoiceItems.get(invoiceItems.size() - 1).getEndDate();
+
+        final Map<String, Usage> knownUsage = new HashMap<String, Usage>();
+        final DefaultTieredBlock block = createDefaultTieredBlock("unit", 100, 1000, BigDecimal.ONE);
+        final DefaultTier tier = createDefaultTier(block);
+        final DefaultUsage usage = createDefaultUsage(usageName, BillingPeriod.MONTHLY, tier);
+        knownUsage.put(usageName, usage);
+
+        final DefaultTieredBlock block2 = createDefaultTieredBlock("unit2", 10, 10000, BigDecimal.TEN);
+        final DefaultTier tier2 = createDefaultTier(block2);
+        final DefaultUsage usage2 = createDefaultUsage("usageName2", BillingPeriod.ANNUAL, tier2);
+        knownUsage.put("usageName2", usage2);
+
+
+        final LocalDate result = rawUsageOptimizer.getOptimizedRawUsageStartDate(firstEventStartDate, targetDate, invoiceItems, knownUsage);
+        // The same reasoning applies as previously because there is no usage items against the annual and
+        // so, the largest endDate for ii is 2014-08-15, and by default org.killbill.invoice.readMaxRawUsagePreviousPeriod == 2 => targetDate =>  2014-06-15
+        Assert.assertEquals(result.compareTo(new LocalDate(2014, 06, 15)), 0, "142 got " + result);
+    }
+
+
+    private InvoiceItem createUsageItem(final LocalDate startDate) {
+        return new UsageInvoiceItem(invoiceId, accountId, bundleId, subscriptionId, planName, phaseName, usageName, startDate, startDate.plusMonths(1), BigDecimal.TEN, Currency.USD);
+    }
+}
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
index 90c20d6..f7e6870 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestSubscriptionConsumableInArrear.java
@@ -24,6 +24,7 @@ import org.joda.time.DateTimeZone;
 import org.joda.time.LocalDate;
 import org.killbill.billing.catalog.DefaultTier;
 import org.killbill.billing.catalog.DefaultTieredBlock;
+import org.killbill.billing.catalog.api.BillingPeriod;
 import org.killbill.billing.catalog.api.Usage;
 import org.killbill.billing.junction.BillingEvent;
 import org.killbill.billing.usage.RawUsage;
@@ -51,12 +52,12 @@ public class TestSubscriptionConsumableInArrear extends TestUsageInArrearBase {
         final String usageName1 = "erw";
         final DefaultTieredBlock block1 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier1 = createDefaultTier(block1);
-        final Usage usage1 = createDefaultUsage(usageName1, tier1);
+        final Usage usage1 = createDefaultUsage(usageName1, BillingPeriod.MONTHLY, tier1);
 
         final String usageName2 = "hghg";
         final DefaultTieredBlock block2 = createDefaultTieredBlock("unit", 100, 10, BigDecimal.ONE);
         final DefaultTier tier2 = createDefaultTier(block2);
-        final Usage usage2 = createDefaultUsage(usageName2, tier2);
+        final Usage usage2 = createDefaultUsage(usageName2, BillingPeriod.MONTHLY, tier2);
 
         final DateTime dt1 = new DateTime(2013, 3, 23, 4, 34, 59, DateTimeZone.UTC);
         final BillingEvent evt1 = createMockBillingEvent(dt1, ImmutableList.<Usage>builder().add(usage1).add(usage2).build());
diff --git a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
index f817cfb..412e96d 100644
--- a/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
+++ b/invoice/src/test/java/org/killbill/billing/invoice/usage/TestUsageInArrearBase.java
@@ -41,7 +41,6 @@ import org.killbill.billing.invoice.InvoiceTestSuiteNoDB;
 import org.killbill.billing.junction.BillingEvent;
 import org.killbill.billing.subscription.api.SubscriptionBase;
 import org.killbill.billing.usage.RawUsage;
-import org.killbill.billing.usage.api.UsageUserApi;
 import org.mockito.Mockito;
 import org.testng.annotations.BeforeClass;
 
@@ -57,8 +56,6 @@ public abstract class TestUsageInArrearBase extends InvoiceTestSuiteNoDB {
     protected Currency currency;
     protected String usageName;
 
-    protected UsageUserApi mockUsageUserApi;
-
     @BeforeClass(groups = "fast")
     protected void beforeClass() throws Exception {
         super.beforeClass();
@@ -82,7 +79,7 @@ public abstract class TestUsageInArrearBase extends InvoiceTestSuiteNoDB {
         return intervalConsumableInArrear;
     }
 
-    protected DefaultUsage createDefaultUsage(final String usageName, final DefaultTier... tiers) {
+    protected DefaultUsage createDefaultUsage(final String usageName, final BillingPeriod billingPeriod, final DefaultTier... tiers) {
         final DefaultUsage usage = new DefaultUsage();
         usage.setName(usageName);
         usage.setBillingMode(BillingMode.IN_ARREAR);
diff --git a/usage/src/main/java/org/killbill/billing/usage/api/svcs/DefaultInternalUserApi.java b/usage/src/main/java/org/killbill/billing/usage/api/svcs/DefaultInternalUserApi.java
index fc20656..e2de011 100644
--- a/usage/src/main/java/org/killbill/billing/usage/api/svcs/DefaultInternalUserApi.java
+++ b/usage/src/main/java/org/killbill/billing/usage/api/svcs/DefaultInternalUserApi.java
@@ -18,7 +18,6 @@
 package org.killbill.billing.usage.api.svcs;
 
 import java.util.List;
-import java.util.UUID;
 
 import javax.annotation.Nullable;
 import javax.inject.Inject;
@@ -44,7 +43,7 @@ public class DefaultInternalUserApi implements InternalUserApi {
     }
 
     @Override
-    public List<RawUsage> getRawUsageForAccount(final UUID accountId, final LocalDate stateDate, final LocalDate endDate, final InternalTenantContext internalTenantContext) {
+    public List<RawUsage> getRawUsageForAccount(final LocalDate stateDate, final LocalDate endDate, final InternalTenantContext internalTenantContext) {
         final List<RolledUpUsageModelDao> usage = rolledUpUsageDao.getRawUsageForAccount(stateDate, endDate, internalTenantContext);
         return ImmutableList.copyOf(Iterables.transform(usage, new Function<RolledUpUsageModelDao, RawUsage>() {
             @Nullable
diff --git a/util/src/main/java/org/killbill/billing/util/config/InvoiceConfig.java b/util/src/main/java/org/killbill/billing/util/config/InvoiceConfig.java
index de5b308..8143de4 100644
--- a/util/src/main/java/org/killbill/billing/util/config/InvoiceConfig.java
+++ b/util/src/main/java/org/killbill/billing/util/config/InvoiceConfig.java
@@ -37,4 +37,11 @@ public interface InvoiceConfig extends KillbillConfig {
     @Default("0s")
     @Description("DryRun invoice notification time before targetDate (ignored if set to 0s)")
     public TimeSpan getDryRunNotificationSchedule();
+
+
+    @Config("org.killbill.invoice.readMaxRawUsagePreviousPeriod")
+    @Default("2")
+    @Description("Maximum number of billingPeriod we read when retrieve raw usage data")
+    public int getMaxRawUsagePreviousPeriod();
+
 }