Fix a TSAN reported race

I think this was the frequent crash in uds_cancel_after_invoke.

The race happens because a channel is deleted concurrently with an address being resolved (and UDS gets the resolution fast enough for this to actually happen).

The fix is to guarantee no callbacks will be made after cancel has been called (which was the original guaranteee that got lost somewhere).

4 files changed, 43 insertions, 2 deletions

diff --git a/src/core/channel/client_setup.c b/src/core/channel/client_setup.c
index bb6d363807..6d892d6c92 100644
--- a/src/core/channel/client_setup.c
+++ b/src/core/channel/client_setup.c
@@ -49,8 +49,11 @@ struct grpc_client_setup {
   grpc_alarm backoff_alarm;
   gpr_timespec current_backoff_interval;
   int in_alarm;
+  int in_cb;
+  int cancelled;
 
   gpr_mu mu;
+  gpr_cv cv;
   grpc_client_setup_request *active_request;
   int refs;
 };
@@ -67,6 +70,7 @@ gpr_timespec grpc_client_setup_request_deadline(grpc_client_setup_request *r) {
 
 static void destroy_setup(grpc_client_setup *s) {
   gpr_mu_destroy(&s->mu);
+  gpr_cv_destroy(&s->cv);
   s->done(s->user_data);
   grpc_channel_args_destroy(s->args);
   gpr_free(s);
@@ -111,6 +115,10 @@ static void setup_cancel(grpc_transport_setup *sp) {
   int cancel_alarm = 0;
 
   gpr_mu_lock(&s->mu);
+  s->cancelled = 1;
+  while (s->in_cb) {
+    gpr_cv_wait(&s->cv, &s->mu, gpr_inf_future);
+  }
 
   GPR_ASSERT(s->refs > 0);
   /* effectively cancels the current request (if any) */
@@ -129,6 +137,24 @@ static void setup_cancel(grpc_transport_setup *sp) {
   }
 }
 
+int grpc_client_setup_cb_begin(grpc_client_setup_request *r) {
+  gpr_mu_lock(&r->setup->mu);
+  if (r->setup->cancelled) {
+    gpr_mu_unlock(&r->setup->mu);
+    return 0;
+  }
+  r->setup->in_cb++;
+  gpr_mu_unlock(&r->setup->mu);
+  return 1;
+}
+
+void grpc_client_setup_cb_end(grpc_client_setup_request *r) {
+  gpr_mu_lock(&r->setup->mu);
+  r->setup->in_cb--;
+  if (r->setup->cancelled) gpr_cv_signal(&r->setup->cv);
+  gpr_mu_unlock(&r->setup->mu);
+}
+
 /* vtable for transport setup */
 static const grpc_transport_setup_vtable setup_vtable = {setup_initiate,
                                                          setup_cancel};
@@ -142,6 +168,7 @@ void grpc_client_setup_create_and_attach(
 
   s->base.vtable = &setup_vtable;
   gpr_mu_init(&s->mu);
+  gpr_cv_init(&s->cv);
   s->refs = 1;
   s->mdctx = mdctx;
   s->initiate = initiate;
@@ -151,6 +178,8 @@ void grpc_client_setup_create_and_attach(
   s->args = grpc_channel_args_copy(args);
   s->current_backoff_interval = gpr_time_from_micros(1000000);
   s->in_alarm = 0;
+  s->in_cb = 0;
+  s->cancelled = 0;
 
   grpc_client_channel_set_transport_setup(newly_minted_channel, &s->base);
 }
diff --git a/src/core/channel/client_setup.h b/src/core/channel/client_setup.h
index 6ac3fe62f1..f2b64265bc 100644
--- a/src/core/channel/client_setup.h
+++ b/src/core/channel/client_setup.h
@@ -58,6 +58,12 @@ void grpc_client_setup_request_finish(grpc_client_setup_request *r,
 const grpc_channel_args *grpc_client_setup_get_channel_args(
     grpc_client_setup_request *r);
 
+/* Call before calling back into the setup listener, and call only if
+   this function returns 1. If it returns 1, also promise to call
+   grpc_client_setup_cb_end */
+int grpc_client_setup_cb_begin(grpc_client_setup_request *r);
+void grpc_client_setup_cb_end(grpc_client_setup_request *r);
+
 /* Get the deadline for a request passed in to initiate. Implementations should
    make a best effort to honor this deadline. */
 gpr_timespec grpc_client_setup_request_deadline(grpc_client_setup_request *r);
diff --git a/src/core/surface/channel_create.c b/src/core/surface/channel_create.c
index 7a5f62ed53..3104b1d00d 100644
--- a/src/core/surface/channel_create.c
+++ b/src/core/surface/channel_create.c
@@ -107,13 +107,16 @@ static void on_connect(void *rp, grpc_endpoint *tcp) {
     } else {
       return;
     }
-  } else {
+  } else if (grpc_client_setup_cb_begin(r->cs_request)) {
     grpc_create_chttp2_transport(
         r->setup->setup_callback, r->setup->setup_user_data,
         grpc_client_setup_get_channel_args(r->cs_request), tcp, NULL, 0,
         grpc_client_setup_get_mdctx(r->cs_request), 1);
+    grpc_client_setup_cb_end(r->cs_request);
     done(r, 1);
     return;
+  } else {
+    done(r, 0);
   }
 }
 
diff --git a/src/core/surface/secure_channel_create.c b/src/core/surface/secure_channel_create.c
index c6968f4b24..8e56868d42 100644
--- a/src/core/surface/secure_channel_create.c
+++ b/src/core/surface/secure_channel_create.c
@@ -97,12 +97,15 @@ static void on_secure_transport_setup_done(void *rp,
   if (status != GRPC_SECURITY_OK) {
     gpr_log(GPR_ERROR, "Secure transport setup failed with error %d.", status);
     done(r, 0);
-  } else {
+  } else if (grpc_client_setup_cb_begin(r->cs_request)) {
     grpc_create_chttp2_transport(
         r->setup->setup_callback, r->setup->setup_user_data,
         grpc_client_setup_get_channel_args(r->cs_request), secure_endpoint,
         NULL, 0, grpc_client_setup_get_mdctx(r->cs_request), 1);
+    grpc_client_setup_cb_end(r->cs_request);
     done(r, 1);
+  } else {
+    done(r, 0);
   }
 }