AddEngagementBudget and SpendBudget added to BudgetDatabase.

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 d7180115ecef043851fef01f2fcf5fe8febf4057..6d9186c5e87a2eb642a34c6c9722f0e3263b3d87 100644
--- a/chrome/browser/budget_service/budget_database.cc
+++ b/chrome/browser/budget_service/budget_database.cc
@@ -27,6 +27,15 @@ constexpr double kBudgetDurationInHours = 72;
 
 }  // namespace
 
+BudgetDatabase::BudgetInfo::BudgetInfo() {}
+
+BudgetDatabase::BudgetInfo::BudgetInfo(const BudgetInfo&& other)
+    : last_engagement_award(other.last_engagement_award) {
+  chunks = std::move(other.chunks);
+}
+
+BudgetDatabase::BudgetInfo::~BudgetInfo() {}
+
 BudgetDatabase::BudgetDatabase(
     const base::FilePath& database_dir,
     const scoped_refptr<base::SequencedTaskRunner>& task_runner)
@@ -72,9 +81,91 @@ void BudgetDatabase::AddBudget(const GURL& origin,
   // 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()].emplace_back(amount, expiration);
+  budget_map_[origin.spec()].chunks.emplace_back(amount, expiration);
+
+  // Now that the cache is updated, write the data to the database.
+  WriteCachedValuesToDatabase(origin, callback);
+}
+
+void BudgetDatabase::AddEngagementBudget(const GURL& origin,
+                                         double sesScore,

[Peter Beverloo, 2016/08/22 18:04:27]

nit: s/sesScore/score/

[harkness, 2016/08/23 09:55:55]

Done.

+                                         const StoreBudgetCallback& callback) {
+  DCHECK_EQ(origin.GetOrigin(), origin);
+
+  // By default we award the "full" award. Then that ratio is decreased if
+  // there have been other awards recently.
+  double awardRatio = 1.0;

[Peter Beverloo, 2016/08/22 18:04:27]

nit: s/awardRatio/ratio/

[harkness, 2016/08/23 09:55:55]

Done.

+
+  // Calculate how much budget should be awarded. If the origin is not cached,
+  // then we award a full amount.
+  if (IsCached(origin)) {
+    base::TimeDelta elapsed =
+        clock_->Now() - budget_map_[origin.spec()].last_engagement_award;
+    int elapsedHours = elapsed.InHours();

[Peter Beverloo, 2016/08/22 18:04:27]

nit: s/elapsedHours/elapsed_hours/

It'd be amazing if clang-format could tell you about this.

[harkness, 2016/08/23 09:55:55]

I'm not sure whether you or I would appreciate that more at this point!

+    if (elapsedHours == 0) {
+      // Don't give engagement awards for periods less than an hour.
+      callback.Run(true);
+      return;
+    }
+    if (elapsedHours < kBudgetDurationInHours)
+      awardRatio = elapsedHours / kBudgetDurationInHours;
+  }
+
+  // Update the last_engagement_award to the current time. If the origin wasn't
+  // already in the map, this adds a new entry for it.
+  budget_map_[origin.spec()].last_engagement_award = clock_->Now();
+
+  // Pass to the base AddBudget to update the cache and write to the database.
+  AddBudget(origin, sesScore * awardRatio, callback);
+}
+
+void BudgetDatabase::SpendBudget(const GURL& origin,
+                                 double amount,
+                                 const StoreBudgetCallback& callback) {
+  DCHECK_EQ(origin.GetOrigin(), origin);
+
+  // First, cleanup any expired budget chunks for the origin.
+  CleanupExpiredBudget(origin);
+
+  if (!IsCached(origin)) {
+    callback.Run(false);
+    return;
+  }
+
+  // Walk the list of budget chunks to see if the origin has enough budget.
+  double total = 0;
+  auto& info = budget_map_[origin.spec()];

[johnme, 2016/08/22 14:50:31]

consistency.

https://google.github.io/styleguide/cppguide.html#auto says (paraphrasing): Use
auto to avoid long/cluttery type names, especially when they involve templates
or namespaces or ::iterator, and only for local variables. But continue to use
manifest type declarations when it helps readability, for example if the type is
not otherwise spelled out nearby.

`auto& info` here seems clearly less legible than `BudgetInfo& info` since
otherwise you have to lookup the declaration of the map to know what type info
is. Ditto with the `const auto& chunk` below. I checked with Peter and he has no
strong opinion on these cases. However conversely, we'd both suggest using `auto
chunkIter` instead of `BudgetChunks::iterator` below ;-)

[harkness, 2016/08/23 09:55:56]

Done.

+  for (const auto& chunk : info.chunks)
+    total += chunk.amount;
+
+  if (total < amount) {
+    callback.Run(false);
+    return;
+  }
+
+  // Walk the chunks and remove enough budget to cover the needed amount.
+  double bill = amount;
+  BudgetChunks::iterator chunkIter;
+  for (chunkIter = info.chunks.begin(); chunkIter != info.chunks.end();
+       chunkIter++) {

[johnme, 2016/08/22 14:50:31]

Always ++it not it++ with iterators to avoid copying them.

[harkness, 2016/08/23 09:55:55]

Done.

+    if (chunkIter->amount > bill) {
+      chunkIter->amount -= bill;
+      bill = 0;
+      break;
+    }
+    bill -= chunkIter->amount;
+  }
+  info.chunks.erase(info.chunks.begin(), chunkIter);

[johnme, 2016/08/22 14:50:31]

Now that you're using a std::list, and since you kept iterators rather than
range-based for, I'd lean towards erasing as you go rather than than all at the
end, since otherwise someone reading the code has to think carefully about the
possible final values of chunkIter at the end of the for loop. Then you can also
inline the declaration of chunkIter into the for loop (and rename it to `it` or
`iter` so it fits on one line).

[harkness, 2016/08/23 09:55:56]

Done.

+
+  // There should have been enough budget to cover the entire bill.
+  DCHECK_EQ(0, bill);
 
   // Now that the cache is updated, write the data to the database.
+  // TODO(harkness): Consider adding a second parameter to the callback so the
+  // caller can distinguish between not enough budget and a failed database
+  // write.
+  // TODO(harkness): If the database write fails, the cache will be out of sync
+  // with the database. Consider ways to mitigate this.
   WriteCachedValuesToDatabase(origin, callback);
 }
 
@@ -99,13 +190,15 @@ void BudgetDatabase::AddToCache(
 
   // Add the data to the cache, converting from the proto format to an STL
   // format which is better for removing things from the list.
-  BudgetChunks chunks;
+  BudgetInfo& info = budget_map_[origin.spec()];
+  BudgetChunks& chunks = info.chunks;

[johnme, 2016/08/22 14:50:31]

Micro-nit: you only use this variable once - inline it?

[harkness, 2016/08/23 09:55:55]

Done.

   for (const auto& chunk : budget_proto->budget()) {
     chunks.emplace_back(chunk.amount(),
                         base::Time::FromInternalValue(chunk.expiration()));
   }
 
-  budget_map_[origin.spec()] = std::move(chunks);
+  info.last_engagement_award =
+      base::Time::FromInternalValue(budget_proto->engagement_last_updated());
 
   callback.Run(success);
 }
@@ -114,9 +207,9 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
                                   const GetBudgetDetailsCallback& callback,
                                   bool success) {
   // If the database wasn't able to read the information, return the
-  // failure and an empty BudgetExpectation.
+  // failure and an empty BudgetPrediction.
   if (!success) {
-    callback.Run(success, BudgetExpectation());
+    callback.Run(success, BudgetPrediction());
     return;
   }
 
@@ -125,25 +218,25 @@ void BudgetDatabase::DidGetBudget(const GURL& 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
+  // when that budget expires. The BudgetPrediction describes a set of times
   // and the budget at those times.
-  BudgetExpectation expectation;
+  BudgetPrediction prediction;
   double total = 0;
 
   if (IsCached(origin)) {
     // Starting with the chunks that expire the farthest in the future, build up
-    // the budget expectations for those future times.
-    const BudgetChunks& chunks = budget_map_[origin.spec()];
+    // the budget predictions for those future times.
+    const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
     for (const auto& chunk : base::Reversed(chunks)) {
-      expectation.emplace_front(total, chunk.expiration);
+      prediction.emplace_front(total, chunk.expiration);
       total += chunk.amount;
     }
   }
 
   // Always add one entry at the front of the list for the total budget now.
-  expectation.emplace_front(total, clock_->Now());
+  prediction.emplace_front(total, clock_->Now());
 
-  callback.Run(true /* success */, expectation);
+  callback.Run(true /* success */, prediction);
 }
 
 void BudgetDatabase::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
@@ -169,12 +262,14 @@ void BudgetDatabase::WriteCachedValuesToDatabase(
   if (IsCached(origin)) {
     // Build the Budget proto object.
     budget_service::Budget budget;
-    const BudgetChunks& chunks = budget_map_[origin.spec()];
-    for (const auto& chunk : chunks) {
+    const BudgetInfo& info = budget_map_[origin.spec()];
+    for (const auto& chunk : info.chunks) {
       budget_service::BudgetChunk* budget_chunk = budget.add_budget();
       budget_chunk->set_amount(chunk.amount);
       budget_chunk->set_expiration(chunk.expiration.ToInternalValue());
     }
+    budget.set_engagement_last_updated(
+        info.last_engagement_award.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.
@@ -191,14 +286,17 @@ void BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
 
   base::Time now = clock_->Now();
 
-  BudgetChunks& chunks = budget_map_[origin.spec()];
+  BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
   auto cleanup_iter = chunks.begin();
 
   // This relies on the list of chunks being in timestamp order.
   while (cleanup_iter != chunks.end() && cleanup_iter->expiration <= now)
     cleanup_iter = chunks.erase(cleanup_iter);
 
-  // If the entire budget is empty now, cleanup the map.
-  if (chunks.empty())
+  // If the entire budget is empty now AND there have been no engagements
+  // in the last kBudgetDurationInHours hours, remove this from the cache.
+  if (chunks.empty() &&
+      budget_map_[origin.spec()].last_engagement_award <
+          clock_->Now() - base::TimeDelta::FromHours(kBudgetDurationInHours))
     budget_map_.erase(origin.spec());
 }