Implement THROTTLED priority semantics.

net/url_request/url_request_unittest.cc

 // Copyright (c) 2012 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 <memory>
 #include <utility>
 
 #include "base/memory/ptr_util.h"
 #include "base/run_loop.h"
 #include "build/build_config.h"
@@ skipping 367 matching lines @@
 
   // Behavior during blocked stages.  During other stages, just
   // returns OK or NetworkDelegate::AUTH_REQUIRED_RESPONSE_NO_ACTION.
   enum BlockMode {
     SYNCHRONOUS,    // No callback, returns specified return values.
     AUTO_CALLBACK,  // |this| posts a task to run the callback using the
                     // specified return codes.
     USER_CALLBACK,  // User takes care of doing a callback.  |retval_| and
                     // |auth_retval_| are ignored. In every blocking stage the
                     // message loop is quit.
+    USER_NOTIFY,    // User is notified by a provided callback of the
+                    // blocking, and synchronously returns instructions
+                    // for handling it.
   };
 
+  using NotificationCallback =
+      base::Callback<Error(const CompletionCallback&, const URLRequest*)>;
+
+  using NotificationAuthCallback =
+      base::Callback<NetworkDelegate::AuthRequiredResponse(const AuthCallback&,
+                                                           const URLRequest*)>;
+
   // Creates a delegate which does not block at all.
   explicit BlockingNetworkDelegate(BlockMode block_mode);
 
   // For users to trigger a callback returning |response|.
   // Side-effects: resets |stage_blocked_for_callback_| and stored callbacks.
   // Only call if |block_mode_| == USER_CALLBACK.
   void DoCallback(int response);
   void DoAuthCallback(NetworkDelegate::AuthRequiredResponse response);
 
   // Setters.
@@ skipping 16 matching lines @@
   }
 
   void set_redirect_url(const GURL& url) {
     redirect_url_ = url;
   }
 
   void set_block_on(int block_on) {
     block_on_ = block_on;
   }
 
+  // Only valid if |block_mode_| == USER_NOTIFY
+  void set_notification_callback(
+      const NotificationCallback& notification_callback) {
+    notification_callback_ = notification_callback;
+  }
+
+  void set_notification_auth_callback(
+      const NotificationAuthCallback& notification_auth_callback) {
+    notification_auth_callback_ = notification_auth_callback;
+  }
+
   // Allows the user to check in which state did we block.
   Stage stage_blocked_for_callback() const {
     EXPECT_EQ(USER_CALLBACK, block_mode_);
     return stage_blocked_for_callback_;
   }
 
  private:
   void RunCallback(int response, const CompletionCallback& callback);
   void RunAuthCallback(AuthRequiredResponse response,
                        const AuthCallback& callback);
@@ skipping 18 matching lines @@
       URLRequest* request,
       const AuthChallengeInfo& auth_info,
       const AuthCallback& callback,
       AuthCredentials* credentials) override;
 
   // Resets the callbacks and |stage_blocked_for_callback_|.
   void Reset();
 
   // Checks whether we should block in |stage|. If yes, returns an error code
   // and optionally sets up callback based on |block_mode_|. If no, returns OK.
-  int MaybeBlockStage(Stage stage, const CompletionCallback& callback);
+  int MaybeBlockStage(Stage stage,
+                      const URLRequest* request,
+                      const CompletionCallback& callback);
 
   // Configuration parameters, can be adjusted by public methods:
   const BlockMode block_mode_;
 
   // Values returned on blocking stages when mode is SYNCHRONOUS or
   // AUTO_CALLBACK. For USER_CALLBACK these are set automatically to IO_PENDING.
   int retval_;  // To be returned in non-auth stages.
   AuthRequiredResponse auth_retval_;
 
   GURL redirect_url_;  // Used if non-empty during OnBeforeURLRequest.
   int block_on_;  // Bit mask: in which stages to block.
 
   // |auth_credentials_| will be copied to |*target_auth_credential_| on
   // callback.
   AuthCredentials auth_credentials_;
   AuthCredentials* target_auth_credentials_;
 
   // Internal variables, not set by not the user:
   // Last blocked stage waiting for user callback (unused if |block_mode_| !=
   // USER_CALLBACK).
   Stage stage_blocked_for_callback_;
 
   // Callback objects stored during blocking stages.
   CompletionCallback callback_;
   AuthCallback auth_callback_;
 
+  // Callback to request user instructions for blocking.
+  NotificationCallback notification_callback_;
+  NotificationAuthCallback notification_auth_callback_;
+
   base::WeakPtrFactory<BlockingNetworkDelegate> weak_factory_;
 
   DISALLOW_COPY_AND_ASSIGN(BlockingNetworkDelegate);
 };
 
 BlockingNetworkDelegate::BlockingNetworkDelegate(BlockMode block_mode)
     : block_mode_(block_mode),
       retval_(OK),
       auth_retval_(AUTH_REQUIRED_RESPONSE_NO_ACTION),
       block_on_(0),
@@ skipping 39 matching lines @@
     const CompletionCallback& callback,
     GURL* new_url) {
   if (redirect_url_ == request->url())
     return OK;  // We've already seen this request and redirected elsewhere.
 
   TestNetworkDelegate::OnBeforeURLRequest(request, callback, new_url);
 
   if (!redirect_url_.is_empty())
     *new_url = redirect_url_;
 
-  return MaybeBlockStage(ON_BEFORE_URL_REQUEST, callback);
+  return MaybeBlockStage(ON_BEFORE_URL_REQUEST, request, callback);
 }
 
 int BlockingNetworkDelegate::OnBeforeStartTransaction(
     URLRequest* request,
     const CompletionCallback& callback,
     HttpRequestHeaders* headers) {
   TestNetworkDelegate::OnBeforeStartTransaction(request, callback, headers);
 
-  return MaybeBlockStage(ON_BEFORE_SEND_HEADERS, callback);
+  return MaybeBlockStage(ON_BEFORE_SEND_HEADERS, request, callback);
 }
 
 int BlockingNetworkDelegate::OnHeadersReceived(
     URLRequest* request,
     const CompletionCallback& callback,
     const HttpResponseHeaders* original_response_headers,
     scoped_refptr<HttpResponseHeaders>* override_response_headers,
     GURL* allowed_unsafe_redirect_url) {
   TestNetworkDelegate::OnHeadersReceived(request,
                                          callback,
                                          original_response_headers,
                                          override_response_headers,
                                          allowed_unsafe_redirect_url);
 
-  return MaybeBlockStage(ON_HEADERS_RECEIVED, callback);
+  return MaybeBlockStage(ON_HEADERS_RECEIVED, request, callback);
 }
 
 NetworkDelegate::AuthRequiredResponse BlockingNetworkDelegate::OnAuthRequired(
     URLRequest* request,
     const AuthChallengeInfo& auth_info,
     const AuthCallback& callback,
     AuthCredentials* credentials) {
   TestNetworkDelegate::OnAuthRequired(request, auth_info, callback,
                                       credentials);
   // Check that the user has provided callback for the previous blocked stage.
@@ skipping 17 matching lines @@
           base::Bind(&BlockingNetworkDelegate::RunAuthCallback,
                      weak_factory_.GetWeakPtr(), auth_retval_, callback));
       return AUTH_REQUIRED_RESPONSE_IO_PENDING;
 
     case USER_CALLBACK:
       auth_callback_ = callback;
       stage_blocked_for_callback_ = ON_AUTH_REQUIRED;
       base::ThreadTaskRunnerHandle::Get()->PostTask(
           FROM_HERE, base::MessageLoop::QuitWhenIdleClosure());
       return AUTH_REQUIRED_RESPONSE_IO_PENDING;
+
+    case USER_NOTIFY:
+      // If the callback returns ERR_IO_PENDING, the user has accepted
+      // responsibility for running the callback in the future.
+      return notification_auth_callback_.Run(callback, request);
   }
   NOTREACHED();
   return AUTH_REQUIRED_RESPONSE_NO_ACTION;  // Dummy value.
 }
 
 void BlockingNetworkDelegate::Reset() {
   EXPECT_NE(NOT_BLOCKED, stage_blocked_for_callback_);
   stage_blocked_for_callback_ = NOT_BLOCKED;
   callback_.Reset();
   auth_callback_.Reset();
 }
 
 int BlockingNetworkDelegate::MaybeBlockStage(
     BlockingNetworkDelegate::Stage stage,
+    const URLRequest* request,
     const CompletionCallback& callback) {
   // Check that the user has provided callback for the previous blocked stage.
   EXPECT_EQ(NOT_BLOCKED, stage_blocked_for_callback_);
 
   if ((block_on_ & stage) == 0) {
     return OK;
   }
 
   switch (block_mode_) {
     case SYNCHRONOUS:
       EXPECT_NE(OK, retval_);
       return retval_;
 
     case AUTO_CALLBACK:
       base::ThreadTaskRunnerHandle::Get()->PostTask(
           FROM_HERE, base::Bind(&BlockingNetworkDelegate::RunCallback,
                                 weak_factory_.GetWeakPtr(), retval_, callback));
       return ERR_IO_PENDING;
 
     case USER_CALLBACK:
       callback_ = callback;
       stage_blocked_for_callback_ = stage;
       base::ThreadTaskRunnerHandle::Get()->PostTask(
           FROM_HERE, base::MessageLoop::QuitWhenIdleClosure());
       return ERR_IO_PENDING;
+
+    case USER_NOTIFY:
+      // If the callback returns ERR_IO_PENDING, the user has accepted
+      // responsibility for running the callback in the future.
+      return notification_callback_.Run(callback, request);
   }
   NOTREACHED();
   return 0;
 }
 
 class TestURLRequestContextWithProxy : public TestURLRequestContext {
  public:
   // Does not own |delegate|.
   TestURLRequestContextWithProxy(const std::string& proxy,
                                  NetworkDelegate* delegate)
@@ skipping 7175 matching lines @@
   std::unique_ptr<URLRequest> req(
       default_context_.CreateRequest(test_url, DEFAULT_PRIORITY, &d));
   req->SetLoadFlags(LOAD_ONLY_FROM_CACHE);
 
   req->Start();
   base::RunLoop().Run();
 
   EXPECT_FALSE(req->response_info().network_accessed);
 }
 
-// Test that a single job with a throttled priority completes
+// Test that a single job with a THROTTLED priority completes
 // correctly in the absence of contention.
 TEST_F(URLRequestTestHTTP, ThrottledPriority) {
   ASSERT_TRUE(http_test_server()->Start());
 
   TestDelegate d;
   GURL test_url(http_test_server()->GetURL("/"));
   std::unique_ptr<URLRequest> req(
       default_context_.CreateRequest(test_url, THROTTLED, &d));
   req->Start();
   base::RunLoop().Run();
 
   EXPECT_TRUE(req->status().is_success());
 }
 
+// A class to hold state for responding to USER_NOTIFY callbacks from
+// BlockingNetworkDelegate.  It also accepts a RunLoop that will be
+// signaled via QuitWhenIdle() when any request is blocked.
+//
+class NotificationCallbackHandler {
+ public:
+  // Default constructed object doesn't block anything.
+  NotificationCallbackHandler() : run_loop_(nullptr) {}
+
+  void AddURLRequestToBlockList(const URLRequest* request) {
+    requests_to_block_.insert(request);
+  }
+
+  Error ShouldBlockRequest(const CompletionCallback& callback,
+                           const URLRequest* request) {
+    if (requests_to_block_.find(request) == requests_to_block_.end()) {
+      return OK;
+    }
+
+    DCHECK(blocked_callbacks_.find(request) == blocked_callbacks_.end());
+    blocked_callbacks_[request] = callback;
+    if (run_loop_)
+      run_loop_->QuitWhenIdle();
+    return ERR_IO_PENDING;
+  }
+
+  // Erases object's memory of blocked callbacks as a side effect.
+  void GetBlockedCallbacks(
+      std::map<const URLRequest*, CompletionCallback>* blocked_callbacks) {
+    blocked_callbacks_.swap(*blocked_callbacks);
+  }
+
+  // Set a RunLoop that, if non-null, will be signaled if any request
+  // is blocked.  It is the callers responsibility to make sure the
+  // passed object lives past the destruction of this class or
+  // next call to SetRunLoop().
+  void SetRunLoop(base::RunLoop* run_loop) { run_loop_ = run_loop; }
+
+ private:
+  std::set<const URLRequest*> requests_to_block_;
+  std::map<const URLRequest*, CompletionCallback> blocked_callbacks_;
+
+  base::RunLoop* run_loop_;
+
+  DISALLOW_COPY_AND_ASSIGN(NotificationCallbackHandler);
+};
+
+TEST_F(URLRequestTestHTTP, MultiThrottledPriority) {
+  ASSERT_TRUE(http_test_server()->Start());
+
+  base::RunLoop run_until_request_blocked;
+
+  NotificationCallbackHandler notification_handler;
+  notification_handler.SetRunLoop(&run_until_request_blocked);
+  BlockingNetworkDelegate network_delegate(
+      BlockingNetworkDelegate::USER_NOTIFY);
+  network_delegate.set_block_on(BlockingNetworkDelegate::ON_HEADERS_RECEIVED);
+  network_delegate.set_notification_callback(
+      base::Bind(&NotificationCallbackHandler::ShouldBlockRequest,
+                 // Both objects are owned by this function, and
+                 // |*network_delegate| will be destroyed first, so
+                 // it's safe to pass it an unretained pointer.
+                 base::Unretained(&notification_handler)));
+
+  TestURLRequestContext context(true);
+  context.set_network_delegate(&network_delegate);
+  context.Init();
+
+  // Use different test URLs to make sure all three requests turn into
+  // HttpNetworkTransacations.  Use different URLRequest::Delegates so that
+  // the requests may be waited on separately.
+  TestDelegate d1;
+  std::unique_ptr<URLRequest> req1(context.CreateRequest(
+      http_test_server()->GetURL("/echoall/l"), THROTTLED, &d1));
+  notification_handler.AddURLRequestToBlockList(req1.get());
+
+  TestDelegate d2;
+  std::unique_ptr<URLRequest> req2(context.CreateRequest(
+      http_test_server()->GetURL("/echoall/2"), THROTTLED, &d2));
+  notification_handler.AddURLRequestToBlockList(req2.get());
+
+  TestDelegate d3;
+  std::unique_ptr<URLRequest> req3(context.CreateRequest(
+      http_test_server()->GetURL("/echoall/3"), THROTTLED, &d3));
+  req1->Start();
+  req2->Start();
+  req3->Start();
+  run_until_request_blocked.Run();
+  notification_handler.SetRunLoop(nullptr);
+
+  // The first two requests should be blocked based on the notification
+  // callback, and their status should have blocked the third request
+  // through throttling.
+  EXPECT_TRUE(req1->status().is_io_pending());
+  EXPECT_TRUE(req2->status().is_io_pending());
+  EXPECT_TRUE(req3->status().is_io_pending());
+
+  std::map<const URLRequest*, CompletionCallback> blocked_callbacks;
+  notification_handler.GetBlockedCallbacks(&blocked_callbacks);
+  ASSERT_EQ(2u, blocked_callbacks.size());
+  ASSERT_TRUE(blocked_callbacks.find(req1.get()) != blocked_callbacks.end());
+  ASSERT_TRUE(blocked_callbacks.find(req2.get()) != blocked_callbacks.end());
+
+  // Unblocking one of the requests blocked on the notification callback
+  // should let it completely, which should then let the third request
+  // complete.  Unblock the second request, then wait for the third
+  // request to complete.
+  // TODO(rdsmith): Find something to wait on other than the third
+  // requests completion; if there's a bug in throttling, that will
+  // result in this test hanging rather than failing quickly.
+  d1.set_quit_on_complete(false);
+  d2.set_quit_on_complete(false);
+  d3.set_quit_on_complete(true);
+  blocked_callbacks[req2.get()].Run(OK);
+  base::RunLoop().Run();
+
+  notification_handler.GetBlockedCallbacks(&blocked_callbacks);
+  EXPECT_EQ(0u, blocked_callbacks.size());
+  EXPECT_TRUE(req1->status().is_io_pending());
+  EXPECT_EQ(URLRequestStatus::SUCCESS, req2->status().status());
+  EXPECT_EQ(URLRequestStatus::SUCCESS, req3->status().status());
+}
+
 TEST_F(URLRequestTestHTTP, RawBodyBytesNoContentEncoding) {
   ASSERT_TRUE(http_test_server()->Start());
 
   TestDelegate d;
   std::unique_ptr<URLRequest> req(default_context().CreateRequest(
       http_test_server()->GetURL("/simple.html"), DEFAULT_PRIORITY, &d));
   req->Start();
   base::RunLoop().Run();
 
   EXPECT_EQ(5, req->GetRawBodyBytes());
@@ skipping 2398 matching lines @@
   AddTestInterceptor()->set_main_intercept_job(std::move(job));
 
   req->Start();
   req->Cancel();
   base::RunLoop().RunUntilIdle();
   EXPECT_EQ(URLRequestStatus::CANCELED, req->status().status());
   EXPECT_EQ(0, d.received_redirect_count());
 }
 
 }  // namespace net