Let QUIC streams write 16k before ceding. Behind FLAG_quic_batch_writes.

net/quic/quic_session_test.cc

Index: net/quic/quic_session_test.cc
diff --git a/net/quic/quic_session_test.cc b/net/quic/quic_session_test.cc
index e4b67b0331d0b4a572551b12974bd02414cb5468..699d588b32634a316435af3f40a8fc9086bbbe43 100644
--- a/net/quic/quic_session_test.cc
+++ b/net/quic/quic_session_test.cc
@@ -96,6 +96,8 @@ class TestStream : public QuicSpdyStream {
     WriteOrBufferData(data, fin, nullptr);
   }
 
+  using QuicSpdyStream::set_priority;
+
   MOCK_METHOD0(OnCanWrite, void());
 };
 
@@ -111,6 +113,10 @@ class StreamBlocker {
     session_->MarkConnectionLevelWriteBlocked(stream_id_, kSomeMiddlePriority);
   }
 
+  void MarkHighPriorityWriteBlocked() {
+    session_->MarkConnectionLevelWriteBlocked(stream_id_, kHighestPriority);
+  }
+
  private:
   QuicSession* const session_;
   const QuicStreamId stream_id_;
@@ -158,13 +164,14 @@ class TestSession : public QuicSpdySession {
       bool fin,
       FecProtection fec_protection,
       QuicAckListenerInterface* ack_notifier_delegate) override {
-    // Always consumes everything.
-    if (writev_consumes_all_data_) {
-      return QuicConsumedData(data.total_length, fin);
-    } else {
-      return QuicSession::WritevData(id, data, offset, fin, fec_protection,
-                                     ack_notifier_delegate);
+    QuicConsumedData consumed(data.total_length, fin);
+    if (!writev_consumes_all_data_) {
+      consumed = QuicSession::WritevData(id, data, offset, fin, fec_protection,
+                                         ack_notifier_delegate);
     }
+    QuicSessionPeer::GetWriteBlockedStreams(this)->UpdateBytesForStream(
+        id, consumed.bytes_consumed);
+    return consumed;
   }
 
   void set_writev_consumes_all_data(bool val) {
@@ -177,6 +184,15 @@ class TestSession : public QuicSpdySession {
                       MAY_FEC_PROTECT, nullptr);
   }
 
+  QuicConsumedData SendLargeFakeData(QuicStreamId id, int bytes) {
+    DCHECK(writev_consumes_all_data_);
+    struct iovec iov;
+    iov.iov_base = nullptr;  // should not be read.
+    iov.iov_len = static_cast<size_t>(bytes);
+    return WritevData(id, QuicIOVector(&iov, 1, bytes), 0, true,
+                      MAY_FEC_PROTECT, nullptr);
+  }
+
   using QuicSession::PostProcessAfterData;
 
  private:
@@ -390,16 +406,102 @@ TEST_P(QuicSessionTestServer, OnCanWrite) {
 
   InSequence s;
   StreamBlocker stream2_blocker(&session_, stream2->id());
-  // Reregister, to test the loop limit.
-  EXPECT_CALL(*stream2, OnCanWrite())
-      .WillOnce(Invoke(&stream2_blocker,
-                       &StreamBlocker::MarkConnectionLevelWriteBlocked));
-  EXPECT_CALL(*stream6, OnCanWrite());
-  EXPECT_CALL(*stream4, OnCanWrite());
+
+  if (FLAGS_quic_batch_writes) {
+    // Reregister, to test the loop limit.
+    EXPECT_CALL(*stream2, OnCanWrite())
+        .WillOnce(Invoke(&stream2_blocker,
+                         &StreamBlocker::MarkConnectionLevelWriteBlocked));
+    // 2 will get called a second time as it didn't finish its block
+    EXPECT_CALL(*stream2, OnCanWrite());
+    EXPECT_CALL(*stream6, OnCanWrite());
+    // 4 will not get called, as we exceeded the loop limit.
+  } else {
+    // Reregister, to test the loop limit.
+    EXPECT_CALL(*stream2, OnCanWrite())
+        .WillOnce(Invoke(&stream2_blocker,
+                         &StreamBlocker::MarkConnectionLevelWriteBlocked));
+    EXPECT_CALL(*stream6, OnCanWrite());
+    EXPECT_CALL(*stream4, OnCanWrite());
+  }
   session_.OnCanWrite();
   EXPECT_TRUE(session_.WillingAndAbleToWrite());
 }
 
+TEST_P(QuicSessionTestServer, TestBatchedWrites) {
+  FLAGS_quic_batch_writes = true;
+  TestStream* stream2 = session_.CreateOutgoingDynamicStream();
+  TestStream* stream4 = session_.CreateOutgoingDynamicStream();
+  TestStream* stream6 = session_.CreateOutgoingDynamicStream();
+
+  session_.set_writev_consumes_all_data(true);
+  session_.MarkConnectionLevelWriteBlocked(stream2->id(), kSomeMiddlePriority);
+  session_.MarkConnectionLevelWriteBlocked(stream4->id(), kSomeMiddlePriority);
+
+  StreamBlocker stream2_blocker(&session_, stream2->id());
+  StreamBlocker stream4_blocker(&session_, stream4->id());
+  StreamBlocker stream6_blocker(&session_, stream6->id());
+  // With two sessions blocked, we should get two write calls.  They should both
+  // go to the first stream as it will only write 6k and mark itself blocked
+  // again.
+  InSequence s;
+  EXPECT_CALL(*stream2, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream2->id(), 6000))),
+                      Invoke(&stream2_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  EXPECT_CALL(*stream2, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream2->id(), 6000))),
+                      Invoke(&stream2_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  session_.OnCanWrite();
+
+  // We should get one more call for stream2, at which point it has used its
+  // write quota and we move over to stream 4.
+  EXPECT_CALL(*stream2, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream2->id(), 6000))),
+                      Invoke(&stream2_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  EXPECT_CALL(*stream4, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream4->id(), 6000))),
+                      Invoke(&stream4_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  session_.OnCanWrite();
+
+  // Now let stream 4 do the 2nd of its 3 writes, but add a block for a high
+  // priority stream 6.  4 should be preempted.  6 will write but *not* block so
+  // will cede back to 4.
+  stream6->set_priority(kHighestPriority);
+  EXPECT_CALL(*stream4, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream4->id(), 6000))),
+                      Invoke(&stream4_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked),
+                      Invoke(&stream6_blocker,
+                             &StreamBlocker::MarkHighPriorityWriteBlocked)));
+  EXPECT_CALL(*stream6, OnCanWrite())
+      .WillOnce(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream4->id(), 6000))));
+    session_.OnCanWrite();
+
+  // Stream4 alread did 6k worth of writes, so after doing another 12k it should
+  // cede and 2 should resume.
+  EXPECT_CALL(*stream4, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream4->id(), 12000))),
+                      Invoke(&stream4_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  EXPECT_CALL(*stream2, OnCanWrite())
+      .WillOnce(DoAll(testing::IgnoreResult(Invoke(CreateFunctor(
+          &session_, &TestSession::SendLargeFakeData, stream2->id(), 6000))),
+                      Invoke(&stream2_blocker,
+                             &StreamBlocker::MarkConnectionLevelWriteBlocked)));
+  session_.OnCanWrite();
+}
+
 TEST_P(QuicSessionTestServer, OnCanWriteBundlesStreams) {
   // Drive congestion control manually.
   MockSendAlgorithm* send_algorithm = new StrictMock<MockSendAlgorithm>;
@@ -516,13 +618,8 @@ TEST_P(QuicSessionTestServer, BufferedHandshake) {
   TestCryptoStream* crypto_stream = session_.GetCryptoStream();
   EXPECT_CALL(*crypto_stream, OnCanWrite());
 
-  // Re-register all other streams, to show they weren't able to proceed.
-  EXPECT_CALL(*stream2, OnCanWrite())
-      .WillOnce(Invoke(&stream2_blocker,
-                       &StreamBlocker::MarkConnectionLevelWriteBlocked));
-  EXPECT_CALL(*stream3, OnCanWrite())
-      .WillOnce(Invoke(&stream3_blocker,
-                       &StreamBlocker::MarkConnectionLevelWriteBlocked));
+  EXPECT_CALL(*stream2, OnCanWrite());
+  EXPECT_CALL(*stream3, OnCanWrite());
   EXPECT_CALL(*stream4, OnCanWrite())
       .WillOnce(Invoke(&stream4_blocker,
                        &StreamBlocker::MarkConnectionLevelWriteBlocked));