Replace GURL in BudgetDatabase/BudgetManager with url::Origin

The Mojo service uses url::Origin, and that's the better way to be
storing origins, so converting over to using that instead of GURL. This
also does a bit of cleanup on the tests, which had a mix of methods which
took an origin argument and methods which created the origin internally.
Now the test has a member which is calculated once.

BUG=617971

Review-Url: https://codereview.chromium.org/2324133002
Cr-Commit-Position: refs/heads/master@{#417593}

chrome/browser/budget_service/budget_database.cc

@@ -14,6 +14,7 @@
 #include "components/leveldb_proto/proto_database_impl.h"
 #include "content/public/browser/browser_thread.h"
 #include "url/gurl.h"
+#include "url/origin.h"
 
 using content::BrowserThread;
 
@@ -55,16 +56,14 @@ BudgetDatabase::BudgetDatabase(
 
 BudgetDatabase::~BudgetDatabase() {}
 
-void BudgetDatabase::GetBudgetDetails(const GURL& origin,
+void BudgetDatabase::GetBudgetDetails(const url::Origin& origin,
                                       const GetBudgetCallback& callback) {
-  DCHECK_EQ(origin.GetOrigin(), origin);
-
   SyncCache(origin,
             base::Bind(&BudgetDatabase::GetBudgetAfterSync,
                        weak_ptr_factory_.GetWeakPtr(), origin, callback));
 }
 
-void BudgetDatabase::SpendBudget(const GURL& origin,
+void BudgetDatabase::SpendBudget(const url::Origin& origin,
                                  double amount,
                                  const StoreBudgetCallback& callback) {
   SyncCache(origin, base::Bind(&BudgetDatabase::SpendBudgetAfterSync,
@@ -80,13 +79,13 @@ void BudgetDatabase::OnDatabaseInit(bool success) {
   // TODO(harkness): Consider caching the budget database now?
 }
 
-bool BudgetDatabase::IsCached(const GURL& origin) const {
-  return budget_map_.find(origin.spec()) != budget_map_.end();
+bool BudgetDatabase::IsCached(const url::Origin& origin) const {
+  return budget_map_.find(origin) != budget_map_.end();
 }
 
-double BudgetDatabase::GetBudget(const GURL& origin) const {
+double BudgetDatabase::GetBudget(const url::Origin& origin) const {
   double total = 0;
-  auto iter = budget_map_.find(origin.spec());
+  auto iter = budget_map_.find(origin);
   if (iter == budget_map_.end())
     return total;
 
@@ -97,7 +96,7 @@ double BudgetDatabase::GetBudget(const GURL& origin) const {
 }
 
 void BudgetDatabase::AddToCache(
-    const GURL& origin,
+    const url::Origin& origin,
     const AddToCacheCallback& callback,
     bool success,
     std::unique_ptr<budget_service::Budget> budget_proto) {
@@ -116,7 +115,7 @@ void BudgetDatabase::AddToCache(
 
   // 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()];
+  BudgetInfo& info = budget_map_[origin];
   for (const auto& chunk : budget_proto->budget()) {
     info.chunks.emplace_back(chunk.amount(),
                              base::Time::FromInternalValue(chunk.expiration()));
@@ -128,7 +127,7 @@ void BudgetDatabase::AddToCache(
   callback.Run(success);
 }
 
-void BudgetDatabase::GetBudgetAfterSync(const GURL& origin,
+void BudgetDatabase::GetBudgetAfterSync(const url::Origin& origin,
                                         const GetBudgetCallback& callback,
                                         bool success) {
   mojo::Array<blink::mojom::BudgetStatePtr> predictions;
@@ -155,7 +154,7 @@ void BudgetDatabase::GetBudgetAfterSync(const GURL& origin,
 
   // Starting with the soonest expiring chunks, add entries for the
   // expiration times going forward.
-  const BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
+  const BudgetChunks& chunks = budget_map_[origin].chunks;
   for (const auto& chunk : chunks) {
     blink::mojom::BudgetStatePtr prediction(blink::mojom::BudgetState::New());
     total -= chunk.amount;
@@ -170,7 +169,7 @@ void BudgetDatabase::GetBudgetAfterSync(const GURL& origin,
                std::move(predictions));
 }
 
-void BudgetDatabase::SpendBudgetAfterSync(const GURL& origin,
+void BudgetDatabase::SpendBudgetAfterSync(const url::Origin& origin,
                                           double amount,
                                           const StoreBudgetCallback& callback,
                                           bool success) {
@@ -181,11 +180,11 @@ void BudgetDatabase::SpendBudgetAfterSync(const GURL& origin,
 
   // Get the current SES score, to generate UMA.
   SiteEngagementService* service = SiteEngagementService::Get(profile_);
-  double score = service->GetScore(origin);
+  double score = service->GetScore(GURL(origin.Serialize()));
 
   // Walk the list of budget chunks to see if the origin has enough budget.
   double total = 0;
-  BudgetInfo& info = budget_map_[origin.spec()];
+  BudgetInfo& info = budget_map_[origin];
   for (const BudgetChunk& chunk : info.chunks)
     total += chunk.amount;
 
@@ -222,7 +221,7 @@ void BudgetDatabase::SpendBudgetAfterSync(const GURL& origin,
 }
 
 void BudgetDatabase::WriteCachedValuesToDatabase(
-    const GURL& origin,
+    const url::Origin& origin,
     const StoreBudgetCallback& callback) {
   // Create the data structures that are passed to the ProtoDatabase.
   std::unique_ptr<
@@ -237,7 +236,7 @@ void BudgetDatabase::WriteCachedValuesToDatabase(
   if (IsCached(origin)) {
     // Build the Budget proto object.
     budget_service::Budget budget;
-    const BudgetInfo& info = budget_map_[origin.spec()];
+    const BudgetInfo& info = budget_map_[origin];
     for (const auto& chunk : info.chunks) {
       budget_service::BudgetChunk* budget_chunk = budget.add_budget();
       budget_chunk->set_amount(chunk.amount);
@@ -245,34 +244,32 @@ void BudgetDatabase::WriteCachedValuesToDatabase(
     }
     budget.set_engagement_last_updated(
         info.last_engagement_award.ToInternalValue());
-    entries->push_back(std::make_pair(origin.spec(), budget));
+    entries->push_back(std::make_pair(origin.Serialize(), budget));
   } else {
     // If the origin doesn't exist in the cache, this is a remove operation.
-    keys_to_remove->push_back(origin.spec());
+    keys_to_remove->push_back(origin.Serialize());
   }
 
   // Send the updates to the database.
   db_->UpdateEntries(std::move(entries), std::move(keys_to_remove), callback);
 }
 
-void BudgetDatabase::SyncCache(const GURL& origin,
+void BudgetDatabase::SyncCache(const url::Origin& 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));
+    db_->GetEntry(origin.Serialize(), base::Bind(&BudgetDatabase::AddToCache,
+                                                 weak_ptr_factory_.GetWeakPtr(),
+                                                 origin, add_callback));
     return;
   }
   SyncLoadedCache(origin, callback, true /* success */);
 }
 
-void BudgetDatabase::SyncLoadedCache(const GURL& origin,
+void BudgetDatabase::SyncLoadedCache(const url::Origin& origin,
                                      const SyncCacheCallback& callback,
                                      bool success) {
   if (!success) {
@@ -292,10 +289,10 @@ void BudgetDatabase::SyncLoadedCache(const GURL& origin,
     callback.Run(success);
 }
 
-void BudgetDatabase::AddEngagementBudget(const GURL& origin) {
+void BudgetDatabase::AddEngagementBudget(const url::Origin& 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);
+  double score = service->GetScore(GURL(origin.Serialize()));
 
   // By default we award the "full" award. Then that ratio is decreased if
   // there have been other awards recently.
@@ -305,7 +302,7 @@ void BudgetDatabase::AddEngagementBudget(const GURL& origin) {
   // cache then we award a full amount.
   if (IsCached(origin)) {
     base::TimeDelta elapsed =
-        clock_->Now() - budget_map_[origin.spec()].last_engagement_award;
+        clock_->Now() - budget_map_[origin].last_engagement_award;
     int elapsed_hours = elapsed.InHours();
     // Don't give engagement awards for periods less than an hour.
     if (elapsed_hours < 1)
@@ -316,12 +313,12 @@ void BudgetDatabase::AddEngagementBudget(const GURL& origin) {
 
   // 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();
+  budget_map_[origin].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);
+  budget_map_[origin].chunks.emplace_back(ratio * score, expiration);
 
   // Any time we award engagement budget, which is done at most once an hour
   // whenever any budget action is taken, record the budget.
@@ -331,12 +328,12 @@ void BudgetDatabase::AddEngagementBudget(const GURL& origin) {
 
 // Cleans up budget in the cache. Relies on the caller eventually writing the
 // cache back to the database.
-bool BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
+bool BudgetDatabase::CleanupExpiredBudget(const url::Origin& origin) {
   if (!IsCached(origin))
     return false;
 
   base::Time now = clock_->Now();
-  BudgetChunks& chunks = budget_map_[origin.spec()].chunks;
+  BudgetChunks& chunks = budget_map_[origin].chunks;
   auto cleanup_iter = chunks.begin();
 
   // This relies on the list of chunks being in timestamp order.
@@ -346,9 +343,9 @@ bool BudgetDatabase::CleanupExpiredBudget(const GURL& origin) {
   // 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 <
+      budget_map_[origin].last_engagement_award <
           clock_->Now() - base::TimeDelta::FromHours(kBudgetDurationInHours)) {
-    budget_map_.erase(origin.spec());
+    budget_map_.erase(origin);
     return true;
   }
 

chrome/browser/budget_service/budget_database.h

@@ -6,8 +6,8 @@
 #define CHROME_BROWSER_BUDGET_SERVICE_BUDGET_DATABASE_H_
 
 #include <list>
+#include <map>
 #include <memory>
-#include <unordered_map>
 
 #include "base/callback_forward.h"
 #include "base/macros.h"
@@ -25,7 +25,10 @@ namespace budget_service {
 class Budget;
 }
 
-class GURL;
+namespace url {
+class Origin;
+}
+
 class Profile;
 
 // A class used to asynchronously read and write details of the budget
@@ -48,13 +51,14 @@ class BudgetDatabase {
 
   // Get the full budget expectation for the origin. This will return a
   // sequence of time points and the expected budget at those times.
-  void GetBudgetDetails(const GURL& origin, const GetBudgetCallback& callback);
+  void GetBudgetDetails(const url::Origin& origin,
+                        const GetBudgetCallback& callback);
 
   // Spend a particular amount of budget for an origin. The callback takes
   // a boolean which indicates whether the origin had enough budget to spend.
   // If it returns success, then the budget was deducted and the result written
   // to the database.
-  void SpendBudget(const GURL& origin,
+  void SpendBudget(const url::Origin& origin,
                    double amount,
                    const StoreBudgetCallback& callback);
 
@@ -98,29 +102,29 @@ class BudgetDatabase {
 
   void OnDatabaseInit(bool success);
 
-  bool IsCached(const GURL& origin) const;
+  bool IsCached(const url::Origin& origin) const;
 
-  double GetBudget(const GURL& origin) const;
+  double GetBudget(const url::Origin& origin) const;
 
-  void AddToCache(const GURL& origin,
+  void AddToCache(const url::Origin& origin,
                   const AddToCacheCallback& callback,
                   bool success,
                   std::unique_ptr<budget_service::Budget> budget);
 
-  void GetBudgetAfterSync(const GURL& origin,
+  void GetBudgetAfterSync(const url::Origin& origin,
                           const GetBudgetCallback& callback,
                           bool success);
 
-  void SpendBudgetAfterSync(const GURL& origin,
+  void SpendBudgetAfterSync(const url::Origin& origin,
                             double amount,
                             const StoreBudgetCallback& callback,
                             bool success);
 
-  void WriteCachedValuesToDatabase(const GURL& origin,
+  void WriteCachedValuesToDatabase(const url::Origin& origin,
                                    const StoreBudgetCallback& callback);
 
-  void SyncCache(const GURL& origin, const SyncCacheCallback& callback);
-  void SyncLoadedCache(const GURL& origin,
+  void SyncCache(const url::Origin& origin, const SyncCacheCallback& callback);
+  void SyncLoadedCache(const url::Origin& origin,
                        const SyncCacheCallback& callback,
                        bool success);
 
@@ -128,9 +132,9 @@ class BudgetDatabase {
   // engagement score of the origin, and then calculates when engagement budget
   // was last awarded and awards a portion of the score based on that.
   // This only writes budget to the cache.
-  void AddEngagementBudget(const GURL& origin);
+  void AddEngagementBudget(const url::Origin& origin);
 
-  bool CleanupExpiredBudget(const GURL& origin);
+  bool CleanupExpiredBudget(const url::Origin& origin);
 
   Profile* profile_;
 
@@ -138,7 +142,7 @@ class BudgetDatabase {
   std::unique_ptr<leveldb_proto::ProtoDatabase<budget_service::Budget>> db_;
 
   // Cached data for the origins which have been loaded.
-  std::unordered_map<std::string, BudgetInfo> budget_map_;
+  std::map<url::Origin, BudgetInfo> budget_map_;
 
   // The clock used to vend times.
   std::unique_ptr<base::Clock> clock_;

chrome/browser/budget_service/budget_database_unittest.cc

@@ -19,6 +19,8 @@
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "mojo/public/cpp/bindings/array.h"
 #include "testing/gtest/include/gtest/gtest.h"
+#include "url/gurl.h"
+#include "url/origin.h"
 
 namespace {
 
@@ -35,7 +37,8 @@ class BudgetDatabaseTest : public ::testing::Test {
       : success_(false),
         db_(&profile_,
             profile_.GetPath().Append(FILE_PATH_LITERAL("BudgetDatabase")),
-            base::ThreadTaskRunnerHandle::Get()) {}
+            base::ThreadTaskRunnerHandle::Get()),
+        origin_(url::Origin(GURL(kTestOrigin))) {}
 
   void WriteBudgetComplete(base::Closure run_loop_closure, bool success) {
     success_ = success;
@@ -43,9 +46,9 @@ class BudgetDatabaseTest : public ::testing::Test {
   }
 
   // Spend budget for the origin.
-  bool SpendBudget(const GURL& origin, double amount) {
+  bool SpendBudget(double amount) {
     base::RunLoop run_loop;
-    db_.SpendBudget(origin, amount,
+    db_.SpendBudget(origin(), amount,
                     base::Bind(&BudgetDatabaseTest::WriteBudgetComplete,
                                base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
@@ -65,13 +68,13 @@ class BudgetDatabaseTest : public ::testing::Test {
   void GetBudgetDetails() {
     base::RunLoop run_loop;
     db_.GetBudgetDetails(
-        GURL(kTestOrigin),
-        base::Bind(&BudgetDatabaseTest::GetBudgetDetailsComplete,
-                   base::Unretained(this), run_loop.QuitClosure()));
+        origin(), base::Bind(&BudgetDatabaseTest::GetBudgetDetailsComplete,
+                             base::Unretained(this), run_loop.QuitClosure()));
     run_loop.Run();
   }
 
   Profile* profile() { return &profile_; }
+  const url::Origin& origin() const { return origin_; }
 
   // Setup a test clock so that the tests can control time.
   base::SimpleTestClock* SetClockForTesting() {
@@ -80,9 +83,9 @@ class BudgetDatabaseTest : public ::testing::Test {
     return clock;
   }
 
-  void SetSiteEngagementScore(const GURL& url, double score) {
+  void SetSiteEngagementScore(double score) {
     SiteEngagementService* service = SiteEngagementService::Get(&profile_);
-    service->ResetScoreForURL(url, score);
+    service->ResetScoreForURL(GURL(kTestOrigin), score);
   }
 
  protected:
@@ -96,10 +99,10 @@ class BudgetDatabaseTest : public ::testing::Test {
   TestingProfile profile_;
   BudgetDatabase db_;
   base::HistogramTester histogram_tester_;
+  const url::Origin origin_;
 };
 
 TEST_F(BudgetDatabaseTest, GetBudgetNoBudgetOrSES) {
-  const GURL origin(kTestOrigin);
   GetBudgetDetails();
   ASSERT_TRUE(success_);
   ASSERT_EQ(2U, prediction_.size());
@@ -107,13 +110,12 @@ TEST_F(BudgetDatabaseTest, GetBudgetNoBudgetOrSES) {
 }
 
 TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
-  const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
   base::Time expiration_time =
       clock->Now() + base::TimeDelta::FromHours(kDefaultExpirationInHours);
 
   // Set the default site engagement.
-  SetSiteEngagementScore(origin, kDefaultEngagement);
+  SetSiteEngagementScore(kDefaultEngagement);
 
   // The budget should include a full share of the engagement.
   GetBudgetDetails();
@@ -152,11 +154,10 @@ TEST_F(BudgetDatabaseTest, AddEngagementBudgetTest) {
 }
 
 TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
-  const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
 
   // Set the default site engagement.
-  SetSiteEngagementScore(origin, kDefaultEngagement);
+  SetSiteEngagementScore(kDefaultEngagement);
 
   // Intialize the budget with several chunks.
   GetBudgetDetails();
@@ -166,7 +167,7 @@ TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
   GetBudgetDetails();
 
   // Spend an amount of budget less than kDefaultEngagement.
-  ASSERT_TRUE(SpendBudget(origin, 1));
+  ASSERT_TRUE(SpendBudget(1));
   GetBudgetDetails();
 
   // There should still be three chunks of budget of size kDefaultEngagement-1,
@@ -181,14 +182,14 @@ TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
 
   // Now spend enough that it will use up the rest of the first chunk and all of
   // the second chunk, but not all of the third chunk.
-  ASSERT_TRUE(SpendBudget(origin, kDefaultEngagement + daily_budget));
+  ASSERT_TRUE(SpendBudget(kDefaultEngagement + daily_budget));
   GetBudgetDetails();
   ASSERT_EQ(2U, prediction_.size());
   ASSERT_DOUBLE_EQ(daily_budget - 1, prediction_[0]->budget_at);
 
   // Validate that the code returns false if SpendBudget tries to spend more
   // budget than the origin has.
-  EXPECT_FALSE(SpendBudget(origin, kDefaultEngagement));
+  EXPECT_FALSE(SpendBudget(kDefaultEngagement));
   GetBudgetDetails();
   ASSERT_EQ(2U, prediction_.size());
   ASSERT_DOUBLE_EQ(daily_budget - 1, prediction_[0]->budget_at);
@@ -196,7 +197,7 @@ TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
   // Advance time until the last remaining chunk should be expired, then query
   // for the full engagement worth of budget.
   clock->Advance(base::TimeDelta::FromHours(kDefaultExpirationInHours + 1));
-  EXPECT_TRUE(SpendBudget(origin, kDefaultEngagement));
+  EXPECT_TRUE(SpendBudget(kDefaultEngagement));
 }
 
 // There are times when a device's clock could move backwards in time, either
@@ -204,11 +205,10 @@ TEST_F(BudgetDatabaseTest, SpendBudgetTest) {
 // time goes backwards and then forwards again, the origin isn't granted extra
 // budget.
 TEST_F(BudgetDatabaseTest, GetBudgetNegativeTime) {
-  const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
 
   // Set the default site engagement.
-  SetSiteEngagementScore(origin, kDefaultEngagement);
+  SetSiteEngagementScore(kDefaultEngagement);
 
   // Initialize the budget with two chunks.
   GetBudgetDetails();
@@ -236,11 +236,10 @@ TEST_F(BudgetDatabaseTest, GetBudgetNegativeTime) {
 }
 
 TEST_F(BudgetDatabaseTest, CheckBackgroundBudgetHistogram) {
-  const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
 
   // Set the default site engagement.
-  SetSiteEngagementScore(origin, kDefaultEngagement);
+  SetSiteEngagementScore(kDefaultEngagement);
 
   // Initialize the budget with some interesting chunks: 30 budget, 3 budget,
   // 0 budget, and then after the first two expire, another 30 budget.
@@ -266,13 +265,12 @@ TEST_F(BudgetDatabaseTest, CheckBackgroundBudgetHistogram) {
 }
 
 TEST_F(BudgetDatabaseTest, CheckEngagementHistograms) {
-  const GURL origin(kTestOrigin);
   base::SimpleTestClock* clock = SetClockForTesting();
 
   // Set the engagement to twice the cost of an action.
   double cost = 2;
   double engagement = cost * 2;
-  SetSiteEngagementScore(origin, engagement);
+  SetSiteEngagementScore(engagement);
 
   // Get the budget, which will award a chunk of budget equal to engagement.
   GetBudgetDetails();
@@ -280,19 +278,19 @@ TEST_F(BudgetDatabaseTest, CheckEngagementHistograms) {
   // Now spend the budget to trigger the UMA recording the SES score. The first
   // call shouldn't write any UMA. The second should write a lowSES entry, and
   // the third should write a noSES entry.
-  ASSERT_TRUE(SpendBudget(origin, cost));
-  ASSERT_TRUE(SpendBudget(origin, cost));
-  ASSERT_FALSE(SpendBudget(origin, cost));
+  ASSERT_TRUE(SpendBudget(cost));
+  ASSERT_TRUE(SpendBudget(cost));
+  ASSERT_FALSE(SpendBudget(cost));
 
   // Advance the clock by 12 days (to guarantee a full new engagement grant)
   // then change the SES score to get a different UMA entry, then spend the
   // budget again.
   clock->Advance(base::TimeDelta::FromDays(12));
   GetBudgetDetails();
-  SetSiteEngagementScore(origin, engagement * 2);
-  ASSERT_TRUE(SpendBudget(origin, cost));
-  ASSERT_TRUE(SpendBudget(origin, cost));
-  ASSERT_FALSE(SpendBudget(origin, cost));
+  SetSiteEngagementScore(engagement * 2);
+  ASSERT_TRUE(SpendBudget(cost));
+  ASSERT_TRUE(SpendBudget(cost));
+  ASSERT_FALSE(SpendBudget(cost));
 
   // Now check the UMA. Both UMA should have 2 buckets with 1 entry each.
   std::vector<base::Bucket> no_budget_buckets =

chrome/browser/budget_service/budget_manager.cc

@@ -18,6 +18,7 @@
 #include "components/prefs/pref_service.h"
 #include "content/public/browser/browser_thread.h"
 #include "third_party/WebKit/public/platform/modules/budget_service/budget_service.mojom.h"
+#include "url/origin.h"
 
 using content::BrowserThread;
 
@@ -65,35 +66,32 @@ double BudgetManager::GetCost(blink::mojom::BudgetOperationType type) {
   return SiteEngagementScore::kMaxPoints + 1.0;
 }
 
-void BudgetManager::GetBudget(const GURL& origin,
+void BudgetManager::GetBudget(const url::Origin& origin,
                               const GetBudgetCallback& callback) {
   db_.GetBudgetDetails(origin, callback);
 }
 
-void BudgetManager::Reserve(const GURL& origin,
+void BudgetManager::Reserve(const url::Origin& origin,
                             blink::mojom::BudgetOperationType type,
                             const ReserveCallback& callback) {
-  DCHECK_EQ(origin, origin.GetOrigin());
-
   db_.SpendBudget(
       origin, GetCost(type),
       base::Bind(&BudgetManager::DidReserve, weak_ptr_factory_.GetWeakPtr(),
                  origin, type, callback));
 }
 
-void BudgetManager::Consume(const GURL& origin,
+void BudgetManager::Consume(const url::Origin& 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());
+  auto count = reservation_map_.find(origin.host());
   if (count != reservation_map_.end()) {
     if (count->second == 1)
-      reservation_map_.erase(origin.spec());
+      reservation_map_.erase(origin.host());
     else
-      reservation_map_[origin.spec()]--;
+      reservation_map_[origin.host()]--;
     found_reservation = true;
   }
 
@@ -107,7 +105,7 @@ void BudgetManager::Consume(const GURL& origin,
   db_.SpendBudget(origin, GetCost(type), callback);
 }
 
-void BudgetManager::DidReserve(const GURL& origin,
+void BudgetManager::DidReserve(const url::Origin& origin,
                                blink::mojom::BudgetOperationType type,
                                const ReserveCallback& callback,
                                bool success) {
@@ -117,6 +115,6 @@ void BudgetManager::DidReserve(const GURL& origin,
   }
 
   // Write the new reservation into the map.
-  reservation_map_[origin.spec()]++;
+  reservation_map_[origin.host()]++;
   callback.Run(true);
 }

chrome/browser/budget_service/budget_manager.h

@@ -13,7 +13,6 @@
 #include "chrome/browser/budget_service/budget_database.h"
 #include "components/keyed_service/core/keyed_service.h"
 #include "third_party/WebKit/public/platform/modules/budget_service/budget_service.mojom.h"
-#include "url/gurl.h"
 
 class Profile;
 
@@ -21,6 +20,10 @@ namespace user_prefs {
 class PrefRegistrySyncable;
 }
 
+namespace url {
+class Origin;
+}
+
 // A budget manager to help Chrome decide how much background work a service
 // worker should be able to do on behalf of the user. The budget is calculated
 // based on the Site Engagment Score and is consumed when a origin does
@@ -42,18 +45,18 @@ class BudgetManager : public KeyedService {
   // Get the budget associated with the origin. This is passed to the
   // callback. Budget will be a sequence of points describing the time and
   // the budget at that time.
-  void GetBudget(const GURL& origin, const GetBudgetCallback& callback);
+  void GetBudget(const url::Origin& origin, const GetBudgetCallback& callback);
 
   // Spend enough budget to cover the cost of the desired action and create
   // a reservation for that action. If this returns true to the callback, then
   // the next action will consume that reservation and not cost any budget.
-  void Reserve(const GURL& origin,
+  void Reserve(const url::Origin& origin,
                blink::mojom::BudgetOperationType type,
                const ReserveCallback& callback);
 
   // Spend budget, first consuming a reservation if one exists, or spend
   // directly from the budget.
-  void Consume(const GURL& origin,
+  void Consume(const url::Origin& origin,
                blink::mojom::BudgetOperationType type,
                const ConsumeCallback& callback);
 
@@ -62,7 +65,7 @@ class BudgetManager : public KeyedService {
 
   // Called as a callback from BudgetDatabase after it has made a reserve
   // decision.
-  void DidReserve(const GURL& origin,
+  void DidReserve(const url::Origin& origin,
                   blink::mojom::BudgetOperationType type,
                   const ReserveCallback& callback,
                   bool success);

chrome/browser/budget_service/budget_manager_unittest.cc

@@ -17,6 +17,7 @@
 #include "content/public/test/test_browser_thread_bundle.h"
 #include "testing/gtest/include/gtest/gtest.h"
 #include "third_party/WebKit/public/platform/modules/budget_service/budget_service.mojom.h"
+#include "url/origin.h"
 
 namespace {
 
@@ -27,18 +28,20 @@ const double kTestSES = 48.0;
 
 class BudgetManagerTest : public testing::Test {
  public:
+  BudgetManagerTest() : origin_(url::Origin(GURL(kTestOrigin))) {}
   ~BudgetManagerTest() override {}
 
   BudgetManager* GetManager() {
     return BudgetManagerFactory::GetForProfile(&profile_);
   }
 
-  void SetSiteEngagementScore(const GURL& url, double score) {
+  void SetSiteEngagementScore(double score) {
     SiteEngagementService* service = SiteEngagementService::Get(&profile_);
-    service->ResetScoreForURL(url, score);
+    service->ResetScoreForURL(GURL(kTestOrigin), score);
   }
 
   Profile* profile() { return &profile_; }
+  const url::Origin origin() const { return origin_; }
 
   void StatusCallback(base::Closure run_loop_closure, bool success) {
     success_ = success;
@@ -46,10 +49,9 @@ class BudgetManagerTest : public testing::Test {
   }
 
   bool ReserveBudget(blink::mojom::BudgetOperationType type) {
-    const GURL origin(kTestOrigin);
     base::RunLoop run_loop;
     GetManager()->Reserve(
-        origin, type,
+        origin(), type,
         base::Bind(&BudgetManagerTest::StatusCallback, base::Unretained(this),
                    run_loop.QuitClosure()));
     run_loop.Run();
@@ -57,10 +59,9 @@ class BudgetManagerTest : public testing::Test {
   }
 
   bool ConsumeBudget(blink::mojom::BudgetOperationType type) {
-    const GURL origin(kTestOrigin);
     base::RunLoop run_loop;
     GetManager()->Consume(
-        origin, type,
+        origin(), type,
         base::Bind(&BudgetManagerTest::StatusCallback, base::Unretained(this),
                    run_loop.QuitClosure()));
     run_loop.Run();
@@ -73,13 +74,13 @@ class BudgetManagerTest : public testing::Test {
  private:
   content::TestBrowserThreadBundle thread_bundle_;
   TestingProfile profile_;
+  const url::Origin origin_;
 };
 
 TEST_F(BudgetManagerTest, GetBudgetConsumedOverTime) {
   // Set initial SES. The first time we try to spend budget, the
   // engagement award will be granted which is 48.0.
-  const GURL origin(kTestOrigin);
-  SetSiteEngagementScore(origin, kTestSES);
+  SetSiteEngagementScore(kTestSES);
   const blink::mojom::BudgetOperationType type =
       blink::mojom::BudgetOperationType::SILENT_PUSH;
 

chrome/browser/budget_service/budget_service_impl.cc

@@ -48,9 +48,7 @@ void BudgetServiceImpl::GetBudget(const url::Origin& origin,
 
   // Query the BudgetManager for the budget.
   content::BrowserContext* context = host->GetBrowserContext();
-  // TODO(harkness): Convert BudgetManager to use url::Origin.
-  const GURL gurl(origin.Serialize());
-  BudgetManagerFactory::GetForProfile(context)->GetBudget(gurl, callback);
+  BudgetManagerFactory::GetForProfile(context)->GetBudget(origin, callback);
 }
 
 void BudgetServiceImpl::Reserve(const url::Origin& origin,

chrome/browser/push_messaging/push_messaging_notification_manager.cc

@@ -192,7 +192,7 @@ void PushMessagingNotificationManager::DidGetNotificationsFromDatabase(
     // push was allowed.
     BudgetManager* manager = BudgetManagerFactory::GetForProfile(profile_);
     manager->Consume(
-        origin, blink::mojom::BudgetOperationType::SILENT_PUSH,
+        url::Origin(origin), blink::mojom::BudgetOperationType::SILENT_PUSH,
         base::Bind(&PushMessagingNotificationManager::ProcessSilentPush,
                    weak_factory_.GetWeakPtr(), origin,
                    service_worker_registration_id, message_handled_closure));