Start plumbing connections from the BudgetManager to the 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 d55682f67258e1cb68b04e0f2d7073f78dfeb670..020846369e9e7229f3d2dcb93bf7c8c36afcd859 100644
--- a/chrome/browser/budget_service/budget_database.cc
+++ b/chrome/browser/budget_service/budget_database.cc
@@ -8,6 +8,9 @@
 #include "base/time/clock.h"
 #include "base/time/default_clock.h"
 #include "chrome/browser/budget_service/budget.pb.h"
+#include "chrome/browser/engagement/site_engagement_score.h"
+#include "chrome/browser/engagement/site_engagement_service.h"
+#include "chrome/browser/profiles/profile.h"
 #include "components/leveldb_proto/proto_database_impl.h"
 #include "content/public/browser/browser_thread.h"
 #include "url/gurl.h"
@@ -37,9 +40,11 @@ BudgetDatabase::BudgetInfo::BudgetInfo(const BudgetInfo&& other)
 BudgetDatabase::BudgetInfo::~BudgetInfo() {}
 
 BudgetDatabase::BudgetDatabase(
+    Profile* profile,
     const base::FilePath& database_dir,
     const scoped_refptr<base::SequencedTaskRunner>& task_runner)
-    : db_(new leveldb_proto::ProtoDatabaseImpl<budget_service::Budget>(
+    : profile_(profile),
+      db_(new leveldb_proto::ProtoDatabaseImpl<budget_service::Budget>(
           task_runner)),
       clock_(base::WrapUnique(new base::DefaultClock)),
       weak_ptr_factory_(this) {
@@ -55,117 +60,25 @@ void BudgetDatabase::GetBudgetDetails(
     const GetBudgetDetailsCallback& callback) {
   DCHECK_EQ(origin.GetOrigin(), origin);
 
-  // If this origin is already in the cache, immediately return the data.
-  if (IsCached(origin)) {
-    BrowserThread::PostTask(BrowserThread::UI, FROM_HERE,
-                            base::Bind(&BudgetDatabase::DidGetBudget,
-                                       weak_ptr_factory_.GetWeakPtr(), origin,
-                                       callback, true /* success */));
-    return;
-  }
-
-  // Otherwise, query for the data, add it to the cache, then return the result.
-  AddToCacheCallback cache_callback =
-      base::Bind(&BudgetDatabase::DidGetBudget, weak_ptr_factory_.GetWeakPtr(),
-                 origin, callback);
-  db_->GetEntry(origin.spec(), base::Bind(&BudgetDatabase::AddToCache,
-                                          weak_ptr_factory_.GetWeakPtr(),
-                                          origin, cache_callback));
-}
-
-void BudgetDatabase::AddBudget(const GURL& origin,
-                               double amount,
-                               const StoreBudgetCallback& callback) {
-  DCHECK_EQ(origin.GetOrigin(), 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()].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 score,
-                                         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 ratio = 1.0;
-
-  // 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 elapsed_hours = elapsed.InHours();
-    if (elapsed_hours == 0) {
-      // Don't give engagement awards for periods less than an hour.
-      callback.Run(true);
-      return;
-    }
-    if (elapsed_hours < kBudgetDurationInHours)
-      ratio = elapsed_hours / 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, score * ratio, callback);
+  // First, synchronize the cache.

[Peter Beverloo, 2016/08/25 15:31:19]

nit: drop. This is implied by the fact that you call "SyncCache" that take a
callback named "GetBudgetAfterSync". (Dito on line 73.)

Should that be called GetBudgetDetailsAfterSync by the way?

[harkness, 2016/08/25 16:36:13]

Done.

+  SyncCacheCallback cache_callback =
+      base::Bind(&BudgetDatabase::GetBudgetAfterSync,
+                 weak_ptr_factory_.GetWeakPtr(), origin, callback);
+  SyncCache(origin, cache_callback);

[Peter Beverloo, 2016/08/25 15:31:19]

nit: just inline the callback, no need to store it in a local first. (Dito on
line 77.)

[harkness, 2016/08/25 16:36:13]

Done.

 }
 
 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;
-  BudgetInfo& info = budget_map_[origin.spec()];
-  for (const BudgetChunk& 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;
-  for (auto iter = info.chunks.begin(); iter != info.chunks.end();) {
-    if (iter->amount > bill) {
-      iter->amount -= bill;
-      bill = 0;
-      ++iter;
-      break;
-    }
-    bill -= iter->amount;
-    iter = info.chunks.erase(iter);
-  }
-
-  // There should have been enough budget to cover the entire bill.
-  DCHECK_EQ(0, bill);
+  // First, synchronize the cache.
+  SyncCacheCallback cache_callback =
+      base::Bind(&BudgetDatabase::SpendBudgetAfterSync,
+                 weak_ptr_factory_.GetWeakPtr(), origin, amount, callback);
+  SyncCache(origin, cache_callback);
+}
 
-  // 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);
+void BudgetDatabase::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
+  clock_ = std::move(clock);
 }
 
 void BudgetDatabase::OnDatabaseInit(bool success) {
@@ -187,6 +100,13 @@ void BudgetDatabase::AddToCache(
     return;
   }
 
+  // If there were two simultaneous loads, don't overwrite the cache value,
+  // which might have been updated after the previous load.
+  if (IsCached(origin)) {
+    callback.Run(success);
+    return;
+  }
+
   // Add the data to the cache, converting from the proto format to an STL
   // format which is better for removing things from the list.
   BudgetInfo& info = budget_map_[origin.spec()];
@@ -201,9 +121,10 @@ void BudgetDatabase::AddToCache(
   callback.Run(success);
 }
 
-void BudgetDatabase::DidGetBudget(const GURL& origin,
-                                  const GetBudgetDetailsCallback& callback,
-                                  bool success) {
+void BudgetDatabase::GetBudgetAfterSync(
+    const GURL& origin,
+    const GetBudgetDetailsCallback& callback,
+    bool success) {

[Peter Beverloo, 2016/08/25 15:31:19]

nit: can this be `const` now?

[harkness, 2016/08/25 16:36:13]

It can't because we use the [] operator on the map, which could insert an
element. It won't because there should be an entry in the cache after
SyncCache(), but it theoretically could.

Actually, this made me realize that there WAS a way it could still not be in the
cache, if it had a score of 0, because of that last little change I made to not
give sites a 0 sized chunk. I've removed that now, it was just going to
complicate things, and if a site has 0 engagement, it shouldn't get budget.

   // If the database wasn't able to read the information, return the
   // failure and an empty BudgetPrediction.
   if (!success) {
@@ -211,9 +132,6 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
     return;
   }
 
-  // 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 BudgetPrediction describes a set of times
@@ -221,14 +139,12 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
   BudgetPrediction prediction;
   double total = 0;
 
-  if (IsCached(origin)) {
-    // Starting with the chunks that expire the farthest in the future, build up
-    // the budget predictions for those future times.
-    const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
-    for (const auto& chunk : base::Reversed(chunks)) {
-      prediction.emplace_front(total, chunk.expiration);
-      total += chunk.amount;
-    }
+  // Starting with the chunks that expire the farthest in the future, build up
+  // the budget predictions for those future times.
+  const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
+  for (const auto& chunk : base::Reversed(chunks)) {
+    prediction.emplace_front(total, chunk.expiration);
+    total += chunk.amount;
   }
 
   // Always add one entry at the front of the list for the total budget now.
@@ -237,16 +153,54 @@ void BudgetDatabase::DidGetBudget(const GURL& origin,
   callback.Run(true /* success */, prediction);
 }
 
-void BudgetDatabase::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
-  clock_ = std::move(clock);
+void BudgetDatabase::SpendBudgetAfterSync(const GURL& origin,
+                                          double amount,
+                                          const StoreBudgetCallback& callback,
+                                          bool success) {
+  if (!success) {
+    callback.Run(false);

[Peter Beverloo, 2016/08/25 15:31:19]

nit (here & lines 172, 256): false -> false /* success */

[harkness, 2016/08/25 16:36:13]

Done.

+    return;
+  }
+
+  // Walk the list of budget chunks to see if the origin has enough budget.
+  double total = 0;
+  BudgetInfo& info = budget_map_[origin.spec()];
+  for (const BudgetChunk& 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;
+  for (auto iter = info.chunks.begin(); iter != info.chunks.end();) {
+    if (iter->amount > bill) {
+      iter->amount -= bill;
+      bill = 0;
+      ++iter;

[Peter Beverloo, 2016/08/25 15:31:19]

nit: delete, this doesn't do anything since you `break`

[harkness, 2016/08/25 16:36:13]

Done.

+      break;
+    }
+    bill -= iter->amount;
+    iter = info.chunks.erase(iter);
+  }
+
+  // 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);
 }
 
 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>
@@ -278,12 +232,86 @@ void BudgetDatabase::WriteCachedValuesToDatabase(
   db_->UpdateEntries(std::move(entries), std::move(keys_to_remove), callback);
 }
 
-void BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
-  if (!IsCached(origin))
+void BudgetDatabase::SyncCache(const GURL& origin,
+                               const SyncCacheCallback& callback) {
+  DCHECK_EQ(origin, origin.GetOrigin());
+
+  // If the origin isn't already cached, add it to the cache.
+  if (!IsCached(origin)) {
+    AddToCacheCallback add_callback =
+        base::Bind(&BudgetDatabase::SyncLoadedCache,
+                   weak_ptr_factory_.GetWeakPtr(), origin, callback);
+    db_->GetEntry(origin.spec(), base::Bind(&BudgetDatabase::AddToCache,
+                                            weak_ptr_factory_.GetWeakPtr(),
+                                            origin, add_callback));
     return;
+  }
+  SyncLoadedCache(origin, callback, true /* success */);
+}
 
-  base::Time now = clock_->Now();
+void BudgetDatabase::SyncLoadedCache(const GURL& origin,
+                                     const SyncCacheCallback& callback,
+                                     bool success) {
+  if (!success) {
+    callback.Run(false);
+    return;
+  }
+
+  // Get the SES score and add engagement budget for the site.
+  AddEngagementBudget(origin);
+
+  // Now, cleanup any expired budget chunks for the origin.
+  bool needs_write = CleanupExpiredBudget(origin);
+
+  if (needs_write)
+    WriteCachedValuesToDatabase(origin, callback);
+  else
+    callback.Run(success);
+}
+
+void BudgetDatabase::AddEngagementBudget(const GURL& origin) {
+  // Get the current SES score, which we'll use to set a new budget.
+  SiteEngagementService* service = SiteEngagementService::Get(profile_);
+  double score = service->GetScore(origin);
+
+  // Don't force new sites to start with a bad chunk.
+  if (score == 0)
+    return;
+
+  // By default we award the "full" award. Then that ratio is decreased if
+  // there have been other awards recently.
+  double ratio = 1.0;
+
+  // Calculate how much budget should be awarded. If there is no entry in the
+  // cache then we award a full amount.
+  if (IsCached(origin)) {
+    base::TimeDelta elapsed =
+        clock_->Now() - budget_map_[origin.spec()].last_engagement_award;
+    int elapsed_hours = elapsed.InHours();
+    // Don't give engagement awards for periods less than an hour.
+    if (elapsed_hours < 1)
+      return;
+    if (elapsed_hours < kBudgetDurationInHours)
+      ratio = elapsed_hours / 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();
+
+  // 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()].chunks.emplace_back(ratio * score, expiration);
+}
 
+// Cleans up budget in the cache. Relies on the caller eventually writing the
+// cache back to the database.
+bool BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
+  if (!IsCached(origin))
+    return false;
+
+  base::Time now = clock_->Now();
   BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
   auto cleanup_iter = chunks.begin();
 
@@ -295,6 +323,13 @@ void BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
   // 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))
+          clock_->Now() - base::TimeDelta::FromHours(kBudgetDurationInHours)) {
     budget_map_.erase(origin.spec());
+    return true;
+  }
+
+  // Although some things may have expired, there are some chunks still valid.
+  // Don't write to the DB now, write either when all chunks expire or when the
+  // origin spends some budget.
+  return false;
 }