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 aaada9642e98d3e5ef40b38914a01cad269269c7..7bb9e548ee87c1978f94ef7f07caa7e501b6ee5c 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);

[johnme, 2016/08/19 18:48:35]

Doesn't make sense to std::move in a copy constructor, since that leaves
other.chunks in an invalid state. Did you mean for this to be a move constructor
(const BudgetInfo&& other) instead? If so, you probably also want to use the
DISALLOW_COPY_AND_ASSIGN macro to mark it move-only.

[harkness, 2016/08/22 13:46:25]

Done.

+}
+
+BudgetDatabase::BudgetInfo::~BudgetInfo() {}
+
 BudgetDatabase::BudgetDatabase(
     const base::FilePath& database_dir,
     const scoped_refptr<base::SequencedTaskRunner>& task_runner)
@@ -72,9 +81,82 @@ 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,

[johnme, 2016/08/19 18:48:35]

Nit: parameter name is different in .h file; it's slightly clearer if they're
the same.

[harkness, 2016/08/22 13:46:25]

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;
+
+  // Calculate how much budget should be awarded. If the origin is not cached,
+  // then we award a full amount.
+  if (IsCached(origin)) {

[johnme, 2016/08/19 18:48:35]

Nit: each call to AddEngagementBudget causes origin to searched for in
budget_map_ 8 times (including string hashing each time). I can't see obvious
ways of optimising that without decreasing legibility, since it's spread across
various methods, but perhaps something to be aware of in future.

[harkness, 2016/08/22 13:46:26]

I'll definitely keep it in mind, and see if there are some ways to optimize in
the future. I have one idea, but it's out of scope for this review.

+    base::TimeDelta elapsed =
+        clock_->Now() - budget_map_[origin.spec()].last_engagement_award;
+    int elapsedHours = elapsed.InHours();

[johnme, 2016/08/19 18:48:36]

If AddEngagementBudget gets called more than once per hour, the origin will
never get any budget, because elapsedHours will always be zero (since it is
rounded down). Use floating point maths for this instead, e.g. |InSecondsF() /
3600.0|?

[harkness, 2016/08/22 13:46:26]

Actually, I always meant to just return if the time was less than an hour to
prevent proliferation of small chunks, I just didn't add the code. I've added it
now.

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

I see. Even so, `awardRatio = elapsedHours / kBudgetDurationInHours;` still
sometimes awards less budget than the site deserves, which seems odd.

+    if (elapsedHours < kBudgetDurationInHours)
+      awardRatio = elapsedHours / kBudgetDurationInHours;
+  }
+
+  // Update the last_engagement_award to the current time.

[johnme, 2016/08/19 18:48:35]

Nit: please reword this comment to clarify that this line sometimes constructs a
new BudgetInfo entry in the map, rather than merely updating an existing one

[harkness, 2016/08/22 13:46:26]

Done.

+  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;
+  for (const auto& chunk : budget_map_[origin.spec()].chunks)

[johnme, 2016/08/19 18:48:35]

Nit: You can probably s/auto/BudgetChunk/ here since the type name is short

[harkness, 2016/08/22 13:46:26]

Peter has been pushing for everything to be auto, so I'd rather keep it auto for
consistency.

+    total += chunk.amount;
+
+  if (total < amount) {
+    callback.Run(false);
+    return;
+  }
+
+  double bill = amount;
+  auto& info = budget_map_[origin.spec()];

[johnme, 2016/08/19 18:48:36]

Please move this up so that the for loop above can use it (again, you may as
well use the full type name since it's short).

[harkness, 2016/08/22 13:46:25]

Done.

+  auto chunkIter = info.chunks.begin();
+  // This is guaranteed to exit because the total budget check above guarantees
+  // it.
+  while (bill > 0 && chunkIter != info.chunks.end()) {

[johnme, 2016/08/19 18:48:35]

Nit: perhaps simpler as a for loop; more efficient anyway since you only rewrite
the array (move assign all the elements) once:

int chunksToErase = 0;
for (BudgetChunk& chunk : info.chunks) {
  if (chunk.amount > bill) {
    chunk.amount -= bill
    break;
  }
  bill -= chunk.amount;
  chunksToErase++;
}
info.chunks.erase(info.chunks.begin(),
                  info.chunks.begin() + chunksToErase);

Or you could make BudgetChunks a std::list instead, then removing elements
wouldn't require the array to be rewritten (and there seem to be no downsides,
since you never index into it).

[harkness, 2016/08/22 13:46:25]

You're right that this use means we should change the vector to a list. I
updated the iteration to a for loop, but continued to use iterators rather than
directly accessing the chunks.

+    if (chunkIter->amount > bill) {
+      chunkIter->amount -= bill;
+      bill = 0;
+    } else {
+      bill -= chunkIter->amount;
+      chunkIter = info.chunks.erase(chunkIter);
+    }
+  }
 
   // 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);
 }
 
@@ -105,7 +187,9 @@ void BudgetDatabase::AddToCache(
                         base::Time::FromInternalValue(chunk.expiration()));
   }
 
-  budget_map_[origin.spec()] = std::move(chunks);
+  budget_map_[origin.spec()].chunks = std::move(chunks);

[johnme, 2016/08/19 18:48:36]

You could just write the chunks directly into budget_map_[origin.spec()].chunks
in the first place. Obviously you'd need to cache a reference to the BudgetInfo
rather than doing a string hash lookup every loop iteration, but that's probably
not a bad idea anyway.

[harkness, 2016/08/22 13:46:25]

Done.

+  budget_map_[origin.spec()].last_engagement_award =
+      base::Time::FromInternalValue(budget_proto->last_updated());
 
   callback.Run(success);
 }
@@ -133,7 +217,7 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
   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()];
+    const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
     for (const auto& chunk : base::Reversed(chunks)) {
       expectation.emplace_front(total, chunk.expiration);
       total += chunk.amount;
@@ -169,12 +253,14 @@ void BudgetDatabase::WriteCachedValuesToDatabase(
   if (IsCached(origin)) {
     // Build the Budget proto object.
     budget_service::Budget budget;
-    const BudgetChunks& chunks = budget_map_[origin.spec()];
+    const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;

[johnme, 2016/08/19 18:48:36]

Nit: store a reference to the BudgetInfo instead. Then you can re-use the
reference for set_last_updated.

[harkness, 2016/08/22 13:46:26]

Done.

     for (const auto& chunk : chunks) {
       budget_service::BudgetChunk* budget_chunk = budget.add_budget();
       budget_chunk->set_amount(chunk.amount);
       budget_chunk->set_expiration(chunk.expiration.ToInternalValue());
     }
+    budget.set_last_updated(
+        budget_map_[origin.spec()].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 +277,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());
 }