Implementation of network level throttler.

net/http/http_network_transaction_unittest.cc

Index: net/http/http_network_transaction_unittest.cc
diff --git a/net/http/http_network_transaction_unittest.cc b/net/http/http_network_transaction_unittest.cc
index 0d467c666b3eee3cb8c4869fa55031fd344d877b..9225256c740facb432f2dbfe6e1a751b3680c3ea 100644
--- a/net/http/http_network_transaction_unittest.cc
+++ b/net/http/http_network_transaction_unittest.cc
@@ -10,6 +10,7 @@
 
 #include <limits>
 #include <memory>
+#include <set>
 #include <string>
 #include <utility>
 #include <vector>
@@ -34,6 +35,7 @@
 #include "net/base/load_timing_info.h"
 #include "net/base/load_timing_info_test_util.h"
 #include "net/base/net_errors.h"
+#include "net/base/network_stream_throttler.h"
 #include "net/base/proxy_delegate.h"
 #include "net/base/request_priority.h"
 #include "net/base/test_completion_callback.h"
@@ -115,6 +117,100 @@ enum TestCase {
   kTestCaseHTTP2PriorityDependencies
 };
 
+class TestNetworkStreamThrottler : public NetworkStreamThrottler {
+ public:
+  TestNetworkStreamThrottler()
+      : throttle_new_requests_(false),
+        num_set_priority_calls_(0),
+        last_priority_set_(IDLE) {}
+
+  ~TestNetworkStreamThrottler() override {
+    EXPECT_TRUE(outstanding_throttles_.empty());
+  }
+
+  // NetworkStreamThrottler
+  std::unique_ptr<Throttle> CreateThrottle(Delegate* delegate,
+                                           RequestPriority priority,
+                                           bool ignore_limits) override {
+    std::unique_ptr<TestThrottle> test_throttle(
+        new TestThrottle(throttle_new_requests_, delegate, this));
+    outstanding_throttles_.insert(test_throttle.get());
+    return std::move(test_throttle);
+  }
+
+  void UnthrottleAllRequests() {
+    std::set<TestThrottle*> outstanding_throttles_copy(outstanding_throttles_);
+    for (auto& throttle : outstanding_throttles_copy) {
+      if (throttle->IsThrottled())
+        throttle->Unthrottle();
+    }
+  }
+
+  void set_throttle_new_requests(bool throttle_new_requests) {
+    throttle_new_requests_ = throttle_new_requests;
+  }
+
+  // Includes both throttled and unthrottled throttles.
+  size_t num_outstanding_requests() const {
+    return outstanding_throttles_.size();
+  }
+
+  int num_set_priority_calls() const { return num_set_priority_calls_; }
+  RequestPriority last_priority_set() const { return last_priority_set_; }
+  void set_priority_change_closure(
+      const base::Closure& priority_change_closure) {
+    priority_change_closure_ = priority_change_closure;
+  }
+
+ private:
+  class TestThrottle : public NetworkStreamThrottler::Throttle {
+   public:
+    ~TestThrottle() override { throttler_->OnThrottleDestroyed(this); }
+
+    // Throttle
+    bool IsThrottled() const override { return throttled_; }
+    void SetPriority(RequestPriority priority) override {
+      throttler_->SetPriorityCalled(priority);
+    }
+
+    TestThrottle(bool throttled,
+                 Delegate* delegate,
+                 TestNetworkStreamThrottler* throttler)
+        : throttled_(throttled), delegate_(delegate), throttler_(throttler) {}
+
+    void Unthrottle() {
+      EXPECT_TRUE(throttled_);
+
+      throttled_ = false;
+      delegate_->OnThrottleStateChanged();
+    }
+
+    bool throttled_;
+    Delegate* delegate_;
+    TestNetworkStreamThrottler* throttler_;
+  };
+
+  void OnThrottleDestroyed(TestThrottle* throttle) {
+    outstanding_throttles_.erase(throttle);
+  }
+
+  void SetPriorityCalled(RequestPriority priority) {
+    ++num_set_priority_calls_;
+    last_priority_set_ = priority;
+    if (!priority_change_closure_.is_null())
+      priority_change_closure_.Run();
+  }
+
+  // Includes both throttled and unthrottled throttles.
+  std::set<TestThrottle*> outstanding_throttles_;
+  bool throttle_new_requests_;
+  int num_set_priority_calls_;
+  RequestPriority last_priority_set_;
+  base::Closure priority_change_closure_;
+
+  DISALLOW_COPY_AND_ASSIGN(TestNetworkStreamThrottler);
+};
+
 const base::string16 kBar(ASCIIToUTF16("bar"));
 const base::string16 kBar2(ASCIIToUTF16("bar2"));
 const base::string16 kBar3(ASCIIToUTF16("bar3"));
@@ -260,6 +356,24 @@ std::unique_ptr<HttpNetworkSession> CreateSession(
   return SpdySessionDependencies::SpdyCreateSession(session_deps);
 }
 
+// Note that the pointer written into |*throttler| will only be valid
+// for the lifetime of the returned HttpNetworkSession.
+std::unique_ptr<HttpNetworkSession> CreateSessionWithThrottler(
+    SpdySessionDependencies* session_deps,
+    TestNetworkStreamThrottler** throttler) {
+  std::unique_ptr<HttpNetworkSession> session(
+      SpdySessionDependencies::SpdyCreateSession(session_deps));
+
+  std::unique_ptr<TestNetworkStreamThrottler> owned_throttler(
+      new TestNetworkStreamThrottler);
+  *throttler = owned_throttler.get();
+
+  HttpNetworkSessionPeer peer(session.get());
+  peer.SetNetworkStreamThrottler(std::move(owned_throttler));
+
+  return session;
+}
+
 }  // namespace
 
 class HttpNetworkTransactionTest
@@ -15941,6 +16055,299 @@ TEST_P(HttpNetworkTransactionTest, DisableNPN) {
   EXPECT_TRUE(trans.server_ssl_config_.npn_protos.empty());
 }
 
+// Confirm that transactions with throttler is created in the unthrottled
+// state (and stays in that state) is not blocked.
+TEST_P(HttpNetworkTransactionTest, ThrottlingUnthrottled) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(DEFAULT_PRIORITY, session.get()));
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  rv = callback.WaitForResult();
+  EXPECT_EQ(OK, rv);
+}
+
+// Confirm requests can be blocked by a throttler, and are resumed
+// when the throttle is unblocked.
+TEST_P(HttpNetworkTransactionTest, ThrottlingBasic) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  // Start a request that will be throttled at start; confirm doesn't continue.
+  throttler->set_throttle_new_requests(true);
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(DEFAULT_PRIORITY, session.get()));
+
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(LOAD_STATE_THROTTLED, trans->GetLoadState());
+  EXPECT_FALSE(callback.have_result());
+
+  // Confirm the request goes on to complete when unthrottled.
+  throttler->UnthrottleAllRequests();
+  base::MessageLoop::current()->RunUntilIdle();
+  ASSERT_TRUE(callback.have_result());
+  EXPECT_EQ(OK, callback.WaitForResult());
+}
+
+// Destroy a request while it's throttled.
+TEST_P(HttpNetworkTransactionTest, ThrottlingDestruction) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  // Start a request that will be throttled at start; confirm doesn't continue.
+  throttler->set_throttle_new_requests(true);
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(DEFAULT_PRIORITY, session.get()));
+
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(LOAD_STATE_THROTTLED, trans->GetLoadState());
+  EXPECT_FALSE(callback.have_result());
+
+  EXPECT_EQ(1u, throttler->num_outstanding_requests());
+  trans.reset();
+  EXPECT_EQ(0u, throttler->num_outstanding_requests());
+}
+
+// Confirm the throttler receives SetPriority calls.
+TEST_P(HttpNetworkTransactionTest, ThrottlingPrioritySet) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(IDLE, session.get()));
+  throttler->set_throttle_new_requests(true);
+  // Start the transaction to associate a throttle with it.
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  EXPECT_EQ(0, throttler->num_set_priority_calls());
+  trans->SetPriority(LOW);
+  EXPECT_EQ(1, throttler->num_set_priority_calls());
+  EXPECT_EQ(LOW, throttler->last_priority_set());
+
+  throttler->UnthrottleAllRequests();
+  base::MessageLoop::current()->RunUntilIdle();
+  ASSERT_TRUE(callback.have_result());
+  EXPECT_EQ(OK, callback.WaitForResult());
+}
+
+// Confirm that unthrottling from a SetPriority call by the
+// throttler works properly.
+TEST_P(HttpNetworkTransactionTest, ThrottlingPrioritySetUnthrottle) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  StaticSocketDataProvider reads1(data_reads, arraysize(data_reads),
+                                  data_writes, arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads1);
+
+  // Start a request that will be throttled at start; confirm doesn't continue.
+  throttler->set_throttle_new_requests(true);
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(DEFAULT_PRIORITY, session.get()));
+
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(LOAD_STATE_THROTTLED, trans->GetLoadState());
+  EXPECT_FALSE(callback.have_result());
+
+  // Create a new request, call SetPriority on it to unthrottle,
+  // and make sure that allows the original request to complete.
+  std::unique_ptr<HttpTransaction> trans1(
+      new HttpNetworkTransaction(LOW, session.get()));
+  throttler->set_priority_change_closure(
+      base::Bind(&TestNetworkStreamThrottler::UnthrottleAllRequests,
+                 base::Unretained(throttler)));
+
+  // Start the transaction to associate a throttle with it.
+  TestCompletionCallback callback1;
+  rv = trans1->Start(&request, callback1.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  trans1->SetPriority(IDLE);
+
+  base::MessageLoop::current()->RunUntilIdle();
+  ASSERT_TRUE(callback.have_result());
+  EXPECT_EQ(OK, callback.WaitForResult());
+}
+
+// Transaction will be destroyed when the unique_ptr goes out of scope.
+void DestroyTransaction(std::unique_ptr<HttpTransaction> transaction) {}
+
+// Confirm that destroying a transaction from a SetPriority call by the
+// throttler works properly.
+TEST_P(HttpNetworkTransactionTest, ThrottlingPrioritySetDestroy) {
+  TestNetworkStreamThrottler* throttler(nullptr);
+  std::unique_ptr<HttpNetworkSession> session(
+      CreateSessionWithThrottler(&session_deps_, &throttler));
+
+  // Send a simple request and make sure it goes through.
+  HttpRequestInfo request;
+  request.method = "GET";
+  request.url = GURL("http://www.example.org/");
+  request.load_flags = 0;
+
+  MockWrite data_writes[] = {
+      MockWrite("GET / HTTP/1.1\r\n"
+                "Host: www.example.org\r\n"
+                "Connection: keep-alive\r\n\r\n"),
+  };
+  MockRead data_reads[] = {
+      MockRead("HTTP/1.0 200 OK\r\n\r\n"), MockRead("hello world"),
+      MockRead(SYNCHRONOUS, OK),
+  };
+  StaticSocketDataProvider reads(data_reads, arraysize(data_reads), data_writes,
+                                 arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads);
+
+  StaticSocketDataProvider reads1(data_reads, arraysize(data_reads),
+                                  data_writes, arraysize(data_writes));
+  session_deps_.socket_factory->AddSocketDataProvider(&reads1);
+
+  // Start a request that will be throttled at start; confirm doesn't continue.
+  throttler->set_throttle_new_requests(true);
+  std::unique_ptr<HttpTransaction> trans(
+      new HttpNetworkTransaction(DEFAULT_PRIORITY, session.get()));
+
+  TestCompletionCallback callback;
+  int rv = trans->Start(&request, callback.callback(), BoundNetLog());
+  EXPECT_EQ(ERR_IO_PENDING, rv);
+
+  base::MessageLoop::current()->RunUntilIdle();
+  EXPECT_EQ(LOAD_STATE_THROTTLED, trans->GetLoadState());
+  EXPECT_FALSE(callback.have_result());
+
+  // Arrange fo the set priority call on the above transaction to delete
+  // the transaction.
+  HttpTransaction* trans_p(trans.get());
+  throttler->set_priority_change_closure(
+      base::Bind(&DestroyTransaction, base::Passed(&trans)));
+
+  // Call it and check results (partially a "doesn't crash" test).
+  trans_p->SetPriority(IDLE);
+  trans_p = nullptr;  // No longer a valid pointer.
+
+  base::MessageLoop::current()->RunUntilIdle();
+  ASSERT_FALSE(callback.have_result());
+}
+
 #if !defined(OS_IOS)
 TEST_P(HttpNetworkTransactionTest, TokenBindingSpdy) {
   const std::string https_url = "https://www.example.com";