Start plumbing connections from the BudgetManager to the BudgetDatabase

chrome/browser/budget_service/budget_manager.cc

Index: chrome/browser/budget_service/budget_manager.cc
diff --git a/chrome/browser/budget_service/budget_manager.cc b/chrome/browser/budget_service/budget_manager.cc
index d11b409ae03ee723d3c1673df91cd8353b9fa694..5f40ec09f6aaf7a46dc9458f013ad8ef60eb2203 100644
--- a/chrome/browser/budget_service/budget_manager.cc
+++ b/chrome/browser/budget_service/budget_manager.cc
@@ -12,6 +12,7 @@
 #include "base/strings/string_number_conversions.h"
 #include "base/strings/string_split.h"
 #include "base/strings/stringprintf.h"
+#include "base/threading/thread_task_runner_handle.h"
 #include "base/time/clock.h"
 #include "base/time/default_clock.h"
 #include "base/time/time.h"
@@ -90,9 +91,15 @@ void SetBudgetDataInPrefs(Profile* profile,
 }  // namespace
 
 BudgetManager::BudgetManager(Profile* profile)
-    : clock_(base::WrapUnique(new base::DefaultClock)), profile_(profile) {
-  DCHECK(profile);
-}
+    : clock_(base::WrapUnique(new base::DefaultClock)),
+      profile_(profile),
+      db_(profile,
+          profile->GetPath().Append(FILE_PATH_LITERAL("BudgetDatabase")),
+          BrowserThread::GetBlockingPool()
+              ->GetSequencedTaskRunnerWithShutdownBehavior(
+                  base::SequencedWorkerPool::GetSequenceToken(),
+                  base::SequencedWorkerPool::CONTINUE_ON_SHUTDOWN)),
+      weak_ptr_factory_(this) {}
 
 BudgetManager::~BudgetManager() {}
 
@@ -190,6 +197,57 @@ void BudgetManager::StoreBudget(const GURL& origin,
   BrowserThread::PostTask(BrowserThread::UI, FROM_HERE, base::Bind(closure));
 }
 
+void BudgetManager::Reserve(const GURL& origin,
+                            blink::mojom::BudgetOperationType type,
+                            const ReserveCallback& callback) {
+  DCHECK_EQ(origin, origin.GetOrigin());
+
+  BudgetDatabase::StoreBudgetCallback reserve_callback =
+      base::Bind(&BudgetManager::DidReserve, weak_ptr_factory_.GetWeakPtr(),
+                 origin, type, callback);
+  db_.SpendBudget(origin, GetCost(type), callback);
+}
+
+void BudgetManager::Consume(const GURL& origin,
+                            blink::mojom::BudgetOperationType type,
+                            const ConsumeCallback& callback) {
+  DCHECK_EQ(origin, origin.GetOrigin());
+  bool found_reservation = false;
+
+  // First, see if there is a reservation already.
+  auto count = reservation_map_.find(origin.spec());
+  if (count != reservation_map_.end()) {
+    if (count->second == 1)
+      reservation_map_.erase(origin.spec());
+    else
+      reservation_map_[origin.spec()]--;
+    found_reservation = true;
+  }
+
+  if (found_reservation) {
+    callback.Run(true);
+    return;
+  }
+
+  // If there wasn't a reservation already, try to directly consume budget.
+  // The callback will return directly to the caller.
+  db_.SpendBudget(origin, GetCost(type), callback);
+}
+
+void BudgetManager::DidReserve(const GURL& origin,
+                               blink::mojom::BudgetOperationType type,
+                               const ReserveCallback& callback,
+                               bool success) {
+  if (!success) {
+    callback.Run(false);
+    return;
+  }
+
+  // Write the new reservation into the map.
+  reservation_map_[origin.spec()]++;
+  callback.Run(true);
+}
+
 // Override the default clock with the specified clock. Only used for testing.
 void BudgetManager::SetClockForTesting(std::unique_ptr<base::Clock> clock) {
   clock_ = std::move(clock);