When QPACK encoding, allow blocking references on a stream that is already blocked.

After this change, blocked_stream_count() would only be used in tests.  This CL
also takes the opportunity to replace it with a more expressive
stream_is_blocked() method.

gfe-relnote: n/a, change to QUIC v99-only code.  Protected by existing disabled gfe2_reloadable_flag_quic_enable_version_99.
PiperOrigin-RevId: 270385673
Change-Id: I2b68be71f2db19e2acbc830b3a4579436fe4fa39
diff --git a/quic/core/qpack/qpack_blocking_manager.cc b/quic/core/qpack/qpack_blocking_manager.cc
index 04d824f..f121aa0 100644
--- a/quic/core/qpack/qpack_blocking_manager.cc
+++ b/quic/core/qpack/qpack_blocking_manager.cc
@@ -71,18 +71,27 @@
   IncreaseReferenceCounts({referred_index});
 }
 
-uint64_t QpackBlockingManager::blocked_stream_count() const {
+bool QpackBlockingManager::blocking_allowed_on_stream(
+    QuicStreamId stream_id,
+    uint64_t maximum_blocked_streams) const {
   uint64_t blocked_stream_count = 0;
   for (const auto& header_blocks_for_stream : header_blocks_) {
     for (const IndexSet& indices : header_blocks_for_stream.second) {
       if (RequiredInsertCount(indices) > known_received_count_) {
+        if (header_blocks_for_stream.first == stream_id) {
+          // Sending blocking references is allowed if stream |stream_id| is
+          // already blocked.
+          return true;
+        }
         ++blocked_stream_count;
         break;
       }
     }
   }
 
-  return blocked_stream_count;
+  // Stream |stream_id| is not blocked, therefore sending blocking references on
+  // it would increase the blocked stream count by one.
+  return blocked_stream_count + 1 <= maximum_blocked_streams;
 }
 
 uint64_t QpackBlockingManager::smallest_blocking_index() const {
diff --git a/quic/core/qpack/qpack_blocking_manager.h b/quic/core/qpack/qpack_blocking_manager.h
index 7acfd5a..6bebf5f 100644
--- a/quic/core/qpack/qpack_blocking_manager.h
+++ b/quic/core/qpack/qpack_blocking_manager.h
@@ -15,6 +15,12 @@
 
 namespace quic {
 
+namespace test {
+
+class QpackBlockingManagerPeer;
+
+}  // namespace test
+
 // Class to keep track of blocked streams and blocking dynamic table entries:
 // https://quicwg.org/base-drafts/draft-ietf-quic-qpack.html#blocked-decoding
 // https://quicwg.org/base-drafts/draft-ietf-quic-qpack.html#blocked-insertion
@@ -47,8 +53,14 @@
   void OnReferenceSentOnEncoderStream(uint64_t inserted_index,
                                       uint64_t referred_index);
 
-  // Returns the number of blocked streams.
-  uint64_t blocked_stream_count() const;
+  // Returns true if sending blocking references on stream |stream_id| would not
+  // increase the total number of blocked streams above
+  // |maximum_blocked_streams|.  Note that if |stream_id| is already blocked
+  // then it is always allowed to send more blocking references on it.
+  // Behavior is undefined if |maximum_blocked_streams| is smaller than number
+  // of currently blocked streams.
+  bool blocking_allowed_on_stream(QuicStreamId stream_id,
+                                  uint64_t maximum_blocked_streams) const;
 
   // Returns the index of the blocking entry with the smallest index,
   // or std::numeric_limits<uint64_t>::max() if there are no blocking entries.
@@ -62,6 +74,8 @@
   static uint64_t RequiredInsertCount(const IndexSet& indices);
 
  private:
+  friend test::QpackBlockingManagerPeer;
+
   // A stream typically has only one header block, except for the rare cases of
   // 1xx responses, trailers, or push promises.  Even if there are multiple
   // header blocks sent on a single stream, they might not be blocked at the
diff --git a/quic/core/qpack/qpack_blocking_manager_test.cc b/quic/core/qpack/qpack_blocking_manager_test.cc
index 6f78d87..f21f1e7 100644
--- a/quic/core/qpack/qpack_blocking_manager_test.cc
+++ b/quic/core/qpack/qpack_blocking_manager_test.cc
@@ -8,6 +8,27 @@
 
 namespace quic {
 namespace test {
+
+class QpackBlockingManagerPeer {
+ public:
+  static bool stream_is_blocked(const QpackBlockingManager* manager,
+                                QuicStreamId stream_id) {
+    for (const auto& header_blocks_for_stream : manager->header_blocks_) {
+      if (header_blocks_for_stream.first != stream_id) {
+        continue;
+      }
+      for (const auto& indices : header_blocks_for_stream.second) {
+        if (QpackBlockingManager::RequiredInsertCount(indices) >
+            manager->known_received_count_) {
+          return true;
+        }
+      }
+    }
+
+    return false;
+  }
+};
+
 namespace {
 
 class QpackBlockingManagerTest : public QuicTest {
@@ -15,11 +36,14 @@
   QpackBlockingManagerTest() = default;
   ~QpackBlockingManagerTest() override = default;
 
+  bool stream_is_blocked(QuicStreamId stream_id) const {
+    return QpackBlockingManagerPeer::stream_is_blocked(&manager_, stream_id);
+  }
+
   QpackBlockingManager manager_;
 };
 
 TEST_F(QpackBlockingManagerTest, Empty) {
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
   EXPECT_EQ(0u, manager_.known_received_count());
   EXPECT_EQ(std::numeric_limits<uint64_t>::max(),
             manager_.smallest_blocking_index());
@@ -34,37 +58,39 @@
   // Stream 0 is not blocked, because it only references entries that are
   // already acknowledged by an Insert Count Increment instruction.
   manager_.OnHeaderBlockSent(0, {1, 0});
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
 }
 
 TEST_F(QpackBlockingManagerTest, UnblockedByInsertCountIncrement) {
   manager_.OnHeaderBlockSent(0, {1, 0});
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
 
   manager_.OnInsertCountIncrement(2);
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
 }
 
 TEST_F(QpackBlockingManagerTest, NotBlockedByHeaderAcknowledgement) {
   manager_.OnHeaderBlockSent(0, {2, 1, 1});
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
 
   EXPECT_TRUE(manager_.OnHeaderAcknowledgement(0));
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
 
   // Stream 1 is not blocked, because it only references entries that are
   // already acknowledged by a Header Acknowledgement instruction.
   manager_.OnHeaderBlockSent(1, {2, 2});
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(1));
 }
 
 TEST_F(QpackBlockingManagerTest, UnblockedByHeaderAcknowledgement) {
   manager_.OnHeaderBlockSent(0, {2, 1, 1});
   manager_.OnHeaderBlockSent(1, {2, 2});
-  EXPECT_EQ(2u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
+  EXPECT_TRUE(stream_is_blocked(1));
 
   EXPECT_TRUE(manager_.OnHeaderAcknowledgement(0));
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
+  EXPECT_FALSE(stream_is_blocked(1));
 }
 
 TEST_F(QpackBlockingManagerTest, KnownReceivedCount) {
@@ -145,38 +171,37 @@
 
 TEST_F(QpackBlockingManagerTest, HeaderAcknowledgementsOnSingleStream) {
   EXPECT_EQ(0u, manager_.known_received_count());
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
   EXPECT_EQ(std::numeric_limits<uint64_t>::max(),
             manager_.smallest_blocking_index());
 
   manager_.OnHeaderBlockSent(0, {2, 1, 1});
   EXPECT_EQ(0u, manager_.known_received_count());
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(1u, manager_.smallest_blocking_index());
 
   manager_.OnHeaderBlockSent(0, {1, 0});
   EXPECT_EQ(0u, manager_.known_received_count());
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(0u, manager_.smallest_blocking_index());
 
   EXPECT_TRUE(manager_.OnHeaderAcknowledgement(0));
   EXPECT_EQ(3u, manager_.known_received_count());
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
   EXPECT_EQ(0u, manager_.smallest_blocking_index());
 
   manager_.OnHeaderBlockSent(0, {3});
   EXPECT_EQ(3u, manager_.known_received_count());
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(0u, manager_.smallest_blocking_index());
 
   EXPECT_TRUE(manager_.OnHeaderAcknowledgement(0));
   EXPECT_EQ(3u, manager_.known_received_count());
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(3u, manager_.smallest_blocking_index());
 
   EXPECT_TRUE(manager_.OnHeaderAcknowledgement(0));
   EXPECT_EQ(4u, manager_.known_received_count());
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
   EXPECT_EQ(std::numeric_limits<uint64_t>::max(),
             manager_.smallest_blocking_index());
 
@@ -185,23 +210,26 @@
 
 TEST_F(QpackBlockingManagerTest, CancelStream) {
   manager_.OnHeaderBlockSent(0, {3});
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(3u, manager_.smallest_blocking_index());
 
   manager_.OnHeaderBlockSent(0, {2});
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
   EXPECT_EQ(2u, manager_.smallest_blocking_index());
 
   manager_.OnHeaderBlockSent(1, {4});
-  EXPECT_EQ(2u, manager_.blocked_stream_count());
+  EXPECT_TRUE(stream_is_blocked(0));
+  EXPECT_TRUE(stream_is_blocked(1));
   EXPECT_EQ(2u, manager_.smallest_blocking_index());
 
   manager_.OnStreamCancellation(0);
-  EXPECT_EQ(1u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
+  EXPECT_TRUE(stream_is_blocked(1));
   EXPECT_EQ(4u, manager_.smallest_blocking_index());
 
   manager_.OnStreamCancellation(1);
-  EXPECT_EQ(0u, manager_.blocked_stream_count());
+  EXPECT_FALSE(stream_is_blocked(0));
+  EXPECT_FALSE(stream_is_blocked(1));
   EXPECT_EQ(std::numeric_limits<uint64_t>::max(),
             manager_.smallest_blocking_index());
 }
@@ -288,6 +316,76 @@
             manager_.smallest_blocking_index());
 }
 
+TEST_F(QpackBlockingManagerTest, BlockingAllowedOnStream) {
+  const QuicStreamId kStreamId1 = 1;
+  const QuicStreamId kStreamId2 = 2;
+  const QuicStreamId kStreamId3 = 3;
+
+  // No stream can block if limit is 0.
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId1, 0));
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId2, 0));
+
+  // Either stream can block if limit is larger.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 1));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 1));
+
+  // Doubly block first stream.
+  manager_.OnHeaderBlockSent(kStreamId1, {0});
+  manager_.OnHeaderBlockSent(kStreamId1, {1});
+
+  // First stream is already blocked so it can carry more blocking references.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 1));
+  // Second stream is not allowed to block if limit is already reached.
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId2, 1));
+
+  // Either stream can block if limit is larger than number of blocked streams.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 2));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 2));
+
+  // Block second stream.
+  manager_.OnHeaderBlockSent(kStreamId2, {2});
+
+  // Streams are already blocked so either can carry more blocking references.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 2));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 2));
+
+  // Third, unblocked stream is not allowed to block unless limit is strictly
+  // larger than number of blocked streams.
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId3, 2));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId3, 3));
+
+  // Acknowledge decoding of first header block on first stream.
+  // Stream is still blocked on its second header block.
+  manager_.OnHeaderAcknowledgement(kStreamId1);
+
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 1));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 1));
+
+  // Acknowledge decoding of second header block on first stream.
+  // This unblocks the stream.
+  manager_.OnHeaderAcknowledgement(kStreamId1);
+
+  // First stream is not allowed to block if limit is already reached.
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId1, 1));
+  // Second stream is already blocked so it can carry more blocking references.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 1));
+
+  // Either stream can block if limit is larger than number of blocked streams.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 2));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 2));
+
+  // Unblock second stream.
+  manager_.OnHeaderAcknowledgement(kStreamId2);
+
+  // No stream can block if limit is 0.
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId1, 0));
+  EXPECT_FALSE(manager_.blocking_allowed_on_stream(kStreamId2, 0));
+
+  // Either stream can block if limit is larger.
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId1, 1));
+  EXPECT_TRUE(manager_.blocking_allowed_on_stream(kStreamId2, 1));
+}
+
 }  // namespace
 }  // namespace test
 }  // namespace quic
diff --git a/quic/core/qpack/qpack_encoder.cc b/quic/core/qpack/qpack_encoder.cc
index 24f12bb..61ee8fc 100644
--- a/quic/core/qpack/qpack_encoder.cc
+++ b/quic/core/qpack/qpack_encoder.cc
@@ -86,6 +86,7 @@
 }
 
 QpackEncoder::Instructions QpackEncoder::FirstPassEncode(
+    QuicStreamId stream_id,
     const spdy::SpdyHeaderBlock& header_list,
     QpackBlockingManager::IndexSet* referred_indices,
     QuicByteCount* encoder_stream_sent_byte_count) {
@@ -106,9 +107,8 @@
       header_table_.draining_index(kDrainingFraction);
   // Blocking references are allowed if the number of blocked streams is less
   // than the limit.
-  // TODO(b/112770235): Also allow blocking if given stream is already blocked.
-  const bool blocking_allowed =
-      maximum_blocked_streams_ > blocking_manager_.blocked_stream_count();
+  const bool blocking_allowed = blocking_manager_.blocking_allowed_on_stream(
+      stream_id, maximum_blocked_streams_);
 
   for (const auto& header : ValueSplittingHeaderList(&header_list)) {
     // These strings are owned by |header_list|.
@@ -314,8 +314,9 @@
   QpackBlockingManager::IndexSet referred_indices;
 
   // First pass: encode into |instructions|.
-  Instructions instructions = FirstPassEncode(header_list, &referred_indices,
-                                              encoder_stream_sent_byte_count);
+  Instructions instructions =
+      FirstPassEncode(stream_id, header_list, &referred_indices,
+                      encoder_stream_sent_byte_count);
 
   const uint64_t required_insert_count =
       referred_indices.empty()
diff --git a/quic/core/qpack/qpack_encoder.h b/quic/core/qpack/qpack_encoder.h
index d64d159..2ffbc99 100644
--- a/quic/core/qpack/qpack_encoder.h
+++ b/quic/core/qpack/qpack_encoder.h
@@ -131,7 +131,8 @@
   // field representations, with all dynamic table entries referred to with
   // absolute indices.  Returned Instructions object may have QuicStringPieces
   // pointing to strings owned by |*header_list|.
-  Instructions FirstPassEncode(const spdy::SpdyHeaderBlock& header_list,
+  Instructions FirstPassEncode(QuicStreamId stream_id,
+                               const spdy::SpdyHeaderBlock& header_list,
                                QpackBlockingManager::IndexSet* referred_indices,
                                QuicByteCount* encoder_stream_sent_byte_count);
 
diff --git a/quic/core/qpack/qpack_encoder_test.cc b/quic/core/qpack/qpack_encoder_test.cc
index f230942..b0464c8 100644
--- a/quic/core/qpack/qpack_encoder_test.cc
+++ b/quic/core/qpack/qpack_encoder_test.cc
@@ -387,10 +387,6 @@
 }
 
 TEST_F(QpackEncoderTest, Draining) {
-  // TODO(b/112770235): Remove when already blocking stream can emit blocking
-  // references.
-  encoder_.SetMaximumBlockedStreams(2);
-
   spdy::SpdyHeaderBlock header_list1;
   header_list1["one"] = "foo";
   header_list1["two"] = "foo";