grpc 1.75.0 → 1.78.1

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

@@ -19,12 +19,35 @@
 #ifndef GRPC_SRC_CORE_EXT_TRANSPORT_CHTTP2_TRANSPORT_STREAM_DATA_QUEUE_H
 #define GRPC_SRC_CORE_EXT_TRANSPORT_CHTTP2_TRANSPORT_STREAM_DATA_QUEUE_H
 
+#include <grpc/support/port_platform.h>
+
+#include <algorithm>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <optional>
 #include <queue>
+#include <string>
+#include <utility>
+#include <variant>
+#include <vector>
 
-#include "absl/log/check.h"
-#include "absl/log/log.h"
+#include "src/core/call/message.h"
+#include "src/core/ext/transport/chttp2/transport/frame.h"
 #include "src/core/ext/transport/chttp2/transport/header_assembler.h"
+#include "src/core/ext/transport/chttp2/transport/hpack_encoder.h"
+#include "src/core/ext/transport/chttp2/transport/http2_status.h"
 #include "src/core/ext/transport/chttp2/transport/message_assembler.h"
+#include "src/core/ext/transport/chttp2/transport/transport_common.h"
+#include "src/core/lib/promise/activity.h"
+#include "src/core/lib/promise/context.h"
+#include "src/core/lib/promise/poll.h"
+#include "src/core/lib/slice/slice_buffer.h"
+#include "src/core/util/grpc_check.h"
+#include "src/core/util/ref_counted.h"
+#include "src/core/util/sync.h"
+#include "absl/log/log.h"
+#include "absl/status/status.h"
 
 namespace grpc_core {
 namespace http2 {
@@ -36,11 +59,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,10 +72,14 @@ 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<bool> Enqueue(T& data, const uint32_t tokens) {
     return PollEnqueue(data, tokens);
   }
 
+  absl::StatusOr<bool> ImmediateEnqueue(T data, const uint32_t tokens) {
+    return ImmediateEnqueueInternal(std::move(data), tokens);
+  }
+
   // Sync function to dequeue the next entry. Returns nullopt if the queue is
   // empty or if the front of the queue has more tokens than
   // allowed_dequeue_tokens. When allow_oversized_dequeue parameter is set to
@@ -77,10 +99,17 @@ class SimpleQueue {
   }
 
   // Returns true if the queue is empty. This function is NOT thread safe.
-  bool TestOnlyIsEmpty() const { return IsEmpty(); }
+  bool IsEmpty() const { return queue_.empty(); }
+  // Clears the queue. This function is NOT thread safe.
+  void Clear() { std::queue<Entry>().swap(queue_); }
+
+  inline std::optional<uint32_t> GetNextEntryTokens() const {
+    return queue_.empty() ? std::nullopt
+                          : std::make_optional(queue_.front().tokens);
+  }
 
  private:
-  Poll<EnqueueResult> PollEnqueue(T& data, const uint32_t tokens) {
+  Poll<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 +120,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
@@ -103,6 +131,21 @@ class SimpleQueue {
     return Pending{};
   }
 
+  inline absl::StatusOr<bool> ImmediateEnqueueInternal(T data,
+                                                       const uint32_t tokens) {
+    GRPC_DCHECK_LE(tokens_consumed_,
+                   std::numeric_limits<uint32_t>::max() - tokens);
+    if (tokens_consumed_ > std::numeric_limits<uint32_t>::max() - tokens) {
+      return absl::InternalError("Tokens consumed overflowed.");
+    }
+    tokens_consumed_ += tokens;
+    queue_.emplace(Entry{std::move(data), tokens});
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "Immediate enqueue successful. Data tokens: " << tokens
+        << " Current tokens consumed: " << tokens_consumed_;
+    return /*became_non_empty*/ (queue_.size() == 1);
+  }
+
   std::optional<T> DequeueInternal(const uint32_t allowed_dequeue_tokens,
                                    const bool allow_oversized_dequeue) {
     if (queue_.empty() || (queue_.front().tokens > allowed_dequeue_tokens &&
@@ -133,8 +176,6 @@ class SimpleQueue {
     return std::move(entry.data);
   }
 
-  bool IsEmpty() const { return queue_.empty(); }
-
   struct Entry {
     T data;
     uint32_t tokens;
@@ -148,10 +189,10 @@ class SimpleQueue {
   // exception to this rule: if the sender tries to enqueue an item when the
   // queue has 0 tokens, the enqueue will always go through regardless of the
   // number of tokens.
-  uint32_t max_tokens_;
+  const uint64_t max_tokens_;
   // The number of tokens that have been enqueued in the queue but not yet
   // dequeued.
-  uint32_t tokens_consumed_ = 0;
+  uint64_t tokens_consumed_ = 0;
   Waker waker_;
 };
 
@@ -161,15 +202,12 @@ class SimpleQueue {
 template <typename MetadataHandle>
 class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
  public:
-  explicit StreamDataQueue(const bool is_client, const uint32_t stream_id,
-                           const uint32_t queue_size)
-      : stream_id_(stream_id),
+  explicit StreamDataQueue(const bool is_client, const uint32_t queue_size)
+      : stream_id_(0),
         is_client_(is_client),
         queue_(queue_size),
-        initial_metadata_disassembler_(stream_id,
-                                       /*is_trailing_metadata=*/false),
-        trailing_metadata_disassembler_(stream_id,
-                                        /*is_trailing_metadata=*/true) {};
+        initial_metadata_disassembler_(/*is_trailing_metadata=*/false),
+        trailing_metadata_disassembler_(/*is_trailing_metadata=*/true) {};
   ~StreamDataQueue() = default;
 
   StreamDataQueue(StreamDataQueue&& rhs) = delete;
@@ -177,6 +215,17 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   StreamDataQueue(const StreamDataQueue&) = delete;
   StreamDataQueue& operator=(const StreamDataQueue&) = delete;
 
+  void SetStreamId(const uint32_t stream_id,
+                   const bool allow_true_binary_metadata_peer) {
+    GRPC_DCHECK_EQ(stream_id_, 0u);
+    GRPC_DCHECK_NE(stream_id, 0u);
+    stream_id_ = stream_id;
+    initial_metadata_disassembler_.Initialize(stream_id,
+                                              allow_true_binary_metadata_peer);
+    trailing_metadata_disassembler_.Initialize(stream_id,
+                                               allow_true_binary_metadata_peer);
+  }
+
   //////////////////////////////////////////////////////////////////////////////
   // Enqueue Helpers
   // These enqueue helpers are based on following assumptions:
@@ -194,66 +243,66 @@ 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 StreamWritabilityUpdate {
+    bool became_writable;
+    WritableStreamPriority 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) {
-    DCHECK(!is_initial_metadata_queued_);
-    DCHECK(!is_trailing_metadata_or_half_close_queued_);
-    DCHECK(metadata != nullptr);
-    DCHECK(!is_reset_stream_queued_);
+  absl::StatusOr<StreamWritabilityUpdate> EnqueueInitialMetadata(
+      MetadataHandle&& metadata) {
+    MutexLock lock(&mu_);
+    GRPC_DCHECK(!is_initial_metadata_queued_);
+    GRPC_DCHECK(!is_trailing_metadata_or_half_close_queued_);
+    GRPC_DCHECK(metadata != nullptr);
+    GRPC_DCHECK(reset_stream_state_ == RstStreamState::kNotQueued);
 
     is_initial_metadata_queued_ = true;
-    return [self = this->Ref(),
-            entry = QueueEntry{InitialMetadataType{
-                std::move(metadata)}}]() mutable -> Poll<absl::StatusOr<bool>> {
-      MutexLock lock(&self->mu_);
-      auto 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);
-          return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
-        }
-        return result.value().status;
-      }
-      return Pending{};
-    };
+    absl::StatusOr<bool> result = queue_.ImmediateEnqueue(
+        QueueEntry{InitialMetadataType{std::move(metadata)}}, /*tokens=*/0);
+    if (GPR_UNLIKELY(!result.ok())) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG
+          << "Immediate enqueueing initial metadata failed with status: "
+          << result.status();
+      return result.status();
+    }
+    return UpdateWritableStateAndPriorityEnqueueLocked(
+        /*became_non_empty*/ result.value(), WritableStreamPriority::kDefault);
   }
 
   // Enqueue Trailing Metadata.
   // 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) {
-    DCHECK(metadata != nullptr);
-    DCHECK(!is_reset_stream_queued_);
-    DCHECK(!is_client_);
-    DCHECK(!is_trailing_metadata_or_half_close_queued_);
+  absl::StatusOr<StreamWritabilityUpdate> EnqueueTrailingMetadata(
+      MetadataHandle&& metadata) {
+    MutexLock lock(&mu_);
+    GRPC_DCHECK(metadata != nullptr);
+    GRPC_DCHECK(!is_client_);
+    GRPC_DCHECK(!is_trailing_metadata_or_half_close_queued_);
+
+    if (GPR_UNLIKELY(IsEnqueueClosed())) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG << "Enqueue closed.";
+      return StreamWritabilityUpdate{/*became_writable=*/false,
+                                     WritableStreamPriority::kStreamClosed};
+    }
 
     is_trailing_metadata_or_half_close_queued_ = true;
-    return [self = this->Ref(),
-            entry = QueueEntry{TrailingMetadataType{
-                std::move(metadata)}}]() mutable -> Poll<absl::StatusOr<bool>> {
-      MutexLock lock(&self->mu_);
-      auto 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()) {
-          return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
-        }
-        return result.value().status;
-      }
-      return Pending{};
-    };
+    absl::StatusOr<bool> result = queue_.ImmediateEnqueue(
+        QueueEntry{TrailingMetadataType{std::move(metadata)}}, /*tokens=*/0);
+    if (GPR_UNLIKELY(!result.ok())) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG
+          << "Immediate enqueueing trailing metadata failed with status: "
+          << result.status();
+      return result.status();
+    }
+    return UpdateWritableStateAndPriorityEnqueueLocked(
+        /*became_non_empty*/ result.value(),
+        WritableStreamPriority::kStreamClosed);
   }
 
   // Returns a promise that resolves when the message is enqueued. There may be
@@ -261,30 +310,33 @@ 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) {
-    DCHECK(is_initial_metadata_queued_);
-    DCHECK(message != nullptr);
-    DCHECK(!is_reset_stream_queued_);
-    DCHECK_LE(message->payload()->Length(),
-              std::numeric_limits<uint32_t>::max() - kGrpcHeaderSizeInBytes);
-    DCHECK(!is_trailing_metadata_or_half_close_queued_);
+  auto EnqueueMessage(MessageHandle&& message) {
+    GRPC_DCHECK(is_initial_metadata_queued_);
+    GRPC_DCHECK(message != nullptr);
+    GRPC_DCHECK_LE(
+        message->payload()->Length(),
+        std::numeric_limits<uint32_t>::max() - kGrpcHeaderSizeInBytes);
+    GRPC_DCHECK(!is_trailing_metadata_or_half_close_queued_);
 
     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<StreamWritabilityUpdate>> {
       MutexLock lock(&self->mu_);
-      auto result = self->queue_.Enqueue(entry, tokens);
+      if (GPR_UNLIKELY(self->IsEnqueueClosed())) {
+        GRPC_STREAM_DATA_QUEUE_DEBUG << "Enqueue closed";
+        return StreamWritabilityUpdate{/*became_writable=*/false,
+                                       WritableStreamPriority::kStreamClosed};
+      }
+      Poll<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;
-        // TODO(akshitpatel) : [PH2][P2] : Add check for flow control tokens.
-        if (result.value().status.ok()) {
-          return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
-        }
-        return result.value().status;
+            << "Enqueued message with tokens: " << tokens
+            << "became_non_empty: " << result.value();
+
+        return self->UpdateWritableStateAndPriorityEnqueueLocked(
+            /*became_non_empty=*/result.value(),
+            WritableStreamPriority::kDefault);
       }
       return Pending{};
     };
@@ -294,66 +346,97 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   // 1. MUST be called at most once.
   // 2. MUST be called only for a client.
   // 3. This function is thread safe.
-  auto EnqueueHalfClosed() {
-    DCHECK(is_initial_metadata_queued_);
-    DCHECK(is_client_);
-    DCHECK(!is_reset_stream_queued_);
-    DCHECK(!is_trailing_metadata_or_half_close_queued_);
+  absl::StatusOr<StreamWritabilityUpdate> EnqueueHalfClosed() {
+    MutexLock lock(&mu_);
+    GRPC_DCHECK(is_initial_metadata_queued_);
+    GRPC_DCHECK(is_client_);
+
+    if (GPR_UNLIKELY(IsEnqueueClosed() ||
+                     is_trailing_metadata_or_half_close_queued_)) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG
+          << "Enqueue closed or trailing metadata/half close queued "
+          << " is_trailing_metadata_or_half_close_queued_ = "
+          << is_trailing_metadata_or_half_close_queued_;
+      return StreamWritabilityUpdate{/*became_writable=*/false,
+                                     WritableStreamPriority::kStreamClosed};
+    }
 
     is_trailing_metadata_or_half_close_queued_ = true;
-    return [self = this->Ref(), entry = QueueEntry{HalfClosed{}}]() mutable
-               -> Poll<absl::StatusOr<bool>> {
-      MutexLock lock(&self->mu_);
-      auto 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()) {
-          return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
-        }
-        return result.value().status;
-      }
-      return Pending{};
-    };
+    absl::StatusOr<bool> result =
+        queue_.ImmediateEnqueue(QueueEntry{HalfClosed{}}, /*tokens=*/0);
+    if (GPR_UNLIKELY(!result.ok())) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG
+          << "Immediate enqueueing half closed failed with status: "
+          << result.status();
+      return result.status();
+    }
+    return UpdateWritableStateAndPriorityEnqueueLocked(
+        /*became_non_empty*/ result.value(),
+        WritableStreamPriority::kStreamClosed);
   }
 
   // Enqueue Reset Stream.
   // 1. MUST be called at most once.
   // 3. This function is thread safe.
-  auto EnqueueResetStream(uint32_t error_code) {
-    DCHECK(is_initial_metadata_queued_);
-    DCHECK(!is_reset_stream_queued_);
-
-    is_reset_stream_queued_ = true;
-    return [self = this->Ref(),
-            entry = QueueEntry{ResetStream{
-                error_code}}]() mutable -> Poll<absl::StatusOr<bool>> {
-      MutexLock lock(&self->mu_);
-      auto 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()) {
-          return self->UpdateWritableStateLocked(
-              result.value().became_non_empty);
-        }
-        return result.value().status;
-      }
-      return Pending{};
-    };
+  absl::StatusOr<StreamWritabilityUpdate> EnqueueResetStream(
+      const uint32_t error_code) {
+    MutexLock lock(&mu_);
+    GRPC_DCHECK(is_initial_metadata_queued_);
+
+    // This can happen when the transport tries to close the stream and the
+    // stream is cancelled from the call stack.
+    if (GPR_UNLIKELY(IsEnqueueClosed())) {
+      GRPC_STREAM_DATA_QUEUE_DEBUG << "Enqueue closed";
+      return StreamWritabilityUpdate{/*became_writable=*/false,
+                                     WritableStreamPriority::kStreamClosed};
+    }
+
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "Immediate enqueueing reset stream with error code: " << error_code;
+    reset_stream_state_ = RstStreamState::kQueued;
+    reset_stream_error_code_ = error_code;
+
+    // became_non_empty is set to true if the queue is empty because we are not
+    // enqueueing reset stream to the queue. In this case, if the queue is
+    // empty, enqueuing reset stream to StreamDataQueue will make the stream
+    // writable.
+    return UpdateWritableStateAndPriorityEnqueueLocked(
+        /*became_non_empty*/ queue_.IsEmpty(),
+        WritableStreamPriority::kStreamClosed);
   }
 
   //////////////////////////////////////////////////////////////////////////////
   // Dequeue Helpers
 
-  // TODO(akshitpatel) : [PH2][P2] : Decide on whether it is needed to return
-  // the number of tokens consumed by one call of this function.
+  static constexpr uint8_t kResetStreamDequeued = 0x1;
+  static constexpr uint8_t kHalfCloseDequeued = 0x2;
+  static constexpr uint8_t kInitialMetadataDequeued = 0x4;
+
   struct DequeueResult {
     std::vector<Http2Frame> frames;
     bool is_writable;
+    WritableStreamPriority priority;
+    // Maybe not be extremely accurate but should be good enough for our
+    // purposes.
+    size_t total_bytes_consumed = 0u;
+    size_t flow_control_tokens_consumed = 0u;
+    // Bitmask of the dequeue flags.
+    uint8_t flags = 0u;
+
+    // Returns true if the reset stream was dequeued.
+    bool ResetStreamDequeued() const {
+      return (flags & kResetStreamDequeued) != 0u;
+    }
+
+    // Returns true if the half close was dequeued.
+    bool HalfCloseDequeued() const {
+      return (flags & kHalfCloseDequeued) != 0u;
+    }
+
+    // Returns true if the initial metadata was dequeued.
+    bool InitialMetadataDequeued() const {
+      return (flags & kInitialMetadataDequeued) != 0u;
+    }
   };
 
   // TODO(akshitpatel) : [PH2][P4] : Measure the performance of this function
@@ -375,16 +458,29 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
   //    the partial first message (sum of payload of all returned frames <=
   //    max_fc_tokens).
   // This function is thread safe.
-  absl::StatusOr<DequeueResult> DequeueFrames(const uint32_t max_fc_tokens,
-                                              const uint32_t max_frame_length,
-                                              HPackCompressor& encoder) {
+  DequeueResult DequeueFrames(const uint32_t max_fc_tokens,
+                              const uint32_t max_frame_length,
+                              const uint32_t stream_fc_tokens,
+                              HPackCompressor& encoder,
+                              const bool can_send_reset_stream) {
     MutexLock lock(&mu_);
-    GRPC_STREAM_DATA_QUEUE_DEBUG << "Dequeueing frames for stream "
-                                 << stream_id_
-                                 << " Max fc tokens: " << max_fc_tokens
-                                 << " Max frame length: " << max_frame_length
-                                 << " Message disassembler buffered length: "
-                                 << message_disassembler_.GetBufferedLength();
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "Dequeueing frames. Max fc tokens: " << max_fc_tokens
+        << " Max frame length: " << max_frame_length
+        << " Message disassembler buffered length: "
+        << message_disassembler_.GetBufferedLength()
+        << " Can send reset stream: " << can_send_reset_stream
+        << " Reset stream state: " << static_cast<uint8_t>(reset_stream_state_);
+    GRPC_DCHECK_GT(stream_id_, 0u)
+        << "Stream id must be set before dequeueing frames.";
+
+    // If a reset stream is queued, we do not want to send any more frames. Any
+    // metadata enqueued has not reached HPACK encoder, so it is safe to drop
+    // all frames.
+    if (std::optional<DequeueResult> result =
+            HandleResetStreamLocked(can_send_reset_stream)) {
+      return std::move(*result);
+    }
 
     HandleDequeue handle_dequeue(max_fc_tokens, max_frame_length, encoder,
                                  *this);
@@ -403,18 +499,41 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
       std::visit(handle_dequeue, std::move(*queue_entry));
     }
 
-    // TODO(akshitpatel) : [PH2][P2] : Add a check for flow control tokens.
-    is_writable_ = false;
-    GRPC_STREAM_DATA_QUEUE_DEBUG << "Stream id: " << stream_id_
-                                 << " writable state changed to "
-                                 << is_writable_;
-    return DequeueResult{handle_dequeue.GetFrames(), is_writable_};
+    GRPC_DCHECK_GE(stream_fc_tokens,
+                   handle_dequeue.GetFlowControlTokensConsumed());
+
+    return DequeueResult{
+        handle_dequeue.GetFrames(),
+        UpdateWritableStateDequeueLocked(
+            stream_fc_tokens - handle_dequeue.GetFlowControlTokensConsumed()),
+        priority_,
+        handle_dequeue.GetTotalBytesConsumed(),
+        handle_dequeue.GetFlowControlTokensConsumed(),
+        handle_dequeue.GetDequeueFlags()};
+  }
+
+  // TODO(tjagtap) : [PH2][P1][FlowControl] : Call this while processing
+  // window update frame.
+  // Needs to be invoked when the peer sends stream flow control window update.
+  // stream_fc_tokens represents the stream flow control (delta) window +
+  // intial_window_size.
+  StreamWritabilityUpdate ReceivedFlowControlWindowUpdate(
+      const uint32_t stream_fc_tokens) {
+    MutexLock lock(&mu_);
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "Received flow control window update. stream_fc_tokens: "
+        << stream_fc_tokens;
+    const bool old_writable_state = is_writable_;
+    const bool new_writable_state =
+        UpdateWritableStateDequeueLocked(stream_fc_tokens);
+    return {/*became_writable=*/(!old_writable_state && new_writable_state),
+            priority_};
   }
 
   // Returns true if the queue is empty. This function is thread safe.
   bool TestOnlyIsEmpty() {
     MutexLock lock(&mu_);
-    return queue_.TestOnlyIsEmpty();
+    return queue_.IsEmpty();
   }
 
  private:
@@ -425,26 +544,24 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
     MetadataHandle metadata;
   };
   struct HalfClosed {};
-  struct ResetStream {
-    uint32_t error_code;
-  };
   using QueueEntry = std::variant<InitialMetadataType, TrailingMetadataType,
-                                  MessageHandle, HalfClosed, ResetStream>;
+                                  MessageHandle, HalfClosed>;
 
   class HandleDequeue {
    public:
-    HandleDequeue(uint32_t max_tokens, uint32_t max_frame_length,
+    HandleDequeue(const uint32_t max_tokens, const uint32_t max_frame_length,
                   HPackCompressor& encoder, StreamDataQueue& queue)
         : queue_(queue),
           max_frame_length_(max_frame_length),
-          max_fc_tokens_(max_tokens),
-          fc_tokens_available_(max_tokens),
+          max_tokens_available_(max_tokens),
+          flow_control_tokens_consumed_(0),
           encoder_(encoder) {}
 
     void operator()(InitialMetadataType initial_metadata) {
       GRPC_STREAM_DATA_QUEUE_DEBUG << "Preparing initial metadata for sending";
       queue_.initial_metadata_disassembler_.PrepareForSending(
           std::move(initial_metadata.metadata), encoder_);
+      dequeue_flags_ |= kInitialMetadataDequeued;
       MaybeAppendInitialMetadataFrames();
     }
 
@@ -462,13 +579,7 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
 
     void operator()(GRPC_UNUSED HalfClosed half_closed) {
       GRPC_STREAM_DATA_QUEUE_DEBUG << "Preparing end of stream for sending";
-      is_half_closed_ = true;
-    }
-
-    void operator()(ResetStream reset_stream) {
-      GRPC_STREAM_DATA_QUEUE_DEBUG << "Preparing reset stream for sending";
-      is_reset_stream_ = true;
-      error_code_ = reset_stream.error_code;
+      dequeue_flags_ |= kHalfCloseDequeued;
     }
 
     std::vector<Http2Frame> GetFrames() {
@@ -483,117 +594,274 @@ class StreamDataQueue : public RefCounted<StreamDataQueue<MetadataHandle>> {
       MaybeAppendMessageFrames();
       MaybeAppendEndOfStreamFrame();
       MaybeAppendTrailingMetadataFrames();
-      MaybeAppendResetStreamFrame();
       return std::move(frames_);
     }
 
+    size_t GetTotalBytesConsumed() const { return total_bytes_consumed_; }
+    size_t GetFlowControlTokensConsumed() const {
+      return flow_control_tokens_consumed_;
+    }
+    uint8_t GetDequeueFlags() const { return dequeue_flags_; }
+
    private:
     inline void MaybeAppendInitialMetadataFrames() {
       while (queue_.initial_metadata_disassembler_.HasMoreData()) {
-        DCHECK(!is_half_closed_);
-        DCHECK(!is_reset_stream_);
+        GRPC_DCHECK(!(dequeue_flags_ & kHalfCloseDequeued));
+        GRPC_DCHECK(!(dequeue_flags_ & kResetStreamDequeued));
         // TODO(akshitpatel) : [PH2][P2] : I do not think we need this.
         // HasMoreData() should be enough.
         bool is_end_headers = false;
-        frames_.emplace_back(queue_.initial_metadata_disassembler_.GetNextFrame(
+        AppendFrame(queue_.initial_metadata_disassembler_.GetNextFrame(
             max_frame_length_, is_end_headers));
       }
     }
 
     inline void MaybeAppendTrailingMetadataFrames() {
       while (queue_.trailing_metadata_disassembler_.HasMoreData()) {
-        DCHECK(!is_half_closed_);
-        DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
-        DCHECK_EQ(queue_.initial_metadata_disassembler_.GetBufferedLength(),
-                  0u);
+        GRPC_DCHECK(!(dequeue_flags_ & kHalfCloseDequeued));
+        GRPC_DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
+        GRPC_DCHECK_EQ(
+            queue_.initial_metadata_disassembler_.GetBufferedLength(), 0u);
         // TODO(akshitpatel) : [PH2][P2] : I do not think we need this.
         // HasMoreData() should be enough.
         bool is_end_headers = false;
-        frames_.emplace_back(
-            queue_.trailing_metadata_disassembler_.GetNextFrame(
-                max_frame_length_, is_end_headers));
+        AppendFrame(queue_.trailing_metadata_disassembler_.GetNextFrame(
+            max_frame_length_, is_end_headers));
       }
     }
 
     inline void MaybeAppendEndOfStreamFrame() {
-      if (is_half_closed_) {
-        DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
-        DCHECK_EQ(queue_.initial_metadata_disassembler_.GetBufferedLength(),
-                  0u);
-        DCHECK_EQ(queue_.trailing_metadata_disassembler_.GetBufferedLength(),
-                  0u);
-        frames_.emplace_back(Http2DataFrame{/*stream_id=*/queue_.stream_id_,
-                                            /*end_stream=*/true,
-                                            /*payload=*/SliceBuffer()});
+      if (dequeue_flags_ & kHalfCloseDequeued) {
+        GRPC_DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
+        GRPC_DCHECK_EQ(
+            queue_.initial_metadata_disassembler_.GetBufferedLength(), 0u);
+        GRPC_DCHECK_EQ(
+            queue_.trailing_metadata_disassembler_.GetBufferedLength(), 0u);
+        AppendFrame(Http2DataFrame{/*stream_id=*/queue_.stream_id_,
+                                   /*end_stream=*/true,
+                                   /*payload=*/SliceBuffer()});
       }
     }
 
     inline void MaybeAppendMessageFrames() {
       while (queue_.message_disassembler_.GetBufferedLength() > 0 &&
-             fc_tokens_available_ > 0) {
-        DCHECK_EQ(queue_.initial_metadata_disassembler_.GetBufferedLength(),
-                  0u);
+             (max_tokens_available_ - flow_control_tokens_consumed_) > 0) {
+        GRPC_DCHECK_EQ(
+            queue_.initial_metadata_disassembler_.GetBufferedLength(), 0u);
         Http2DataFrame frame = queue_.message_disassembler_.GenerateNextFrame(
             queue_.stream_id_,
-            std::min(fc_tokens_available_, max_frame_length_));
-        fc_tokens_available_ -= frame.payload.Length();
+            std::min(max_tokens_available_ - flow_control_tokens_consumed_,
+                     max_frame_length_));
+        flow_control_tokens_consumed_ += frame.payload.Length();
         GRPC_STREAM_DATA_QUEUE_DEBUG
             << "Appending message frame with length " << frame.payload.Length()
-            << "Available tokens: " << fc_tokens_available_;
-        frames_.emplace_back(std::move(frame));
+            << " Consumed tokens: " << flow_control_tokens_consumed_
+            << " Max tokens: " << max_tokens_available_;
+        AppendFrame(std::move(frame));
       }
     }
 
     inline void MaybeAppendResetStreamFrame() {
-      if (is_reset_stream_) {
+      if (dequeue_flags_ & kResetStreamDequeued) {
         // TODO(akshitpatel) : [PH2][P2] : Consider if we can send reset stream
         // frame without flushing all the messages enqueued until now.
-        DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
-        DCHECK_EQ(queue_.initial_metadata_disassembler_.GetBufferedLength(),
-                  0u);
-        DCHECK_EQ(queue_.trailing_metadata_disassembler_.GetBufferedLength(),
-                  0u);
-        frames_.emplace_back(
-            Http2RstStreamFrame{queue_.stream_id_, error_code_});
+        GRPC_DCHECK_EQ(queue_.message_disassembler_.GetBufferedLength(), 0u);
+        GRPC_DCHECK_EQ(
+            queue_.initial_metadata_disassembler_.GetBufferedLength(), 0u);
+        GRPC_DCHECK_EQ(
+            queue_.trailing_metadata_disassembler_.GetBufferedLength(), 0u);
+        AppendFrame(Http2RstStreamFrame{queue_.stream_id_, error_code_});
       }
     }
 
+    inline void AppendFrame(Http2Frame&& frame) {
+      total_bytes_consumed_ += GetFrameMemoryUsage(frame);
+      frames_.emplace_back(std::move(frame));
+    }
+
     StreamDataQueue& queue_;
     const uint32_t max_frame_length_;
-    const uint32_t max_fc_tokens_;
-    uint32_t fc_tokens_available_;
-    bool is_half_closed_ = false;
-    bool is_reset_stream_ = false;
+    const uint32_t max_tokens_available_;
+    uint32_t flow_control_tokens_consumed_;
     uint32_t error_code_ = static_cast<uint32_t>(Http2ErrorCode::kNoError);
     std::vector<Http2Frame> frames_;
     HPackCompressor& encoder_;
+    size_t total_bytes_consumed_ = 0u;
+    uint8_t dequeue_flags_ = 0u;
   };
 
-  // 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 writable state and priority of the stream. MUST only be called
+  // from the enqueue functions.
+  // became_non_empty: True if the queue was empty and became non-empty as a
+  //                   result of this enqueue operation.
+  // priority: The new priority of the stream after this enqueue operation.
+  // Returns the result of the enqueue operation
+  //
+  // High level flow:
+  // Priority is simply updated to the new priority.
+  // Writable state is updated as follows:
+  // 1. If the stream was not writable before and became non-empty as a result
+  //    of this enqueue operation, then the stream is marked as writable.
+  // 2. If the stream was already writable before, it remains writable.
+  // 3. The case where the stream was not writable before and the queue already
+  //    contained data implies that the stream is blocked on stream flow control
+  //    tokens. When the transport receives stream flow control window update,
+  //    the stream is marked as writable.
+  // For enqueue operations there is no easy way to query
+  // stream_flow_control_tokens. So it is assumed that flow control tokens are
+  // always available for an enqueue operation. This can cause a stream to be
+  // marked as writable when it is not but this will correct itself in the next
+  // dequeue operation (which returns an accurate is_writable).
+  StreamWritabilityUpdate UpdateWritableStateAndPriorityEnqueueLocked(
+      const bool became_non_empty, const WritableStreamPriority priority)
       ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    // Update priority.
+    priority_ = priority;
+
+    // Update writable state.
     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 became writable with priority: "
+          << GetWritableStreamPriorityString(priority_);
+      return StreamWritabilityUpdate{/*became_writable=*/true, priority_};
+    }
+
+    GRPC_STREAM_DATA_QUEUE_DEBUG
+        << "UpdateWritableStateAndPriorityEnqueueLocked with priority: "
+        << GetWritableStreamPriorityString(priority_)
+        << " is_writable: " << is_writable_;
+    return StreamWritabilityUpdate{/*became_writable=*/false, priority_};
+  }
+
+  // Updates the writable state of the stream. Returns true if the
+  // stream became writable.
+  // Writable state is updated as follows:
+  // 1. If the next message to dequeue is a grpc message, then the stream is
+  //    writable if and only if we have available stream flow control tokens.
+  // 2. If the next message to dequeue is not a grpc message, then the stream is
+  //    writable if and only if the queue is not empty.
+  // Unlike UpdateWritableStateAndPriorityEnqueueLocked, this function the
+  // become_writable returned by this function is `accurate` as it considers
+  // the whether the stream has bytes to write and the flow control tokens
+  // available.
+  bool UpdateWritableStateDequeueLocked(
+      const uint32_t available_stream_fc_tokens)
+      ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    // The stream is writable if the queue is not empty. If the next bytes to
+    // write are part of a gRPC message, then check if we have stream flow
+    // control tokens.
+    is_writable_ = (!queue_.IsEmpty());
+
+    // Next bytes to write are part of a gRPC message.
+    if (message_disassembler_.GetBufferedLength() > 0 ||
+        IsNextQueueEntryMessage()) {
+      is_writable_ = (available_stream_fc_tokens > 0);
+    }
+
+    GRPC_STREAM_DATA_QUEUE_DEBUG << "UpdateWritableStateLocked with priority: "
+                                 << GetWritableStreamPriorityString(priority_)
+                                 << " is_writable: " << is_writable_;
+    return is_writable_;
+  }
+
+  inline bool IsNextQueueEntryMessage() const {
+    return (!queue_.IsEmpty() && queue_.GetNextEntryTokens().value() > 0);
+  }
+
+  // Handles the case where a reset stream is queued.
+  // If a reset stream is queued or has been dequeued, this function returns a
+  // DequeueResult. Otherwise, it returns std::nullopt.
+  // This function must be called with mu_ held.
+  std::optional<DequeueResult> HandleResetStreamLocked(
+      const bool can_send_reset_stream) ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    switch (reset_stream_state_) {
+      case RstStreamState::kDequeued:
+        GRPC_STREAM_DATA_QUEUE_DEBUG
+            << "Reset stream is already dequeued. Returning empty frames.";
+        GRPC_DCHECK(queue_.IsEmpty());
+        is_writable_ = false;
+        return DequeueResult{
+            std::vector<Http2Frame>(),           is_writable_, priority_,
+            /*total_bytes_consumed=*/0u,
+            /*flow_control_tokens_consumed=*/0u, /*flags=*/0u};
+      case RstStreamState::kQueued: {
+        GRPC_STREAM_DATA_QUEUE_DEBUG
+            << "Reset stream is queued. Skipping all frames (if any) for "
+               "dequeuing.";
+        is_writable_ = false;
+        std::vector<Http2Frame> frames;
+        uint8_t flags = 0u;
+        if (can_send_reset_stream) {
+          frames.emplace_back(
+              Http2RstStreamFrame{stream_id_, reset_stream_error_code_});
+          flags = kResetStreamDequeued;
+        }
+        queue_.Clear();
+        reset_stream_state_ = RstStreamState::kDequeued;
+        return DequeueResult{std::move(frames),
+                             is_writable_,
+                             priority_,
+                             /*total_bytes_consumed=*/0u,
+                             /*flow_control_tokens_consumed=*/0u,
+                             flags};
+      }
+      case RstStreamState::kNotQueued:
+        return std::nullopt;
+      default:
+        GRPC_CHECK(false) << "Invalid reset stream state: "
+                          << static_cast<uint8_t>(reset_stream_state_);
+        return std::nullopt;
     }
-    return false;
   }
 
-  const uint32_t stream_id_;
+  bool IsEnqueueClosed() const ABSL_EXCLUSIVE_LOCKS_REQUIRED(mu_) {
+    switch (reset_stream_state_) {
+      case RstStreamState::kNotQueued:
+        return false;
+      case RstStreamState::kQueued:
+      case RstStreamState::kDequeued:
+        // This can happen when the transport tries to close the stream and the
+        // stream is cancelled from the call stack.
+        GRPC_STREAM_DATA_QUEUE_DEBUG << "Reset stream already queued.";
+        return true;
+      default:
+        GRPC_CHECK(false) << "Invalid reset stream state: "
+                          << static_cast<uint8_t>(reset_stream_state_);
+    }
+
+    GPR_UNREACHABLE_CODE(return false);
+  }
+
+  uint32_t stream_id_;
+
+  // This is only used for DCHECKs. Not actually used for any business logic.
   const bool is_client_;
 
+  enum class RstStreamState : uint8_t {
+    kNotQueued = 0,
+    kQueued,
+    kDequeued,
+  };
+
   // Accessed only during enqueue.
   bool is_initial_metadata_queued_ = false;
   bool is_trailing_metadata_or_half_close_queued_ = false;
-  bool is_reset_stream_queued_ = false;
 
   // Access both during enqueue and dequeue.
   Mutex mu_;
+  // This variable tracks whether the stream is writable. 'Writable' represents
+  // that the stream has bytes to send and the stream has flow control tokens
+  // (if needed) to send them. This variable also has 1-1 correspondence with
+  // whether the stream is in the list of writable streams in the transport.
   bool is_writable_ ABSL_GUARDED_BY(mu_) = false;
+  RstStreamState reset_stream_state_ ABSL_GUARDED_BY(mu_) =
+      RstStreamState::kNotQueued;
   SimpleQueue<QueueEntry> queue_;
+  WritableStreamPriority priority_ ABSL_GUARDED_BY(mu_) =
+      WritableStreamPriority::kDefault;
+  uint32_t reset_stream_error_code_ ABSL_GUARDED_BY(mu_) =
+      static_cast<uint32_t>(Http2ErrorCode::kNoError);
 
   // Accessed only during dequeue.
   HeaderDisassembler initial_metadata_disassembler_;