Add in expiring budget

chrome/browser/budget_service/budget_database.cc

Index: chrome/browser/budget_service/budget_database.cc
diff --git a/chrome/browser/budget_service/budget_database.cc b/chrome/browser/budget_service/budget_database.cc
index 587ed448b563cdc2b89dbd1b40f51b1c9bd90545..bfb672059591ec8859d5fb55978303254c3fad17 100644
--- a/chrome/browser/budget_service/budget_database.cc
+++ b/chrome/browser/budget_service/budget_database.cc
@@ -47,7 +47,7 @@ void BudgetDatabase::GetBudgetDetails(
   DCHECK_EQ(origin.GetOrigin(), origin);
 
   // If this origin is already in the cache, immediately return the data.
-  if (budget_map_.find(origin.spec()) != budget_map_.end()) {
+  if (IsCached(origin)) {
     BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
                             base::Bind(&BudgetDatabase::DidGetBudget,
                                        weak_ptr_factory_.GetWeakPtr(), origin,
@@ -69,56 +69,24 @@ void BudgetDatabase::AddBudget(const GURL& origin,
                                const StoreBudgetCallback& callback) {
   DCHECK_EQ(origin.GetOrigin(), origin);
 
-  // Look up the origin in our cache. Adding budget without first querying the
-  // existing budget is not suported.
-  DCHECK_GT(budget_map_.count(origin.spec()), 0U);
-
+  // Add a new chunk of budget for the origin at the default expiration time.
   base::Time expiration =
       clock_->Now() + base::TimeDelta::FromHours(kBudgetDurationInHours);
   budget_map_[origin.spec()].second.push_back(
-      std::make_pair(amount, expiration.ToInternalValue()));
+      std::make_pair(amount, expiration));
 
   // Now that the cache is updated, write the data to the database.
   WriteCachedValuesToDatabase(origin, callback);
 }
 
-void BudgetDatabase::WriteCachedValuesToDatabase(
-    const GURL& origin,
-    const StoreBudgetCallback& callback) {
-  // Create the data structures that are passed to the ProtoDatabase.
-  std::unique_ptr<
-      leveldb_proto::ProtoDatabase<budget_service::Budget>::KeyEntryVector>
-      entries(new leveldb_proto::ProtoDatabase<
-              budget_service::Budget>::KeyEntryVector());
-  std::unique_ptr<std::vector<std::string>> keys_to_remove(
-      new std::vector<std::string>());
-
-  // Each operation can either update the existing budget or remove the origin's
-  // budget information.
-  if (budget_map_.find(origin.spec()) == budget_map_.end()) {
-    // If the origin doesn't exist in the cache, this is a remove operation.
-    keys_to_remove->push_back(origin.spec());
-  } else {
-    // Build the Budget proto object.
-    budget_service::Budget budget;
-    const BudgetInfo& info = budget_map_[origin.spec()];
-    budget.set_debt(info.first);
-    for (const auto& chunk : info.second) {
-      budget_service::BudgetChunk* budget_chunk = budget.add_budget();
-      budget_chunk->set_amount(chunk.first);
-      budget_chunk->set_expiration(chunk.second);
-    }
-    entries->push_back(std::make_pair(origin.spec(), budget));
-  }
-
-  // Send the updates to the database.
-  db_->UpdateEntries(std::move(entries), std::move(keys_to_remove), callback);
-}
-
 void BudgetDatabase::OnDatabaseInit(bool success) {
   // TODO(harkness): Consider caching the budget database now?
 }
 
+bool BudgetDatabase::IsCached(const GURL& origin) const {
+  return budget_map_.find(origin.spec()) != budget_map_.end();
+}
+
 void BudgetDatabase::AddToCache(
     const GURL& origin,
     const AddToCacheCallback& callback,
@@ -134,7 +102,8 @@ void BudgetDatabase::AddToCache(
   // format which is better for removing things from the list.
   BudgetChunks chunks;
   for (const auto& chunk : budget_proto->budget())
-    chunks.push_back(std::make_pair(chunk.amount(), chunk.expiration()));
+    chunks.push_back(std::make_pair(
+        chunk.amount(), base::Time::FromInternalValue(chunk.expiration())));
 
   DCHECK(budget_proto->has_debt());
   budget_map_[origin.spec()] =
@@ -153,28 +122,88 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
     return;
   }
 
-  // Otherwise, build up the BudgetExpection. This is different from the format
+  // First, cleanup any expired budget chunks for the origin.
+  CleanupExpiredBudget(origin);
+
+  // Now, build up the BudgetExpection. This is different from the format
   // in which the cache stores the data. The cache stores chunks of budget and
   // when that budget expires. The BudgetExpectation describes a set of times
   // and the budget at those times.
-  const BudgetInfo& info = budget_map_[origin.spec()];
   BudgetExpectation expectation;
   double total = 0;
+  double debt = 0;
 
-  // Starting with the chunks that expire the farthest in the future, build up
-  // the budget expectations for those future times.
-  for (const auto& chunk : base::Reversed(info.second)) {
-    expectation.push_front(std::make_pair(total, chunk.second));
-    total += chunk.first;
+  if (IsCached(origin)) {
+    // Starting with the chunks that expire the farthest in the future, build up
+    // the budget expectations for those future times.
+    const BudgetInfo& info = budget_map_[origin.spec()];
+    for (const auto& chunk : base::Reversed(info.second)) {
+      expectation.push_front(std::make_pair(total, chunk.second));
+      total += chunk.first;
+    }
+    debt = info.first;
   }
 
-  // Always add one entry at the front of the list for the total budget right
-  // now.
-  expectation.push_front(std::make_pair(total, 0));
+  // Always add one entry at the front of the list for the total budget now.
+  expectation.push_front(std::make_pair(total, clock_->Now()));
 
-  callback.Run(true /* success */, info.first, expectation);
+  callback.Run(true /* success */, debt, expectation);
 }
 
 void BudgetDatabase::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
   clock_ = std::move(clock);
 }
+
+void BudgetDatabase::WriteCachedValuesToDatabase(
+    const GURL& origin,
+    const StoreBudgetCallback& callback) {
+  // First, cleanup any expired budget chunks for the origin.
+  CleanupExpiredBudget(origin);
+
+  // Create the data structures that are passed to the ProtoDatabase.
+  std::unique_ptr<
+      leveldb_proto::ProtoDatabase<budget_service::Budget>::KeyEntryVector>
+      entries(new leveldb_proto::ProtoDatabase<
+              budget_service::Budget>::KeyEntryVector());
+  std::unique_ptr<std::vector<std::string>> keys_to_remove(
+      new std::vector<std::string>());
+
+  // Each operation can either update the existing budget or remove the origin's
+  // budget information.
+  if (IsCached(origin)) {
+    // Build the Budget proto object.
+    budget_service::Budget budget;
+    const BudgetInfo& info = budget_map_[origin.spec()];
+    budget.set_debt(info.first);
+    for (const auto& chunk : info.second) {
+      budget_service::BudgetChunk* budget_chunk = budget.add_budget();
+      budget_chunk->set_amount(chunk.first);
+      budget_chunk->set_expiration(chunk.second.ToInternalValue());
+    }
+    entries->push_back(std::make_pair(origin.spec(), budget));
+  } else {
+    // If the origin doesn't exist in the cache, this is a remove operation.
+    keys_to_remove->push_back(origin.spec());
+  }
+
+  // Send the updates to the database.
+  db_->UpdateEntries(std::move(entries), std::move(keys_to_remove), callback);
+}
+
+void BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
+  if (!IsCached(origin))
+    return;
+
+  base::Time now = clock_->Now();
+
+  BudgetChunks& chunks = budget_map_[origin.spec()].second;
+  auto cleanup_iter = chunks.begin();
+
+  // This relies on the list of chunks being in timestamp order.
+  while (cleanup_iter != chunks.end() && cleanup_iter->second <= now)
+    cleanup_iter = chunks.erase(cleanup_iter);
+
+  // If the entire budget is empty now, cleanup the map.
+  if (chunks.empty())
+    budget_map_.erase(origin.spec());
+}