From f17fbd4d539b3f88fd729730fb718d6c5455684f Mon Sep 17 00:00:00 2001
From: Yash Tibrewal <yashkt@google.com>
Date: Fri, 20 Sep 2024 00:02:32 +0000
Subject: [PATCH] [Chttp2Server] Fix race between connection manager updates
 and handshake

---
 .../transport/chttp2/server/chttp2_server.cc  | 57 ++++++++++++-------
 1 file changed, 36 insertions(+), 21 deletions(-)

diff --git a/src/core/ext/transport/chttp2/server/chttp2_server.cc b/src/core/ext/transport/chttp2/server/chttp2_server.cc
index 50f165cd325..7e17cfed7d1 100644
--- a/src/core/ext/transport/chttp2/server/chttp2_server.cc
+++ b/src/core/ext/transport/chttp2/server/chttp2_server.cc
@@ -183,24 +183,27 @@ class Chttp2ServerListener : public Server::ListenerInterface {
       void Start(OrphanablePtr<grpc_endpoint> endpoint,
                  const ChannelArgs& args);
 
+      void ShutdownLocked(absl::Status status)
+          ABSL_EXCLUSIVE_LOCKS_REQUIRED(&ActiveConnection::mu_);
+
       // Needed to be able to grab an external ref in
       // ActiveConnection::Start()
       using InternallyRefCounted<HandshakingState>::Ref;
 
      private:
-      void OnTimeout() ABSL_LOCKS_EXCLUDED(&connection_->mu_);
+      void OnTimeout() ABSL_LOCKS_EXCLUDED(&ActiveConnection::mu_);
       static void OnReceiveSettings(void* arg, grpc_error_handle /* error */);
       void OnHandshakeDone(absl::StatusOr<HandshakerArgs*> result);
       RefCountedPtr<ActiveConnection> const connection_;
       grpc_pollset* const accepting_pollset_;
       AcceptorPtr acceptor_;
       RefCountedPtr<HandshakeManager> handshake_mgr_
-          ABSL_GUARDED_BY(&connection_->mu_);
+          ABSL_GUARDED_BY(&ActiveConnection::mu_);
       // State for enforcing handshake timeout on receiving HTTP/2 settings.
       Timestamp const deadline_;
       absl::optional<EventEngine::TaskHandle> timer_handle_
-          ABSL_GUARDED_BY(&connection_->mu_);
-      grpc_closure on_receive_settings_ ABSL_GUARDED_BY(&connection_->mu_);
+          ABSL_GUARDED_BY(&ActiveConnection::mu_);
+      grpc_closure on_receive_settings_ ABSL_GUARDED_BY(&ActiveConnection::mu_);
       grpc_pollset_set* const interested_parties_;
     };
 
@@ -400,9 +403,7 @@ Chttp2ServerListener::ActiveConnection::HandshakingState::~HandshakingState() {
 void Chttp2ServerListener::ActiveConnection::HandshakingState::Orphan() {
   {
     MutexLock lock(&connection_->mu_);
-    if (handshake_mgr_ != nullptr) {
-      handshake_mgr_->Shutdown(GRPC_ERROR_CREATE("Listener stopped serving."));
-    }
+    ShutdownLocked(absl::UnavailableError("Listener stopped serving."));
   }
   Unref();
 }
@@ -422,6 +423,13 @@ void Chttp2ServerListener::ActiveConnection::HandshakingState::Start(
       });
 }
 
+void Chttp2ServerListener::ActiveConnection::HandshakingState::ShutdownLocked(
+    absl::Status status) {
+  if (handshake_mgr_ != nullptr) {
+    handshake_mgr_->Shutdown(std::move(status));
+  }
+}
+
 void Chttp2ServerListener::ActiveConnection::HandshakingState::OnTimeout() {
   grpc_chttp2_transport* transport = nullptr;
   {
@@ -584,20 +592,27 @@ void Chttp2ServerListener::ActiveConnection::SendGoAway() {
   grpc_chttp2_transport* transport = nullptr;
   {
     MutexLock lock(&mu_);
-    if (transport_ != nullptr && !shutdown_) {
-      transport = transport_.get();
-      drain_grace_timer_handle_ = event_engine_->RunAfter(
-          std::max(Duration::Zero(),
-                   listener_->args_
-                       .GetDurationFromIntMillis(
-                           GRPC_ARG_SERVER_CONFIG_CHANGE_DRAIN_GRACE_TIME_MS)
-                       .value_or(Duration::Minutes(10))),
-          [self = Ref(DEBUG_LOCATION, "drain_grace_timer")]() mutable {
-            ApplicationCallbackExecCtx callback_exec_ctx;
-            ExecCtx exec_ctx;
-            self->OnDrainGraceTimeExpiry();
-            self.reset(DEBUG_LOCATION, "drain_grace_timer");
-          });
+    if (!shutdown_) {
+      // Send a GOAWAY if the transport exists
+      if (transport_ != nullptr) {
+        transport = transport_.get();
+        drain_grace_timer_handle_ = event_engine_->RunAfter(
+            std::max(Duration::Zero(),
+                     listener_->args_
+                         .GetDurationFromIntMillis(
+                             GRPC_ARG_SERVER_CONFIG_CHANGE_DRAIN_GRACE_TIME_MS)
+                         .value_or(Duration::Minutes(10))),
+            [self = Ref(DEBUG_LOCATION, "drain_grace_timer")]() mutable {
+              ApplicationCallbackExecCtx callback_exec_ctx;
+              ExecCtx exec_ctx;
+              self->OnDrainGraceTimeExpiry();
+              self.reset(DEBUG_LOCATION, "drain_grace_timer");
+            });
+      }
+      if (handshaking_state_ != nullptr) {
+        handshaking_state_->ShutdownLocked(
+            absl::UnavailableError("Connection going away"));
+      }
       shutdown_ = true;
     }
   }