Separate construction from start for SafeBrowsingTokenFetcher

Real time URL lookups may be running multiple token fetches at once, and
needs to manage these requests in parallel. This CL moves that
functionality into the SafeBrowsingTokenFetcher.

Bug: 1051679, 1041912
Change-Id: I6a667ac86422e5c31c14107d45bf495cc0de4e8c
Reviewed-on: https://chromium-review.googlesource.com/c/chromium/src/+/2072911Reviewed-by: Daniel Rubery <drubery@chromium.org>
Reviewed-by: Xinghui Lu <xinghuilu@chromium.org>
Commit-Queue: Daniel Rubery <drubery@chromium.org>
Cr-Commit-Position: refs/heads/master@{#745160}

components/safe_browsing/core/browser/safe_browsing_token_fetcher.cc

@@ -12,6 +12,7 @@
 #include "components/safe_browsing/core/common/thread_utils.h"
 #include "components/signin/public/identity_manager/access_token_fetcher.h"
 #include "components/signin/public/identity_manager/access_token_info.h"
+#include "components/signin/public/identity_manager/consent_level.h"
 #include "components/signin/public/identity_manager/identity_manager.h"
 #include "google_apis/gaia/core_account_id.h"
 #include "google_apis/gaia/google_service_auth_error.h"
@@ -28,45 +29,63 @@ const int kTimeoutDelaySeconds = 5 * 60;
 }  // namespace
 
 SafeBrowsingTokenFetcher::SafeBrowsingTokenFetcher(
-    signin::IdentityManager* identity_manager,
-    signin::ConsentLevel consent,
-    Callback callback)
-    : callback_(std::move(callback)), weak_ptr_factory_(this) {
+    signin::IdentityManager* identity_manager)
+    : identity_manager_(identity_manager),
+      requests_sent_(0),
+      weak_ptr_factory_(this) {
   DCHECK(CurrentlyOnThread(ThreadID::UI));
-  CoreAccountId account_id = identity_manager->GetPrimaryAccountId(consent);
-  token_fetcher_ = identity_manager->CreateAccessTokenFetcherForAccount(
+}
+
+SafeBrowsingTokenFetcher::~SafeBrowsingTokenFetcher() {
+  for (auto& id_and_callback : callbacks_) {
+    std::move(id_and_callback.second).Run(base::nullopt);
+  }
+}
+
+void SafeBrowsingTokenFetcher::Start(signin::ConsentLevel consent_level,
+                                     Callback callback) {
+  DCHECK(CurrentlyOnThread(ThreadID::UI));
+  const int request_id = requests_sent_;
+  requests_sent_++;
+  CoreAccountId account_id =
+      identity_manager_->GetPrimaryAccountId(consent_level);
+  callbacks_[request_id] = std::move(callback);
+  token_fetchers_[request_id] =
+      identity_manager_->CreateAccessTokenFetcherForAccount(
           account_id, "safe_browsing_service", {kAPIScope},
           base::BindOnce(&SafeBrowsingTokenFetcher::OnTokenFetched,
-                     weak_ptr_factory_.GetWeakPtr()),
+                         weak_ptr_factory_.GetWeakPtr(), request_id),
           signin::AccessTokenFetcher::Mode::kImmediate);
   base::PostDelayedTask(
       FROM_HERE, CreateTaskTraits(ThreadID::UI),
       base::BindOnce(&SafeBrowsingTokenFetcher::OnTokenTimeout,
-                     weak_ptr_factory_.GetWeakPtr()),
+                     weak_ptr_factory_.GetWeakPtr(), request_id),
       base::TimeDelta::FromSeconds(kTimeoutDelaySeconds));
 }
 
-SafeBrowsingTokenFetcher::~SafeBrowsingTokenFetcher() = default;
-
 void SafeBrowsingTokenFetcher::OnTokenFetched(
+    int request_id,
     GoogleServiceAuthError error,
     signin::AccessTokenInfo access_token_info) {
   if (error.state() == GoogleServiceAuthError::NONE)
-    Finish(access_token_info);
+    Finish(request_id, access_token_info);
   else
-    Finish(base::nullopt);
+    Finish(request_id, base::nullopt);
 }
 
-void SafeBrowsingTokenFetcher::OnTokenTimeout() {
-  Finish(base::nullopt);
+void SafeBrowsingTokenFetcher::OnTokenTimeout(int request_id) {
+  Finish(request_id, base::nullopt);
 }
 
 void SafeBrowsingTokenFetcher::Finish(
+    int request_id,
     base::Optional<signin::AccessTokenInfo> token_info) {
-  std::move(callback_).Run(token_info);
+  if (callbacks_.contains(request_id)) {
+    std::move(callbacks_[request_id]).Run(token_info);
+  }
 
-  // Cancel any pending callbacks
-  weak_ptr_factory_.InvalidateWeakPtrs();
+  token_fetchers_.erase(request_id);
+  callbacks_.erase(request_id);
 }
 
 }  // namespace safe_browsing

components/safe_browsing/core/browser/safe_browsing_token_fetcher.h

@@ -8,6 +8,7 @@
 #include <memory>
 
 #include "base/callback.h"
+#include "base/containers/flat_map.h"
 #include "base/memory/weak_ptr.h"
 #include "base/optional.h"
 #include "components/signin/public/identity_manager/access_token_info.h"
@@ -30,23 +31,37 @@ class SafeBrowsingTokenFetcher {
   using Callback =
       base::OnceCallback<void(base::Optional<signin::AccessTokenInfo>)>;
 
-  // Create a SafeBrowsingTokenFetcher and immediately begin fetching the token
-  // for the primary account of |identity_manager|, with consent level
-  // |consent|. |callback| will be called with the result.
-  SafeBrowsingTokenFetcher(signin::IdentityManager* identity_manager,
-                           signin::ConsentLevel consent,
-                           Callback callback);
+  // Create a SafeBrowsingTokenFetcher for the primary account of
+  // |identity_manager|. |identity_manager| is unowned, and must outlive this
+  // object.
+  explicit SafeBrowsingTokenFetcher(signin::IdentityManager* identity_manager);
 
   ~SafeBrowsingTokenFetcher();
 
+  // Begin fetching a token for the account with the given |consent_level|. The
+  // result will be returned in |callback|. Must be called on the UI thread.
+  void Start(signin::ConsentLevel consent_level, Callback callback);
+
  private:
-  void OnTokenFetched(GoogleServiceAuthError error,
+  void OnTokenFetched(int request_id,
+                      GoogleServiceAuthError error,
                       signin::AccessTokenInfo access_token_info);
-  void OnTokenTimeout();
-  void Finish(base::Optional<signin::AccessTokenInfo> token_info);
+  void OnTokenTimeout(int request_id);
+  void Finish(int request_id,
+              base::Optional<signin::AccessTokenInfo> token_info);
+
+  // Reference to the identity manager to fetch from.
+  signin::IdentityManager* identity_manager_;
+
+  // The count of requests sent. This is used as an ID for requests.
+  int requests_sent_;
+
+  // Active fetchers, keyed by ID.
+  base::flat_map<int, std::unique_ptr<signin::AccessTokenFetcher>>
+      token_fetchers_;
 
-  std::unique_ptr<signin::AccessTokenFetcher> token_fetcher_;
-  Callback callback_;
+  // Active callbacks, keyed by ID.
+  base::flat_map<int, Callback> callbacks_;
 
   base::WeakPtrFactory<SafeBrowsingTokenFetcher> weak_ptr_factory_;
 };

components/safe_browsing/core/browser/safe_browsing_token_fetcher_unittest.cc

@@ -28,8 +28,8 @@ TEST_F(SafeBrowsingTokenFetcherTest, Success) {
       "test@example.com");
   base::Optional<signin::AccessTokenInfo> maybe_account_info;
   SafeBrowsingTokenFetcher fetcher(
-      identity_test_environment_.identity_manager(),
-      signin::ConsentLevel::kNotRequired,
+      identity_test_environment_.identity_manager());
+  fetcher.Start(signin::ConsentLevel::kNotRequired,
                 base::BindOnce(
                     [](base::Optional<signin::AccessTokenInfo>* target_info,
                        base::Optional<signin::AccessTokenInfo> info) {
@@ -48,8 +48,8 @@ TEST_F(SafeBrowsingTokenFetcherTest, Failure) {
       "test@example.com");
   base::Optional<signin::AccessTokenInfo> maybe_account_info;
   SafeBrowsingTokenFetcher fetcher(
-      identity_test_environment_.identity_manager(),
-      signin::ConsentLevel::kNotRequired,
+      identity_test_environment_.identity_manager());
+  fetcher.Start(signin::ConsentLevel::kNotRequired,
                 base::BindOnce(
                     [](base::Optional<signin::AccessTokenInfo>* target_info,
                        base::Optional<signin::AccessTokenInfo> info) {
@@ -67,8 +67,8 @@ TEST_F(SafeBrowsingTokenFetcherTest, NoSyncingAccount) {
       "test@example.com");
   base::Optional<signin::AccessTokenInfo> maybe_account_info;
   SafeBrowsingTokenFetcher fetcher(
-      identity_test_environment_.identity_manager(),
-      signin::ConsentLevel::kSync,
+      identity_test_environment_.identity_manager());
+  fetcher.Start(signin::ConsentLevel::kSync,
                 base::BindOnce(
                     [](base::Optional<signin::AccessTokenInfo>* target_info,
                        base::Optional<signin::AccessTokenInfo> info) {
@@ -85,8 +85,8 @@ TEST_F(SafeBrowsingTokenFetcherTest, SyncSuccess) {
   identity_test_environment_.MakePrimaryAccountAvailable("test@example.com");
   base::Optional<signin::AccessTokenInfo> maybe_account_info;
   SafeBrowsingTokenFetcher fetcher(
-      identity_test_environment_.identity_manager(),
-      signin::ConsentLevel::kSync,
+      identity_test_environment_.identity_manager());
+  fetcher.Start(signin::ConsentLevel::kSync,
                 base::BindOnce(
                     [](base::Optional<signin::AccessTokenInfo>* target_info,
                        base::Optional<signin::AccessTokenInfo> info) {