[PH2] Defer assigning ID to streams till the first write.

PiperOrigin-RevId: 815642269

Author: ac-patel

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

@@ -146,6 +146,8 @@ class HeaderAssembler {
 
     // Validate
     GRPC_DCHECK_EQ(is_ready_, true);
+    GRPC_DCHECK_GT(stream_id_, 0u)
+        << "Stream id must be set before reading metadata.";
 
     // Generate the gRPC Metadata from buffer_
     // RFC9113 :  A receiver MUST terminate the connection with a connection
@@ -198,12 +200,11 @@ class HeaderAssembler {
   // This value MUST be checked before calling ReadMetadata()
   bool IsReady() const { return is_ready_; }
 
-  explicit HeaderAssembler(const uint32_t stream_id,
-                           const bool allow_true_binary_metadata_acked)
+  explicit HeaderAssembler(const bool allow_true_binary_metadata_acked)
       : header_in_progress_(false),
         is_ready_(false),
         allow_true_binary_metadata_acked_(allow_true_binary_metadata_acked),
-        stream_id_(stream_id) {}
+        stream_id_(0) {}
 
   ~HeaderAssembler() = default;
 
@@ -212,6 +213,12 @@ class HeaderAssembler {
   HeaderAssembler(const HeaderAssembler&) = delete;
   HeaderAssembler& operator=(const HeaderAssembler&) = delete;
 
+  void SetStreamId(const uint32_t stream_id) {
+    GRPC_DCHECK_EQ(stream_id_, 0u);
+    GRPC_DCHECK_NE(stream_id, 0u);
+    stream_id_ = stream_id;
+  }
+
  private:
   void Cleanup() {
     buffer_.Clear();
@@ -222,7 +229,7 @@ class HeaderAssembler {
   bool header_in_progress_;
   bool is_ready_;
   GRPC_UNUSED const bool allow_true_binary_metadata_acked_;
-  const uint32_t stream_id_;
+  uint32_t stream_id_;
   SliceBuffer buffer_;
 };
 
@@ -244,6 +251,9 @@ class HeaderDisassembler {
 
   Http2Frame GetNextFrame(const uint32_t max_frame_length,
                           bool& out_end_headers) {
+    GRPC_DCHECK_GT(stream_id_, 0u) << "Stream id must be set before getting "
+                                      "next Header/Continuation frame.";
+
     if (buffer_.Length() == 0 || is_done_) {
       GRPC_DCHECK(false) << "Calling code must check size using HasMoreData() "
                             "before GetNextFrame()";
@@ -273,10 +283,9 @@ class HeaderDisassembler {
 
   // A separate HeaderDisassembler object MUST be made for Initial Metadata and
   // Trailing Metadata
-  explicit HeaderDisassembler(const uint32_t stream_id,
-                              const bool is_trailing_metadata,
+  explicit HeaderDisassembler(const bool is_trailing_metadata,
                               const bool allow_true_binary_metadata_peer)
-      : stream_id_(stream_id),
+      : stream_id_(0),
         end_stream_(is_trailing_metadata),
         did_send_header_frame_(false),
         is_done_(false),
@@ -290,9 +299,14 @@ class HeaderDisassembler {
   HeaderDisassembler& operator=(const HeaderDisassembler&) = delete;
 
   size_t TestOnlyGetMainBufferLength() const { return buffer_.Length(); }
+  void SetStreamId(const uint32_t stream_id) {
+    GRPC_DCHECK_EQ(stream_id_, 0u);
+    GRPC_DCHECK_NE(stream_id, 0u);
+    stream_id_ = stream_id;
+  }
 
  private:
-  const uint32_t stream_id_;
+  uint32_t stream_id_;
   const bool end_stream_;
   bool did_send_header_frame_;
   bool is_done_;  // Protect against the same disassembler from being used twice

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

@@ -31,6 +31,7 @@
 #include "absl/strings/string_view.h"
 #include "src/core/call/call_spine.h"
 #include "src/core/call/message.h"
+#include "src/core/call/metadata.h"
 #include "src/core/call/metadata_batch.h"
 #include "src/core/ext/transport/chttp2/transport/flow_control.h"
 #include "src/core/ext/transport/chttp2/transport/flow_control_manager.h"
@@ -855,6 +856,25 @@ auto Http2ClientTransport::MultiplexerLoop() {
                 << stream->GetStreamId()
                 << " is_closed_for_writes = " << stream->IsClosedForWrites();
 
+            if (stream->GetStreamId() == kInvalidStreamId) {
+              // TODO(akshitpatel) : [PH2][P4] : We will waste a stream id in
+              // the rare scenario where the stream is aborted before it can be
+              // written to. This is a possible area to optimize in future.
+              absl::Status status =
+                  self->AssignStreamIdAndAddToStreamList(stream);
+              if (!status.ok()) {
+                GRPC_HTTP2_CLIENT_DLOG
+                    << "Http2ClientTransport MultiplexerLoop "
+                       "Failed to assign stream id and add to stream list for "
+                       "stream: "
+                    << stream.get() << " closing this stream.";
+                self->BeginCloseStream(
+                    stream, /*reset_stream_error_code=*/std::nullopt,
+                    CancelledServerMetadataFromStatus(status));
+                continue;
+              }
+            }
+
             if (GPR_LIKELY(!stream->IsClosedForWrites())) {
               auto stream_frames = self->DequeueStreamFrames(stream);
               if (GPR_UNLIKELY(!stream_frames.ok())) {
@@ -909,6 +929,28 @@ auto Http2ClientTransport::OnMultiplexerLoopEnded() {
       };
 }
 
+absl::Status Http2ClientTransport::AssignStreamIdAndAddToStreamList(
+    RefCountedPtr<Stream> stream) {
+  absl::StatusOr<uint32_t> next_stream_id = NextStreamId();
+  if (!next_stream_id.ok()) {
+    GRPC_HTTP2_CLIENT_DLOG
+        << "Http2ClientTransport AssignStreamIdAndAddToStreamList "
+           "Failed to get next stream id for stream: "
+        << stream.get();
+    return std::move(next_stream_id).status();
+  }
+  GRPC_HTTP2_CLIENT_DLOG
+      << "Http2ClientTransport AssignStreamIdAndAddToStreamList "
+         "Assigned stream id: "
+      << next_stream_id.value() << " to stream: " << stream.get();
+  stream->SetStreamId(next_stream_id.value());
+  {
+    MutexLock lock(&transport_mutex_);
+    stream_list_.emplace(next_stream_id.value(), stream);
+  }
+  return absl::OkStatus();
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // Constructor Destructor
 
@@ -919,7 +961,7 @@ Http2ClientTransport::Http2ClientTransport(
     : channelz::DataSource(http2::CreateChannelzSocketNode(
           endpoint.GetEventEngineEndpoint(), channel_args)),
       endpoint_(std::move(endpoint)),
-      stream_id_mutex_(/*Initial Stream Id*/ 1),
+      next_stream_id_(/*Initial Stream ID*/ 1),
       should_reset_ping_clock_(false),
       incoming_header_in_progress_(false),
       incoming_header_end_stream_(false),
@@ -1308,37 +1350,26 @@ bool Http2ClientTransport::SetOnDone(CallHandler call_handler,
 }
 
 std::optional<RefCountedPtr<Stream>> Http2ClientTransport::MakeStream(
-    CallHandler call_handler,
-    InterActivityMutex<uint32_t>::Lock& next_stream_id_lock) {
+    CallHandler call_handler) {
   // https://datatracker.ietf.org/doc/html/rfc9113#name-stream-identifiers
-  // TODO(akshitpatel) : [PH2][P1] : Probably do not need this lock. This
-  // function is always called under the stream_id_mutex_. The issue is the
-  // OnDone needs to be synchronous and hence InterActivityMutex might not be
-  // an option to protect the stream_list_.
-  MutexLock lock(&transport_mutex_);
-
-  std::optional<uint32_t> stream_id = NextStreamId(next_stream_id_lock);
-  if (!stream_id.has_value()) {
-    return std::nullopt;
+  RefCountedPtr<Stream> stream;
+  {
+    MutexLock lock(&transport_mutex_);
+    stream = MakeRefCounted<Stream>(
+        call_handler, settings_.peer().allow_true_binary_metadata(),
+        settings_.acked().allow_true_binary_metadata(), flow_control_);
   }
-
-  RefCountedPtr<Stream> stream = MakeRefCounted<Stream>(
-      call_handler, stream_id.value(),
-      settings_.peer().allow_true_binary_metadata(),
-      settings_.acked().allow_true_binary_metadata(), flow_control_);
   const bool on_done_added = SetOnDone(call_handler, stream);
   if (!on_done_added) return std::nullopt;
-  stream_list_.emplace(stream_id.value(), stream);
   return stream;
 }
 
 ///////////////////////////////////////////////////////////////////////////////
 // Call Spine related operations
 
-auto Http2ClientTransport::CallOutboundLoop(
-    CallHandler call_handler, RefCountedPtr<Stream> stream,
-    InterActivityMutex<uint32_t>::Lock lock /* Locked stream_id_mutex */,
-    ClientMetadataHandle metadata) {
+auto Http2ClientTransport::CallOutboundLoop(CallHandler call_handler,
+                                            RefCountedPtr<Stream> stream,
+                                            ClientMetadataHandle metadata) {
   GRPC_HTTP2_CLIENT_DLOG << "Http2ClientTransport CallOutboundLoop";
   GRPC_DCHECK(stream != nullptr);
 
@@ -1384,7 +1415,7 @@ auto Http2ClientTransport::CallOutboundLoop(
       "Ph2CallOutboundLoop",
       TrySeq(
           send_initial_metadata(),
-          [call_handler, send_message, lock = std::move(lock)]() {
+          [call_handler, send_message]() {
             // The lock will be released once the promise is constructed from
             // this factory. ForEach will be polled after the lock is
             // released.
@@ -1409,43 +1440,38 @@ void Http2ClientTransport::StartCall(CallHandler call_handler) {
   GRPC_HTTP2_CLIENT_DLOG << "Http2ClientTransport StartCall Begin";
   call_handler.SpawnGuarded(
       "OutboundLoop",
-      TrySeq(
-          call_handler.PullClientInitialMetadata(),
-          [self = RefAsSubclass<Http2ClientTransport>()](
-              ClientMetadataHandle metadata) {
-            // Lock the stream_id_mutex_
-            return Staple(self->stream_id_mutex_.Acquire(),
-                          std::move(metadata));
-          },
-          [self = RefAsSubclass<Http2ClientTransport>(),
-           call_handler](auto args /* Locked stream_id_mutex */) mutable {
-            // For a gRPC Client, we only need to check the
-            // MAX_CONCURRENT_STREAMS setting compliance at the time of
-            // sending (that is write path). A gRPC Client will never
-            // receive a stream initiated by a server, so we dont have to
-            // check MAX_CONCURRENT_STREAMS compliance on the Read-Path.
-            //
-            // TODO(tjagtap) : [PH2][P1] Check for MAX_CONCURRENT_STREAMS
-            // sent by peer before making a stream. Decide behaviour if we are
-            // crossing this threshold.
-            //
-            // TODO(tjagtap) : [PH2][P1] : For a server we will have to do
-            // this for incoming streams only. If a server receives more streams
-            // from a client than is allowed by the clients settings, whether or
-            // not we should fail is debatable.
-            std::optional<RefCountedPtr<Stream>> stream =
-                self->MakeStream(call_handler, std::get<0>(args));
-            return If(
-                stream.has_value(),
-                [self, call_handler, stream, args = std::move(args)]() mutable {
-                  return Map(
-                      self->CallOutboundLoop(call_handler, stream.value(),
-                                             std::move(std::get<0>(args)),
-                                             std::move(std::get<1>(args))),
-                      [](absl::Status status) { return status; });
-                },
-                []() { return absl::InternalError("Failed to make stream"); });
-          }));
+      TrySeq(call_handler.PullClientInitialMetadata(),
+             [self = RefAsSubclass<Http2ClientTransport>(),
+              call_handler](ClientMetadataHandle metadata) mutable {
+               // For a gRPC Client, we only need to check the
+               // MAX_CONCURRENT_STREAMS setting compliance at the time of
+               // sending (that is write path). A gRPC Client will never
+               // receive a stream initiated by a server, so we dont have to
+               // check MAX_CONCURRENT_STREAMS compliance on the Read-Path.
+               //
+               // TODO(tjagtap) : [PH2][P1] Check for MAX_CONCURRENT_STREAMS
+               // sent by peer before making a stream. Decide behaviour if we
+               // are crossing this threshold.
+               //
+               // TODO(tjagtap) : [PH2][P1] : For a server we will have to do
+               // this for incoming streams only. If a server receives more
+               // streams from a client than is allowed by the clients settings,
+               // whether or not we should fail is debatable.
+               std::optional<RefCountedPtr<Stream>> stream =
+                   self->MakeStream(call_handler);
+               return If(
+                   stream.has_value(),
+                   [self, call_handler, stream,
+                    initial_metadata = std::move(metadata)]() mutable {
+                     return Map(
+                         self->CallOutboundLoop(call_handler, stream.value(),
+                                                std::move(initial_metadata)),
+                         [](absl::Status status) { return status; });
+                   },
+                   []() {
+                     return absl::InternalError("Failed to make stream");
+                   });
+             }));
   GRPC_HTTP2_CLIENT_DLOG << "Http2ClientTransport StartCall End";
 }
 

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

@@ -234,7 +234,6 @@ class Http2ClientTransport final : public ClientTransport,
   // Returns a promise to fetch data from the callhandler and pass it further
   // down towards the endpoint.
   auto CallOutboundLoop(CallHandler call_handler, RefCountedPtr<Stream> stream,
-                        InterActivityMutex<uint32_t>::Lock lock,
                         ClientMetadataHandle metadata);
 
   // Force triggers a transport write cycle
@@ -280,11 +279,10 @@ class Http2ClientTransport final : public ClientTransport,
     return stream_list_.size();
   }
 
-  std::optional<uint32_t> NextStreamId(
-      InterActivityMutex<uint32_t>::Lock& next_stream_id_lock)
-      ABSL_EXCLUSIVE_LOCKS_REQUIRED(transport_mutex_) {
-    const uint32_t stream_id = *next_stream_id_lock;
-    if (stream_id > RFC9113::kMaxStreamId31Bit) {
+  // Returns the next stream id. If the next stream id is not available, it
+  // returns std::nullopt. MUST be called from the transport party.
+  absl::StatusOr<uint32_t> NextStreamId() {
+    if (next_stream_id_ > RFC9113::kMaxStreamId31Bit) {
       // TODO(tjagtap) : [PH2][P3] : Handle case if transport runs out of stream
       // ids
       // RFC9113 : Stream identifiers cannot be reused. Long-lived connections
@@ -294,29 +292,36 @@ class Http2ClientTransport final : public ClientTransport,
       // that is unable to establish a new stream identifier can send a GOAWAY
       // frame so that the client is forced to open a new connection for new
       // streams.
-      return std::nullopt;
+      return absl::ResourceExhaustedError("No more stream ids available");
     }
     // TODO(akshitpatel) : [PH2][P3] : There is a channel arg to delay
     // starting new streams instead of failing them. This needs to be
     // implemented.
-    if (GetActiveStreamCount() >= settings_.peer().max_concurrent_streams()) {
-      return std::nullopt;
+    {
+      MutexLock lock(&transport_mutex_);
+      if (GetActiveStreamCount() >= settings_.peer().max_concurrent_streams()) {
+        return absl::ResourceExhaustedError("Reached max concurrent streams");
+      }
     }
+
     // RFC9113 : Streams initiated by a client MUST use odd-numbered stream
     // identifiers.
-    (*next_stream_id_lock) += 2;
-    return stream_id;
+    return std::exchange(next_stream_id_, next_stream_id_ + 2);
   }
 
+  // Returns the next stream id without incrementing it. MUST be called from the
+  // transport party.
+  uint32_t PeekNextStreamId() const { return next_stream_id_; }
+
+  absl::Status AssignStreamIdAndAddToStreamList(RefCountedPtr<Stream> stream);
+
   Mutex transport_mutex_;
   // TODO(tjagtap) : [PH2][P2] : Add to map in StartCall and clean this
   // mapping up in the on_done of the CallInitiator or CallHandler
   absl::flat_hash_map<uint32_t, RefCountedPtr<Stream>> stream_list_
       ABSL_GUARDED_BY(transport_mutex_);
 
-  // Mutex to preserve the order of headers being sent out for new streams.
-  // This also tracks the stream_id for creating new streams.
-  InterActivityMutex<uint32_t> stream_id_mutex_;
+  uint32_t next_stream_id_;
   HPackCompressor encoder_;
   HPackParser parser_;
   bool is_transport_closed_ ABSL_GUARDED_BY(transport_mutex_) = false;
@@ -335,9 +340,8 @@ class Http2ClientTransport final : public ClientTransport,
   ConnectivityStateTracker state_tracker_ ABSL_GUARDED_BY(transport_mutex_){
       "http2_client", GRPC_CHANNEL_READY};
 
-  std::optional<RefCountedPtr<Stream>> MakeStream(
-      CallHandler call_handler,
-      InterActivityMutex<uint32_t>::Lock& next_stream_id_lock);
+  // Runs on the call party.
+  std::optional<RefCountedPtr<Stream>> MakeStream(CallHandler call_handler);
 
   struct CloseStreamArgs {
     bool close_reads;