[PH2][StreamQ] Adding priority for streams (grpc#40456)

Streams now track the priority in the writable stream list based on what type of messages have been enqueued.

Closes grpc#40456

COPYBARA_INTEGRATE_REVIEW=grpc#40456 from ac-patel:streamq8 9510b31
PiperOrigin-RevId: 797595198

Author: ac-patel

src/core/BUILD

@@ -8182,6 +8182,7 @@ grpc_cc_library(
         "message",
         "message_assembler",
         "metadata_batch",
+        "writable_streams",
     ],
 )
 

src/core/ext/transport/chttp2/transport/http2_client_transport.cc

@@ -62,6 +62,7 @@ namespace grpc_core {
 namespace http2 {
 
 using grpc_event_engine::experimental::EventEngine;
+using EnqueueResult = StreamDataQueue<ClientMetadataHandle>::EnqueueResult;
 
 // Experimental : This is just the initial skeleton of class
 // and it is functions. The code will be written iteratively.
@@ -781,7 +782,8 @@ auto Http2ClientTransport::StreamMultiplexerLoop() {
               LOG(ERROR) << "Failed to enqueue stream " << stream_id
                          << " with status: " << status;
               // Close transport if we fail to enqueue stream.
-              return absl::InternalError("Failed to enqueue stream");
+              return absl::UnavailableError(
+                  "Failed to enqueue stream to writable stream list");
             }
           } else if (GPR_UNLIKELY(!result.ok())) {
             // Close the corresponding stream if we fail to dequeue frames from
@@ -1137,12 +1139,12 @@ auto Http2ClientTransport::CallOutboundLoop(
         stream != nullptr,
         [self, stream, message = std::move(message), stream_id]() mutable {
           return TrySeq(stream->EnqueueMessage(std::move(message)),
-                        [self, stream_id](bool became_writable) {
+                        [self, stream_id](const EnqueueResult result) {
                           GRPC_HTTP2_CLIENT_DLOG
                               << "Http2ClientTransport CallOutboundLoop "
                                  "Enqueued Message";
                           return self->MaybeAddStreamToWritableStreamList(
-                              stream_id, became_writable);
+                              stream_id, result);
                         });
         },
         []() {
@@ -1160,12 +1162,12 @@ auto Http2ClientTransport::CallOutboundLoop(
         [self, stream, metadata = std::move(metadata), stream_id]() mutable {
           return TrySeq(
               stream->EnqueueInitialMetadata(std::move(metadata)),
-              [self, stream_id](bool became_writable) {
+              [self, stream_id](const EnqueueResult result) {
                 GRPC_HTTP2_CLIENT_DLOG
                     << "Http2ClientTransport CallOutboundLoop "
                        "Enqueued Initial Metadata";
-                return self->MaybeAddStreamToWritableStreamList(
-                    stream_id, became_writable);
+                return self->MaybeAddStreamToWritableStreamList(stream_id,
+                                                                result);
               },
               [stream] {
                 // TODO(akshitpatel) : [PH2][P2] : Think how to handle stream
@@ -1188,12 +1190,12 @@ auto Http2ClientTransport::CallOutboundLoop(
         stream != nullptr,
         [self, stream, stream_id]() mutable {
           return TrySeq(stream->EnqueueHalfClosed(),
-                        [self, stream_id](bool became_writable) {
+                        [self, stream_id](const EnqueueResult result) {
                           GRPC_HTTP2_CLIENT_DLOG
                               << "Http2ClientTransport CallOutboundLoop "
                                  "Enqueued Half Closed";
                           return self->MaybeAddStreamToWritableStreamList(
-                              stream_id, became_writable);
+                              stream_id, result);
                         });
         },
         []() {

src/core/ext/transport/chttp2/transport/http2_client_transport.h

@@ -739,18 +739,20 @@ class Http2ClientTransport final : public ClientTransport {
 
   WritableStreams writable_stream_list_;
 
-  auto MaybeAddStreamToWritableStreamList(const uint32_t stream_id,
-                                          const bool became_writable) {
-    if (became_writable) {
+  absl::Status MaybeAddStreamToWritableStreamList(
+      const uint32_t stream_id,
+      const StreamDataQueue<ClientMetadataHandle>::EnqueueResult result) {
+    if (result.became_writable) {
       GRPC_HTTP2_CLIENT_DLOG
           << "Http2ClientTransport MaybeAddStreamToWritableStreamList "
              " Stream id: "
           << stream_id << " became writable";
-      absl::Status status = writable_stream_list_.Enqueue(
-          stream_id, WritableStreams::StreamPriority::kDefault);
+      absl::Status status =
+          writable_stream_list_.Enqueue(stream_id, result.priority);
       if (!status.ok()) {
         return HandleError(Http2Status::Http2ConnectionError(
-            Http2ErrorCode::kProtocolError, "Failed to enqueue stream"));
+            Http2ErrorCode::kRefusedStream,
+            "Failed to enqueue stream to writable stream list"));
       }
     }
     return absl::OkStatus();

src/core/ext/transport/chttp2/transport/stream_data_queue.h

@@ -25,6 +25,7 @@
 #include "absl/log/log.h"
 #include "src/core/ext/transport/chttp2/transport/header_assembler.h"
 #include "src/core/ext/transport/chttp2/transport/message_assembler.h"
+#include "src/core/ext/transport/chttp2/transport/writable_streams.h"
 
 namespace grpc_core {
 namespace http2 {
@@ -36,11 +37,6 @@ namespace http2 {
 template <typename T>
 class SimpleQueue {
  public:
-  struct EnqueueResult {
-    absl::Status status;
-    bool became_non_empty;
-  };
-
   explicit SimpleQueue(const uint32_t max_tokens) : max_tokens_(max_tokens) {}
   SimpleQueue(SimpleQueue&& rhs) = delete;
   SimpleQueue& operator=(SimpleQueue&& rhs) = delete;
@@ -54,7 +50,7 @@ class SimpleQueue {
   // with tokens = 0. Enqueues with tokens = 0 are primarily for sending
   // metadata as flow control does not apply to them. This function is NOT
   // thread safe.
-  auto Enqueue(T& data, const uint32_t tokens) {
+  Poll<absl::StatusOr<bool>> Enqueue(T& data, const uint32_t tokens) {
     return PollEnqueue(data, tokens);
   }
 
@@ -80,7 +76,7 @@ class SimpleQueue {
   bool TestOnlyIsEmpty() const { return IsEmpty(); }
 
  private:
-  Poll<EnqueueResult> PollEnqueue(T& data, const uint32_t tokens) {
+  Poll<absl::StatusOr<bool>> PollEnqueue(T& data, const uint32_t tokens) {
     GRPC_STREAM_DATA_QUEUE_DEBUG << "Enqueueing data. Data tokens: " << tokens;
     const uint32_t max_tokens_consumed_threshold =
         max_tokens_ >= tokens ? max_tokens_ - tokens : 0;
@@ -91,8 +87,7 @@ class SimpleQueue {
       GRPC_STREAM_DATA_QUEUE_DEBUG
           << "Enqueue successful. Data tokens: " << tokens
           << " Current tokens consumed: " << tokens_consumed_;
-      return EnqueueResult{absl::OkStatus(),
-                           /*became_non_empty=*/queue_.size() == 1};
+      return /*became_non_empty*/ (queue_.size() == 1);
     }
 
     GRPC_STREAM_DATA_QUEUE_DEBUG
@@ -197,33 +192,40 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   //    end_stream set. If the stream needs to be half closed, the client should
   //    enqueue a half close message.
 
+  struct EnqueueResult {
+    bool became_writable;
+    WritableStreams::StreamPriority priority;
+  };
+
   // Enqueue Initial Metadata.
   // 1. MUST be called at most once.
   // 2. This MUST be called before any messages are enqueued.
   // 3. MUST not be called after trailing metadata is enqueued.
   // 4. This function is thread safe.
-  auto EnqueueInitialMetadata(MetadataHandle metadata) {
+  auto EnqueueInitialMetadata(MetadataHandle&& metadata) {
     DCHECK(!is_initial_metadata_queued_);
     DCHECK(!is_trailing_metadata_or_half_close_queued_);
     DCHECK(metadata != nullptr);
     DCHECK(!is_reset_stream_queued_);
 
     is_initial_metadata_queued_ = true;
     return [self = this->Ref(),
-            entry = QueueEntry{InitialMetadataType{
-                std::move(metadata)}}]() mutable -> Poll<absl::StatusOr<bool>> {
+            entry = QueueEntry{InitialMetadataType{std::move(
+                metadata)}}]() mutable -> Poll<absl::StatusOr<EnqueueResult>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, /*tokens=*/0);
+      Poll<absl::StatusOr<bool>> result =
+          self->queue_.Enqueue(entry, /*tokens=*/0);
       if (result.ready()) {
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Enqueued initial metadata for stream " << self->stream_id_
-            << " with status: " << result.value().status;
-        if (result.value().status.ok()) {
-          DCHECK(result.value().became_non_empty);
+            << " with status: " << result.value().status();
+        if (result.value().ok()) {
+          DCHECK(/*became_non_empty*/ result.value().value());
           return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
+              /*became_non_empty*/ result.value().value(),
+              WritableStreams::StreamPriority::kDefault);
         }
-        return result.value().status;
+        return result.value().status();
       }
       return Pending{};
     };
@@ -233,27 +235,29 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   // 1. MUST be called at most once.
   // 2. MUST be called only for a server.
   // 3. This function is thread safe.
-  auto EnqueueTrailingMetadata(MetadataHandle metadata) {
+  auto EnqueueTrailingMetadata(MetadataHandle&& metadata) {
     DCHECK(metadata != nullptr);
     DCHECK(!is_reset_stream_queued_);
     DCHECK(!is_client_);
     DCHECK(!is_trailing_metadata_or_half_close_queued_);
 
     is_trailing_metadata_or_half_close_queued_ = true;
     return [self = this->Ref(),
-            entry = QueueEntry{TrailingMetadataType{
-                std::move(metadata)}}]() mutable -> Poll<absl::StatusOr<bool>> {
+            entry = QueueEntry{TrailingMetadataType{std::move(
+                metadata)}}]() mutable -> Poll<absl::StatusOr<EnqueueResult>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, /*tokens=*/0);
+      Poll<absl::StatusOr<bool>> result =
+          self->queue_.Enqueue(entry, /*tokens=*/0);
       if (result.ready()) {
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Enqueued trailing metadata for stream " << self->stream_id_
-            << " with status: " << result.value().status;
-        if (result.value().status.ok()) {
+            << " with status: " << result.value().status();
+        if (result.value().ok()) {
           return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
+              /*became_non_empty*/ result.value().value(),
+              WritableStreams::StreamPriority::kStreamClosed);
         }
-        return result.value().status;
+        return result.value().status();
       }
       return Pending{};
     };
@@ -264,7 +268,7 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   // 1. MUST be called after initial metadata is enqueued.
   // 2. MUST not be called after trailing metadata is enqueued.
   // 3. This function is thread safe.
-  auto EnqueueMessage(MessageHandle message) {
+  auto EnqueueMessage(MessageHandle&& message) {
     DCHECK(is_initial_metadata_queued_);
     DCHECK(message != nullptr);
     DCHECK(!is_reset_stream_queued_);
@@ -275,19 +279,20 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
     const uint32_t tokens =
         message->payload()->Length() + kGrpcHeaderSizeInBytes;
     return [self = this->Ref(), entry = QueueEntry{std::move(message)},
-            tokens]() mutable -> Poll<absl::StatusOr<bool>> {
+            tokens]() mutable -> Poll<absl::StatusOr<EnqueueResult>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, tokens);
+      Poll<absl::StatusOr<bool>> result = self->queue_.Enqueue(entry, tokens);
       if (result.ready()) {
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Enqueued message for stream " << self->stream_id_
-            << " with status: " << result.value().status;
+            << " with status: " << result.value().status();
         // TODO(akshitpatel) : [PH2][P2] : Add check for flow control tokens.
-        if (result.value().status.ok()) {
+        if (result.value().ok()) {
           return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
+              /*became_non_empty*/ result.value().value(),
+              WritableStreams::StreamPriority::kDefault);
         }
-        return result.value().status;
+        return result.value().status();
       }
       return Pending{};
     };
@@ -305,18 +310,20 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
 
     is_trailing_metadata_or_half_close_queued_ = true;
     return [self = this->Ref(), entry = QueueEntry{HalfClosed{}}]() mutable
-               -> Poll<absl::StatusOr<bool>> {
+               -> Poll<absl::StatusOr<EnqueueResult>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, /*tokens=*/0);
+      Poll<absl::StatusOr<bool>> result =
+          self->queue_.Enqueue(entry, /*tokens=*/0);
       if (result.ready()) {
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Marking stream " << self->stream_id_ << " as half closed"
-            << " with status: " << result.value().status;
-        if (result.value().status.ok()) {
+            << " with status: " << result.value().status();
+        if (result.value().ok()) {
           return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
+              /*became_non_empty*/ result.value().value(),
+              WritableStreams::StreamPriority::kStreamClosed);
         }
-        return result.value().status;
+        return result.value().status();
       }
       return Pending{};
     };
@@ -332,18 +339,20 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
     is_reset_stream_queued_ = true;
     return [self = this->Ref(),
             entry = QueueEntry{ResetStream{
-                error_code}}]() mutable -> Poll<absl::StatusOr<bool>> {
+                error_code}}]() mutable -> Poll<absl::StatusOr<EnqueueResult>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, /*tokens=*/0);
+      Poll<absl::StatusOr<bool>> result =
+          self->queue_.Enqueue(entry, /*tokens=*/0);
       if (result.ready()) {
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Enqueueing reset stream for stream " << self->stream_id_
-            << " with status: " << result.value().status;
-        if (result.value().status.ok()) {
+            << " with status: " << result.value().status();
+        if (result.value().ok()) {
           return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
+              /*became_non_empty*/ result.value().value(),
+              WritableStreams::StreamPriority::kStreamClosed);
         }
-        return result.value().status;
+        return result.value().status();
       }
       return Pending{};
     };
@@ -357,6 +366,7 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   struct DequeueResult {
     std::vector<Http2Frame> frames;
     bool is_writable;
+    WritableStreams::StreamPriority priority;
   };
 
   // TODO(akshitpatel) : [PH2][P4] : Measure the performance of this function
@@ -411,7 +421,7 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
     GRPC_STREAM_DATA_QUEUE_DEBUG << "Stream id: " << stream_id_
                                  << " writable state changed to "
                                  << is_writable_;
-    return DequeueResult{handle_dequeue.GetFrames(), is_writable_};
+    return DequeueResult{handle_dequeue.GetFrames(), is_writable_, priority_};
   }
 
   // Returns true if the queue is empty. This function is thread safe.
@@ -572,17 +582,29 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
     HPackCompressor& encoder_;
   };
 
-  // Returns true if the queue is now writable. It is expected that the caller
-  // will hold the lock on the queue when calling this function.
-  bool UpdateWritableStateLocked(const bool became_non_empty)
+  // Updates the stream priority. Also sets the writable state to true if the
+  // stream has become writable. Returns if the stream became writable and
+  // updated priority. It is expected that the caller will hold the lock on the
+  // queue when calling this function.
+  EnqueueResult UpdateWritableStateLocked(
+      const bool became_non_empty,
+      const WritableStreams::StreamPriority priority)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    priority_ = priority;
     if (!is_writable_ && became_non_empty) {
-      GRPC_STREAM_DATA_QUEUE_DEBUG << "Stream id: " << stream_id_
-                                   << " writeable state changed to true";
       is_writable_ = true;
-      return true;
+      GRPC_STREAM_DATA_QUEUE_DEBUG
+          << "UpdateWritableStateLocked for stream id: " << stream_id_
+          << " became writable with priority: "
+          << WritableStreams::GetPriorityString(priority_);
+      return EnqueueResult{/*became_writable=*/true, priority_};
     }
-    return false;
+
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "UpdateWritableStateLocked for stream id: " << stream_id_
+        << " with priority: " << WritableStreams::GetPriorityString(priority_)
+        << " is_writable: " << is_writable_;
+    return EnqueueResult{/*became_writable=*/false, priority_};
   }
 
   const uint32_t stream_id_;
@@ -597,6 +619,8 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   Mutex mu_;
   bool is_writable_ ABSL_GUARDED_BY(mu_) = false;
   SimpleQueue<QueueEntry> queue_;
+  WritableStreams::StreamPriority priority_ ABSL_GUARDED_BY(mu_) =
+      WritableStreams::StreamPriority::kDefault;
 
   // Accessed only during dequeue.
   HeaderDisassembler initial_metadata_disassembler_;

test/core/transport/chttp2/stream_data_queue_fuzz_test.cc

@@ -48,7 +48,9 @@ class SimpleQueueFuzzTest : public YodelTest {
   auto EnqueueAndCheckSuccess(SimpleQueue<int>& queue, int data, int tokens) {
     return Map([&queue, data,
                 tokens]() mutable { return queue.Enqueue(data, tokens); },
-               [](auto result) { EXPECT_EQ(result.status, absl::OkStatus()); });
+               [](absl::StatusOr<bool> result) {
+                 EXPECT_EQ(result.status(), absl::OkStatus());
+               });
   }
 
   bool DequeueAndCheck(SimpleQueue<int>& queue, int data,