grpc 1.40.0.pre1 → 1.42.0.pre1

data/include/grpc/impl/codegen/slice.h

@@ -19,6 +19,8 @@
 #ifndef GRPC_IMPL_CODEGEN_SLICE_H
 #define GRPC_IMPL_CODEGEN_SLICE_H
 
+// IWYU pragma: private, include <grpc/slice.h>
+
 #include <grpc/impl/codegen/port_platform.h>
 
 #include <stddef.h>

data/include/grpc/impl/codegen/status.h

@@ -19,6 +19,8 @@
 #ifndef GRPC_IMPL_CODEGEN_STATUS_H
 #define GRPC_IMPL_CODEGEN_STATUS_H
 
+// IWYU pragma: private, include <grpc/status.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif

data/include/grpc/impl/codegen/sync.h

@@ -18,6 +18,9 @@
 
 #ifndef GRPC_IMPL_CODEGEN_SYNC_H
 #define GRPC_IMPL_CODEGEN_SYNC_H
+
+// IWYU pragma: private, include <grpc/support/sync.h>
+
 /** Synchronization primitives for GPR.
 
    The type  gpr_mu              provides a non-reentrant mutex (lock).
@@ -44,16 +47,16 @@ extern "C" {
 /* Platform-specific type declarations of gpr_mu and gpr_cv.   */
 #include <grpc/impl/codegen/port_platform.h>
 
-#include <grpc/impl/codegen/sync_generic.h>
+#include <grpc/impl/codegen/sync_generic.h>  // IWYU pragma: export
 
 #if defined(GPR_CUSTOM_SYNC)
-#include <grpc/impl/codegen/sync_custom.h>
+#include <grpc/impl/codegen/sync_custom.h>  // IWYU pragma: export
 #elif defined(GPR_ABSEIL_SYNC)
-#include <grpc/impl/codegen/sync_abseil.h>
+#include <grpc/impl/codegen/sync_abseil.h>  // IWYU pragma: export
 #elif defined(GPR_POSIX_SYNC)
-#include <grpc/impl/codegen/sync_posix.h>
+#include <grpc/impl/codegen/sync_posix.h>  // IWYU pragma: export
 #elif defined(GPR_WINDOWS)
-#include <grpc/impl/codegen/sync_windows.h>
+#include <grpc/impl/codegen/sync_windows.h>  // IWYU pragma: export
 #else
 #error Unable to determine platform for sync
 #endif

data/include/grpc/impl/codegen/sync_abseil.h

@@ -19,6 +19,8 @@
 #ifndef GRPC_IMPL_CODEGEN_SYNC_ABSEIL_H
 #define GRPC_IMPL_CODEGEN_SYNC_ABSEIL_H
 
+// IWYU pragma: private, include <grpc/support/sync.h>
+
 #include <grpc/impl/codegen/port_platform.h>
 
 #include <grpc/impl/codegen/sync_generic.h>

data/include/grpc/impl/codegen/sync_custom.h

@@ -19,6 +19,8 @@
 #ifndef GRPC_IMPL_CODEGEN_SYNC_CUSTOM_H
 #define GRPC_IMPL_CODEGEN_SYNC_CUSTOM_H
 
+// IWYU pragma: private, include <grpc/support/sync.h>
+
 #include <grpc/impl/codegen/port_platform.h>
 
 #include <grpc/impl/codegen/sync_generic.h>

data/include/grpc/impl/codegen/sync_generic.h

@@ -18,6 +18,9 @@
 
 #ifndef GRPC_IMPL_CODEGEN_SYNC_GENERIC_H
 #define GRPC_IMPL_CODEGEN_SYNC_GENERIC_H
+
+// IWYU pragma: private, include <grpc/support/sync.h>
+
 /* Generic type definitions for gpr_sync. */
 
 #include <grpc/impl/codegen/port_platform.h>

data/include/grpc/impl/codegen/sync_posix.h

@@ -19,12 +19,14 @@
 #ifndef GRPC_IMPL_CODEGEN_SYNC_POSIX_H
 #define GRPC_IMPL_CODEGEN_SYNC_POSIX_H
 
-#include <grpc/impl/codegen/port_platform.h>
+// IWYU pragma: private, include <grpc/support/sync.h>
 
-#include <grpc/impl/codegen/sync_generic.h>
+#include <grpc/impl/codegen/port_platform.h>
 
 #include <pthread.h>
 
+#include <grpc/impl/codegen/sync_generic.h>
+
 #ifdef GRPC_ASAN_ENABLED
 /* The member |leak_checker| is used to check whether there is a memory leak
  * caused by upper layer logic that's missing the |gpr_xx_destroy| call

data/include/grpc/impl/codegen/sync_windows.h

@@ -19,6 +19,8 @@
 #ifndef GRPC_IMPL_CODEGEN_SYNC_WINDOWS_H
 #define GRPC_IMPL_CODEGEN_SYNC_WINDOWS_H
 
+// IWYU pragma: private, include <grpc/support/sync.h>
+
 #include <grpc/impl/codegen/port_platform.h>
 
 #ifdef GPR_WINDOWS

data/include/grpc/slice.h

@@ -21,7 +21,7 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/slice.h>
+#include <grpc/impl/codegen/slice.h>  // IWYU pragma: export
 #include <grpc/support/sync.h>
 
 #ifdef __cplusplus

data/include/grpc/status.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/status.h>
+#include <grpc/impl/codegen/status.h>  // IWYU pragma: export
 
 #endif /* GRPC_STATUS_H */

data/include/grpc/support/atm.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/atm.h>
+#include <grpc/impl/codegen/atm.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_ATM_H */

data/include/grpc/support/atm_gcc_atomic.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/atm_gcc_atomic.h>
+#include <grpc/impl/codegen/atm_gcc_atomic.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_ATM_GCC_ATOMIC_H */

data/include/grpc/support/atm_gcc_sync.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/atm_gcc_sync.h>
+#include <grpc/impl/codegen/atm_gcc_sync.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_ATM_GCC_SYNC_H */

data/include/grpc/support/atm_windows.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/atm_windows.h>
+#include <grpc/impl/codegen/atm_windows.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_ATM_WINDOWS_H */

data/include/grpc/support/log.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/log.h>
+#include <grpc/impl/codegen/log.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_LOG_H */

data/include/grpc/support/port_platform.h

@@ -19,6 +19,6 @@
 #ifndef GRPC_SUPPORT_PORT_PLATFORM_H
 #define GRPC_SUPPORT_PORT_PLATFORM_H
 
-#include <grpc/impl/codegen/port_platform.h>
+#include <grpc/impl/codegen/port_platform.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_PORT_PLATFORM_H */

data/include/grpc/support/sync.h

@@ -22,7 +22,7 @@
 #include <grpc/support/port_platform.h>
 
 #include <grpc/impl/codegen/gpr_types.h> /* for gpr_timespec */
-#include <grpc/impl/codegen/sync.h>
+#include <grpc/impl/codegen/sync.h>      // IWYU pragma: export
 
 #ifdef __cplusplus
 extern "C" {

data/include/grpc/support/sync_abseil.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/sync_abseil.h>
+#include <grpc/impl/codegen/sync_abseil.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_SYNC_ABSEIL_H */

data/include/grpc/support/sync_custom.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/sync_custom.h>
+#include <grpc/impl/codegen/sync_custom.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_SYNC_CUSTOM_H */

data/include/grpc/support/sync_generic.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/sync_generic.h>
+#include <grpc/impl/codegen/sync_generic.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_SYNC_GENERIC_H */

data/include/grpc/support/sync_posix.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/sync_posix.h>
+#include <grpc/impl/codegen/sync_posix.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_SYNC_POSIX_H */

data/include/grpc/support/sync_windows.h

@@ -21,6 +21,6 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/sync_windows.h>
+#include <grpc/impl/codegen/sync_windows.h>  // IWYU pragma: export
 
 #endif /* GRPC_SUPPORT_SYNC_WINDOWS_H */

data/include/grpc/support/time.h

@@ -21,11 +21,11 @@
 
 #include <grpc/support/port_platform.h>
 
-#include <grpc/impl/codegen/gpr_types.h>
-
 #include <stddef.h>
 #include <time.h>
 
+#include <grpc/impl/codegen/gpr_types.h>
+
 #ifdef __cplusplus
 extern "C" {
 #endif

data/src/core/ext/filters/census/grpc_context.cc

@@ -20,6 +20,7 @@
 
 #include <grpc/census.h>
 #include <grpc/grpc.h>
+
 #include "src/core/lib/surface/api_trace.h"
 #include "src/core/lib/surface/call.h"
 

data/src/core/ext/filters/client_channel/backend_metric.cc

@@ -19,9 +19,8 @@
 #include "src/core/ext/filters/client_channel/backend_metric.h"
 
 #include "absl/strings/string_view.h"
-
-#include "udpa/data/orca/v1/orca_load_report.upb.h"
 #include "upb/upb.hpp"
+#include "xds/data/orca/v3/orca_load_report.upb.h"
 
 namespace grpc_core {
 
@@ -29,8 +28,8 @@ namespace {
 
 template <typename EntryType>
 std::map<absl::string_view, double> ParseMap(
-    udpa_data_orca_v1_OrcaLoadReport* msg,
-    const EntryType* (*entry_func)(const udpa_data_orca_v1_OrcaLoadReport*,
+    xds_data_orca_v3_OrcaLoadReport* msg,
+    const EntryType* (*entry_func)(const xds_data_orca_v3_OrcaLoadReport*,
                                    size_t*),
     upb_strview (*key_func)(const EntryType*),
     double (*value_func)(const EntryType*), Arena* arena) {
@@ -52,30 +51,29 @@ std::map<absl::string_view, double> ParseMap(
 const LoadBalancingPolicy::BackendMetricData* ParseBackendMetricData(
     const grpc_slice& serialized_load_report, Arena* arena) {
   upb::Arena upb_arena;
-  udpa_data_orca_v1_OrcaLoadReport* msg =
-      udpa_data_orca_v1_OrcaLoadReport_parse(
-          reinterpret_cast<const char*>(
-              GRPC_SLICE_START_PTR(serialized_load_report)),
-          GRPC_SLICE_LENGTH(serialized_load_report), upb_arena.ptr());
+  xds_data_orca_v3_OrcaLoadReport* msg = xds_data_orca_v3_OrcaLoadReport_parse(
+      reinterpret_cast<const char*>(
+          GRPC_SLICE_START_PTR(serialized_load_report)),
+      GRPC_SLICE_LENGTH(serialized_load_report), upb_arena.ptr());
   if (msg == nullptr) return nullptr;
   LoadBalancingPolicy::BackendMetricData* backend_metric_data =
       arena->New<LoadBalancingPolicy::BackendMetricData>();
   backend_metric_data->cpu_utilization =
-      udpa_data_orca_v1_OrcaLoadReport_cpu_utilization(msg);
+      xds_data_orca_v3_OrcaLoadReport_cpu_utilization(msg);
   backend_metric_data->mem_utilization =
-      udpa_data_orca_v1_OrcaLoadReport_mem_utilization(msg);
+      xds_data_orca_v3_OrcaLoadReport_mem_utilization(msg);
   backend_metric_data->requests_per_second =
-      udpa_data_orca_v1_OrcaLoadReport_rps(msg);
+      xds_data_orca_v3_OrcaLoadReport_rps(msg);
   backend_metric_data->request_cost =
-      ParseMap<udpa_data_orca_v1_OrcaLoadReport_RequestCostEntry>(
-          msg, udpa_data_orca_v1_OrcaLoadReport_request_cost_next,
-          udpa_data_orca_v1_OrcaLoadReport_RequestCostEntry_key,
-          udpa_data_orca_v1_OrcaLoadReport_RequestCostEntry_value, arena);
+      ParseMap<xds_data_orca_v3_OrcaLoadReport_RequestCostEntry>(
+          msg, xds_data_orca_v3_OrcaLoadReport_request_cost_next,
+          xds_data_orca_v3_OrcaLoadReport_RequestCostEntry_key,
+          xds_data_orca_v3_OrcaLoadReport_RequestCostEntry_value, arena);
   backend_metric_data->utilization =
-      ParseMap<udpa_data_orca_v1_OrcaLoadReport_UtilizationEntry>(
-          msg, udpa_data_orca_v1_OrcaLoadReport_utilization_next,
-          udpa_data_orca_v1_OrcaLoadReport_UtilizationEntry_key,
-          udpa_data_orca_v1_OrcaLoadReport_UtilizationEntry_value, arena);
+      ParseMap<xds_data_orca_v3_OrcaLoadReport_UtilizationEntry>(
+          msg, xds_data_orca_v3_OrcaLoadReport_utilization_next,
+          xds_data_orca_v3_OrcaLoadReport_UtilizationEntry_key,
+          xds_data_orca_v3_OrcaLoadReport_UtilizationEntry_value, arena);
   return backend_metric_data;
 }
 

data/src/core/ext/filters/client_channel/backup_poller.cc

@@ -28,6 +28,7 @@
 #include "src/core/ext/filters/client_channel/client_channel.h"
 #include "src/core/lib/gpr/string.h"
 #include "src/core/lib/gprpp/global_config.h"
+#include "src/core/lib/gprpp/memory.h"
 #include "src/core/lib/iomgr/error.h"
 #include "src/core/lib/iomgr/iomgr.h"
 #include "src/core/lib/iomgr/pollset.h"
@@ -138,7 +139,7 @@ static void run_poller(void* arg, grpc_error_handle error) {
 
 static void g_poller_init_locked() {
   if (g_poller == nullptr) {
-    g_poller = static_cast<backup_poller*>(gpr_zalloc(sizeof(backup_poller)));
+    g_poller = grpc_core::Zalloc<backup_poller>();
     g_poller->pollset =
         static_cast<grpc_pollset*>(gpr_zalloc(grpc_pollset_size()));
     g_poller->shutting_down = false;

data/src/core/ext/filters/client_channel/backup_poller.h

@@ -22,6 +22,7 @@
 #include <grpc/support/port_platform.h>
 
 #include <grpc/grpc.h>
+
 #include "src/core/lib/channel/channel_stack.h"
 #include "src/core/lib/gprpp/global_config.h"
 

data/src/core/ext/filters/client_channel/channel_connectivity.cc

@@ -16,14 +16,24 @@
 
 #include <grpc/support/port_platform.h>
 
-#include "src/core/lib/surface/channel.h"
-
 #include <grpc/support/log.h>
 
 #include "src/core/ext/filters/client_channel/client_channel.h"
 #include "src/core/lib/iomgr/timer.h"
 #include "src/core/lib/surface/api_trace.h"
+#include "src/core/lib/surface/channel.h"
 #include "src/core/lib/surface/completion_queue.h"
+#include "src/core/lib/surface/lame_client.h"
+
+namespace {
+
+bool IsLameChannel(grpc_channel* channel) {
+  grpc_channel_element* elem =
+      grpc_channel_stack_last_element(grpc_channel_get_channel_stack(channel));
+  return elem->filter == &grpc_lame_filter;
+}
+
+}  // namespace
 
 grpc_connectivity_state grpc_channel_check_connectivity_state(
     grpc_channel* channel, int try_to_connect) {
@@ -36,6 +46,7 @@ grpc_connectivity_state grpc_channel_check_connectivity_state(
   grpc_core::ClientChannel* client_channel =
       grpc_core::ClientChannel::GetFromChannel(channel);
   if (GPR_UNLIKELY(client_channel == nullptr)) {
+    if (IsLameChannel(channel)) return GRPC_CHANNEL_TRANSIENT_FAILURE;
     gpr_log(GPR_ERROR,
             "grpc_channel_check_connectivity_state called on something that is "
             "not a client channel");
@@ -48,9 +59,11 @@ int grpc_channel_num_external_connectivity_watchers(grpc_channel* channel) {
   grpc_core::ClientChannel* client_channel =
       grpc_core::ClientChannel::GetFromChannel(channel);
   if (client_channel == nullptr) {
-    gpr_log(GPR_ERROR,
-            "grpc_channel_num_external_connectivity_watchers called on "
-            "something that is not a client channel");
+    if (!IsLameChannel(channel)) {
+      gpr_log(GPR_ERROR,
+              "grpc_channel_num_external_connectivity_watchers called on "
+              "something that is not a client channel");
+    }
     return 0;
   }
   return client_channel->NumExternalConnectivityWatchers();
@@ -63,7 +76,7 @@ int grpc_channel_support_connectivity_watcher(grpc_channel* channel) {
 namespace grpc_core {
 namespace {
 
-class StateWatcher {
+class StateWatcher : public DualRefCounted<StateWatcher> {
  public:
   StateWatcher(grpc_channel* channel, grpc_completion_queue* cq, void* tag,
                grpc_connectivity_state last_observed_state,
@@ -73,16 +86,35 @@ class StateWatcher {
     GRPC_CHANNEL_INTERNAL_REF(channel, "watch_channel_connectivity");
     GRPC_CLOSURE_INIT(&on_complete_, WatchComplete, this, nullptr);
     GRPC_CLOSURE_INIT(&on_timeout_, TimeoutComplete, this, nullptr);
+    ClientChannel* client_channel = ClientChannel::GetFromChannel(channel);
+    if (client_channel == nullptr) {
+      // If the target URI used to create the channel was invalid, channel
+      // stack initialization failed, and that caused us to create a lame
+      // channel.  In that case, connectivity state will never change (it
+      // will always be TRANSIENT_FAILURE), so we don't actually start a
+      // watch, but we are hiding that fact from the application.
+      if (IsLameChannel(channel)) {
+        // Ref from object creation is held by timer callback.
+        StartTimer(grpc_timespec_to_millis_round_up(deadline));
+        return;
+      }
+      gpr_log(GPR_ERROR,
+              "grpc_channel_watch_connectivity_state called on "
+              "something that is not a client channel");
+      GPR_ASSERT(false);
+    }
+    // Take an addition ref, so we have two (the first one is from the
+    // creation of this object).  One will be held by the timer callback,
+    // the other by the watcher callback.
+    Ref().release();
     auto* watcher_timer_init_state = new WatcherTimerInitState(
         this, grpc_timespec_to_millis_round_up(deadline));
-    ClientChannel* client_channel = ClientChannel::GetFromChannel(channel);
-    GPR_ASSERT(client_channel != nullptr);
     client_channel->AddExternalConnectivityWatcher(
         grpc_polling_entity_create_from_pollset(grpc_cq_pollset(cq)), &state_,
         &on_complete_, watcher_timer_init_state->closure());
   }
 
-  ~StateWatcher() {
+  ~StateWatcher() override {
     GRPC_CHANNEL_INTERNAL_UNREF(channel_, "watch_channel_connectivity");
   }
 
@@ -101,8 +133,7 @@ class StateWatcher {
    private:
     static void WatcherTimerInit(void* arg, grpc_error_handle /*error*/) {
       auto* self = static_cast<WatcherTimerInitState*>(arg);
-      grpc_timer_init(&self->state_watcher_->timer_, self->deadline_,
-                      &self->state_watcher_->on_timeout_);
+      self->state_watcher_->StartTimer(self->deadline_);
       delete self;
     }
 
@@ -111,94 +142,46 @@ class StateWatcher {
     grpc_closure closure_;
   };
 
-  enum CallbackPhase { kWaiting, kReadyToCallBack, kCallingBackAndFinished };
+  void StartTimer(grpc_millis deadline) {
+    grpc_timer_init(&timer_, deadline, &on_timeout_);
+  }
 
-  // Called when the completion is returned to the CQ.
-  static void FinishedCompletion(void* arg, grpc_cq_completion* /*ignored*/) {
+  static void WatchComplete(void* arg, grpc_error_handle error) {
     auto* self = static_cast<StateWatcher*>(arg);
-    bool should_delete = false;
-    {
-      MutexLock lock(&self->mu_);
-      switch (self->phase_) {
-        case kWaiting:
-        case kReadyToCallBack:
-          GPR_UNREACHABLE_CODE(return );
-        case kCallingBackAndFinished:
-          should_delete = true;
-      }
+    if (GRPC_TRACE_FLAG_ENABLED(grpc_trace_operation_failures)) {
+      GRPC_LOG_IF_ERROR("watch_completion_error", GRPC_ERROR_REF(error));
     }
-    if (should_delete) delete self;
+    grpc_timer_cancel(&self->timer_);
+    self->Unref();
   }
 
-  void PartlyDone(bool due_to_completion, grpc_error_handle error) {
-    bool end_op = false;
-    void* end_op_tag = nullptr;
-    grpc_error_handle end_op_error = GRPC_ERROR_NONE;
-    grpc_completion_queue* end_op_cq = nullptr;
-    grpc_cq_completion* end_op_completion_storage = nullptr;
-    if (due_to_completion) {
-      grpc_timer_cancel(&timer_);
-    } else {
-      grpc_core::ClientChannel* client_channel =
-          grpc_core::ClientChannel::GetFromChannel(channel_);
-      GPR_ASSERT(client_channel != nullptr);
-      client_channel->CancelExternalConnectivityWatcher(&on_complete_);
-    }
-    {
-      MutexLock lock(&mu_);
-      if (due_to_completion) {
-        if (GRPC_TRACE_FLAG_ENABLED(grpc_trace_operation_failures)) {
-          GRPC_LOG_IF_ERROR("watch_completion_error", GRPC_ERROR_REF(error));
-        }
-        GRPC_ERROR_UNREF(error);
-        error = GRPC_ERROR_NONE;
-      } else {
-        if (error == GRPC_ERROR_NONE) {
-          error = GRPC_ERROR_CREATE_FROM_STATIC_STRING(
-              "Timed out waiting for connection state change");
-        } else if (error == GRPC_ERROR_CANCELLED) {
-          error = GRPC_ERROR_NONE;
-        }
-      }
-      switch (phase_) {
-        case kWaiting:
-          GRPC_ERROR_REF(error);
-          error_ = error;
-          phase_ = kReadyToCallBack;
-          break;
-        case kReadyToCallBack:
-          if (error != GRPC_ERROR_NONE) {
-            GPR_ASSERT(!due_to_completion);
-            GRPC_ERROR_UNREF(error_);
-            GRPC_ERROR_REF(error);
-            error_ = error;
-          }
-          phase_ = kCallingBackAndFinished;
-          end_op = true;
-          end_op_cq = cq_;
-          end_op_tag = tag_;
-          end_op_error = error_;
-          end_op_completion_storage = &completion_storage_;
-          break;
-        case kCallingBackAndFinished:
-          GPR_UNREACHABLE_CODE(return );
-      }
-    }
-    if (end_op) {
-      grpc_cq_end_op(end_op_cq, end_op_tag, end_op_error, FinishedCompletion,
-                     this, end_op_completion_storage);
+  static void TimeoutComplete(void* arg, grpc_error_handle error) {
+    auto* self = static_cast<StateWatcher*>(arg);
+    self->timer_fired_ = error == GRPC_ERROR_NONE;
+    // If this is a client channel (not a lame channel), cancel the watch.
+    ClientChannel* client_channel =
+        ClientChannel::GetFromChannel(self->channel_);
+    if (client_channel != nullptr) {
+      client_channel->CancelExternalConnectivityWatcher(&self->on_complete_);
     }
-    GRPC_ERROR_UNREF(error);
+    self->Unref();
   }
 
-  static void WatchComplete(void* arg, grpc_error_handle error) {
-    auto* self = static_cast<StateWatcher*>(arg);
-    self->PartlyDone(/*due_to_completion=*/true, GRPC_ERROR_REF(error));
+  // Invoked when both strong refs are released.
+  void Orphan() override {
+    WeakRef().release();  // Take a weak ref until completion is finished.
+    grpc_error_handle error =
+        timer_fired_ ? GRPC_ERROR_CREATE_FROM_STATIC_STRING(
+                           "Timed out waiting for connection state change")
+                     : GRPC_ERROR_NONE;
+    grpc_cq_end_op(cq_, tag_, error, FinishedCompletion, this,
+                   &completion_storage_);
   }
 
-  static void TimeoutComplete(void* arg, grpc_error_handle error) {
+  // Called when the completion is returned to the CQ.
+  static void FinishedCompletion(void* arg, grpc_cq_completion* /*ignored*/) {
     auto* self = static_cast<StateWatcher*>(arg);
-    self->PartlyDone(/*due_to_completion=*/false, GRPC_ERROR_REF(error));
+    self->WeakUnref();
   }
 
   grpc_channel* channel_;
@@ -213,9 +196,7 @@ class StateWatcher {
   grpc_timer timer_;
   grpc_closure on_timeout_;
 
-  Mutex mu_;
-  CallbackPhase phase_ ABSL_GUARDED_BY(mu_) = kWaiting;
-  grpc_error_handle error_ ABSL_GUARDED_BY(mu_) = GRPC_ERROR_NONE;
+  bool timer_fired_ = false;
 };
 
 }  // namespace