Start plumbing connections from the BudgetManager to the BudgetDatabase

chrome/browser/budget_service/budget_manager.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_manager.h"
 
 #include <stdint.h>
 
 #include "base/callback.h"
 #include "base/memory/ptr_util.h"
 #include "base/metrics/histogram_macros.h"
 #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"
 #include "chrome/browser/engagement/site_engagement_score.h"
 #include "chrome/browser/engagement/site_engagement_service.h"
 #include "chrome/browser/profiles/profile.h"
 #include "chrome/common/pref_names.h"
 #include "components/pref_registry/pref_registry_syncable.h"
 #include "components/prefs/pref_service.h"
 #include "components/prefs/scoped_user_pref_update.h"
@@ skipping 58 matching lines @@
                                      budget, kSeparator, ses);
   DictionaryPrefUpdate update(profile->GetPrefs(), prefs::kBackgroundBudgetMap);
   base::DictionaryValue* map = update.Get();
 
   map->SetStringWithoutPathExpansion(origin.spec(), s);
 }
 
 }  // 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")),
+          base::ThreadTaskRunnerHandle::Get()),
+      weak_ptr_factory_(this) {}
 
 BudgetManager::~BudgetManager() {}
 
 // static
 void BudgetManager::RegisterProfilePrefs(
     user_prefs::PrefRegistrySyncable* registry) {
   registry->RegisterDictionaryPref(prefs::kBackgroundBudgetMap);
 }
 
 // static
@@ skipping 77 matching lines @@
   // Get the current SES score to write into the prefs with the new budget.
   SiteEngagementService* service = SiteEngagementService::Get(profile_);
   double ses_score = service->GetScore(origin);
 
   base::Time time = clock_->Now();
   SetBudgetDataInPrefs(profile_, origin, time.ToDoubleT(), budget, ses_score);
 
   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);
 }