From dbad2db8482f264944ba3b33b405dfbac1ebdef3 Mon Sep 17 00:00:00 2001
From: "Mark D. Roth" <roth@google.com>
Date: Mon, 30 Sep 2019 12:17:34 -0700
Subject: [PATCH] Immediately orphan watcher if state is SHUTDOWN when it is
 added.

---
 src/core/lib/transport/connectivity_state.cc  |  6 +++-
 src/core/lib/transport/connectivity_state.h   |  4 +++
 .../core/transport/connectivity_state_test.cc | 34 +++++++++++++++++++
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/src/core/lib/transport/connectivity_state.cc b/src/core/lib/transport/connectivity_state.cc
index 45ebdca4f0d..32aa99306be 100644
--- a/src/core/lib/transport/connectivity_state.cc
+++ b/src/core/lib/transport/connectivity_state.cc
@@ -120,7 +120,11 @@ void ConnectivityStateTracker::AddWatcher(
     }
     watcher->Notify(current_state);
   }
-  watchers_.insert(MakePair(watcher.get(), std::move(watcher)));
+  // If we're in state SHUTDOWN, don't add the watcher, so that it will
+  // be orphaned immediately.
+  if (current_state != GRPC_CHANNEL_SHUTDOWN) {
+    watchers_.insert(MakePair(watcher.get(), std::move(watcher)));
+  }
 }
 
 void ConnectivityStateTracker::RemoveWatcher(
diff --git a/src/core/lib/transport/connectivity_state.h b/src/core/lib/transport/connectivity_state.h
index 41f7bf08e41..25b43561364 100644
--- a/src/core/lib/transport/connectivity_state.h
+++ b/src/core/lib/transport/connectivity_state.h
@@ -76,6 +76,10 @@ class AsyncConnectivityStateWatcherInterface
 
 // Tracks connectivity state.  Maintains a list of watchers that are
 // notified whenever the state changes.
+//
+// Note that once the state becomes SHUTDOWN, watchers will be notified
+// and then automatically orphaned (i.e., RemoveWatcher() does not need
+// to be called).
 class ConnectivityStateTracker {
  public:
   ConnectivityStateTracker(const char* name,
diff --git a/test/core/transport/connectivity_state_test.cc b/test/core/transport/connectivity_state_test.cc
index e0141cd02d0..6493a62ff3c 100644
--- a/test/core/transport/connectivity_state_test.cc
+++ b/test/core/transport/connectivity_state_test.cc
@@ -113,6 +113,40 @@ TEST(StateTracker, SubscribeThenUnsubscribe) {
   EXPECT_EQ(state, GRPC_CHANNEL_IDLE);
 }
 
+TEST(StateTracker, OrphanUponShutdown) {
+  int count = 0;
+  grpc_connectivity_state state = GRPC_CHANNEL_IDLE;
+  bool destroyed = false;
+  ConnectivityStateTracker tracker("xxx", GRPC_CHANNEL_IDLE);
+  ConnectivityStateWatcherInterface* watcher =
+      New<Watcher>(&count, &state, &destroyed);
+  tracker.AddWatcher(GRPC_CHANNEL_IDLE,
+                     OrphanablePtr<ConnectivityStateWatcherInterface>(watcher));
+  // No initial notification, since we started the watch from the
+  // current state.
+  EXPECT_EQ(count, 0);
+  EXPECT_EQ(state, GRPC_CHANNEL_IDLE);
+  // Set state to SHUTDOWN.
+  tracker.SetState(GRPC_CHANNEL_SHUTDOWN, "shutting down");
+  EXPECT_TRUE(destroyed);
+  EXPECT_EQ(count, 1);
+  EXPECT_EQ(state, GRPC_CHANNEL_SHUTDOWN);
+}
+
+TEST(StateTracker, AddWhenAlreadyShutdown) {
+  int count = 0;
+  grpc_connectivity_state state = GRPC_CHANNEL_IDLE;
+  bool destroyed = false;
+  ConnectivityStateTracker tracker("xxx", GRPC_CHANNEL_SHUTDOWN);
+  ConnectivityStateWatcherInterface* watcher =
+      New<Watcher>(&count, &state, &destroyed);
+  tracker.AddWatcher(GRPC_CHANNEL_IDLE,
+                     OrphanablePtr<ConnectivityStateWatcherInterface>(watcher));
+  EXPECT_TRUE(destroyed);
+  EXPECT_EQ(count, 1);
+  EXPECT_EQ(state, GRPC_CHANNEL_SHUTDOWN);
+}
+
 TEST(StateTracker, NotifyShutdownAtDestruction) {
   int count = 0;
   grpc_connectivity_state state = GRPC_CHANNEL_IDLE;