Adding more functionality to BudgetDatabase.

chrome/browser/budget_service/budget_database.cc

 // Copyright 2016 The Chromium Authors. All rights reserved.
 // Use of this source code is governed by a BSD-style license that can be
 // found in the LICENSE file.
 
 #include "chrome/browser/budget_service/budget_database.h"
 
 #include "base/containers/adapters.h"
+#include "base/time/clock.h"
+#include "base/time/default_clock.h"
 #include "chrome/browser/budget_service/budget.pb.h"
 #include "components/leveldb_proto/proto_database_impl.h"
 #include "content/public/browser/browser_thread.h"
 #include "url/gurl.h"
 
 using content::BrowserThread;
 
 namespace {
 
 // UMA are logged for the database with this string as part of the name.
 // They will be LevelDB.*.BackgroundBudgetService. Changes here should be
 // synchronized with histograms.xml.
 const char kDatabaseUMAName[] = "BackgroundBudgetService";
 
+// The default amount of time during which a budget will be valid.
+// This is 3 days = 72 hours.
+constexpr double kBudgetDurationInHours = 72.0;

[Peter Beverloo, 2016/08/01 17:33:08]

nit: no ".0"

[harkness, 2016/08/03 11:17:47]

Done.

+
 }  // namespace
 
 BudgetDatabase::BudgetDatabase(
     const base::FilePath& database_dir,
     const scoped_refptr<base::SequencedTaskRunner>& task_runner)
     : db_(new leveldb_proto::ProtoDatabaseImpl<budget_service::Budget>(
           task_runner)),
+      clock_(base::WrapUnique(new base::DefaultClock)),
       weak_ptr_factory_(this) {
   db_->Init(kDatabaseUMAName, database_dir,
             base::Bind(&BudgetDatabase::OnDatabaseInit,
                        weak_ptr_factory_.GetWeakPtr()));
 }
 
 BudgetDatabase::~BudgetDatabase() {}
 
-// TODO(harkness): Remove this method once the replacement is available.
-void BudgetDatabase::GetValue(const GURL& origin,
-                              const GetValueCallback& callback) {
-  DCHECK_EQ(origin.GetOrigin(), origin);
-  db_->GetEntry(origin.spec(), callback);
-}
-
-void BudgetDatabase::SetValue(const GURL& origin,
-                              const budget_service::Budget& budget,
-                              const SetValueCallback& callback) {
-  DCHECK_EQ(origin.GetOrigin(), origin);
-
-  // TODO(harkness) Remove this method once the replacement is available.
-
-  // Build structures to hold the updated values.
-  std::unique_ptr<
-      leveldb_proto::ProtoDatabase<budget_service::Budget>::KeyEntryVector>
-      entries(new leveldb_proto::ProtoDatabase<
-              budget_service::Budget>::KeyEntryVector());
-  entries->push_back(std::make_pair(origin.spec(), budget));
-  std::unique_ptr<std::vector<std::string>> keys_to_remove(
-      new std::vector<std::string>());
-
-  // Send the updates to the database.
-  db_->UpdateEntries(std::move(entries), std::move(keys_to_remove), callback);
-}
-
 void BudgetDatabase::GetBudgetDetails(
     const GURL& origin,
     const GetBudgetDetailsCallback& callback) {
   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()) {
     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);
+
+  // Look up the origin in our cache. We don't support adding budget without
+  // having first queried for the existing budget.

[Peter Beverloo, 2016/08/01 17:33:08]

nit: try to avoid "we"

[harkness, 2016/08/03 11:17:48]

Done.

+  const auto& info_iter = budget_map_.find(origin.spec());
+  DCHECK(info_iter != budget_map_.end());
+
+  base::Time now =

[Peter Beverloo, 2016/08/01 17:33:09]

now + 72 hours != |now| :-)

[harkness, 2016/08/03 11:17:47]

tsk tsk.

+      clock_->Now() + base::TimeDelta::FromHours(kBudgetDurationInHours);
+  info_iter->second.second.push_back(
+      std::make_pair(amount, now.ToInternalValue()));

[Peter Beverloo, 2016/08/01 17:33:09]

"info_iter->second.second" is rather unreadable.

How about the following:

  (1) swap the DCHECK with DCHECK_GT(budget_map_.count(origin.spec()), 0);
  (2) Just access budget_map_[origin.spec()] directly here.

[harkness, 2016/08/03 11:17:48]

Good idea, updated.

+
+  // 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());

[Peter Beverloo, 2016/08/01 17:33:08]

Let's remove this case until there is a ClearBudget() method? It's unreachable
code right now.

[harkness, 2016/08/03 11:17:48]

This case is possible once expiration is around, which is just the next CL. I'd
rather leave it in unless you feel strongly about it.

[Peter Beverloo, 2016/08/03 13:40:12]

Ok. In general, try to minimize unused code :).

+  } 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?
 }
 
 void BudgetDatabase::AddToCache(
     const GURL& origin,
     const AddToCacheCallback& callback,
     bool success,
     std::unique_ptr<budget_service::Budget> budget_proto) {
   // If the database read failed, there's nothing to add to the cache.
-  if (!success) {
+  if (!success || !budget_proto) {
     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.
   BudgetChunks chunks;
   for (const auto& chunk : budget_proto->budget())
     chunks.push_back(std::make_pair(chunk.amount(), chunk.expiration()));
 
-  budget_map_[origin.spec()] =
-      std::make_pair(budget_proto->debt(), std::move(chunks));
+  double debt = budget_proto->has_debt() ? budget_proto->debt() : 0;

[Peter Beverloo, 2016/08/01 17:33:08]

DCHECK(budget_proto->has_debt())?

When wouldn't it have debt? (You always set it on line 106.)

[harkness, 2016/08/03 11:17:47]

That's true, maybe I was just being overly cautious. Technically in the proto,
debt is an optional, so I was being defensive.

[Peter Beverloo, 2016/08/03 13:40:11]

So let's have a DCHECK() and continue on that assumption?

[harkness, 2016/08/08 12:59:35]

Done.

+  budget_map_[origin.spec()] = std::make_pair(debt, std::move(chunks));
 
   callback.Run(success);
 }
 
 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.
   if (!success) {
@@ skipping 15 matching lines @@
     expectation.push_front(std::make_pair(total, chunk.second));
     total += chunk.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));
 
   callback.Run(true /* success */, info.first, expectation);
 }
+
+// Override the default clock with the specified clock. Only used for testing.

[Peter Beverloo, 2016/08/01 17:33:09]

nit: redundant with header file

[harkness, 2016/08/03 11:17:47]

Generally when a method has a gotcha like "SHould only be used for testing", I
prefer extra commenting. Do you really feel that having the comment there hurts
the code?

[Peter Beverloo, 2016/08/03 13:40:11]

It's doubly redundant with the comment in the header file and the *ForTesting
suffix in the method name. I don't think it's adding any value. Up to you.

[harkness, 2016/08/08 12:59:35]

Done.

+void BudgetDatabase::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
+  clock_ = std::move(clock);
+}