Fixes

4 files changed, 13 insertions, 5 deletions

diff --git a/src/core/channel/client_channel.c b/src/core/channel/client_channel.c
index fe43c50018..955b390dbf 100644
--- a/src/core/channel/client_channel.c
+++ b/src/core/channel/client_channel.c
@@ -359,7 +359,8 @@ static int cc_pick_subchannel(grpc_exec_ctx *exec_ctx, void *elemp,
 /* Constructor for call_data */
 static void init_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
                            grpc_call_element_args *args) {
-  grpc_subchannel_call_holder_init(elem->call_data, cc_pick_subchannel, elem);
+  grpc_subchannel_call_holder_init(elem->call_data, cc_pick_subchannel, elem,
+                                   args->call_stack);
 }
 
 /* Destructor for call_data */
@@ -381,6 +382,7 @@ static void init_channel_elem(grpc_exec_ctx *exec_ctx,
 
   gpr_mu_init(&chand->mu_config);
   grpc_closure_init(&chand->on_config_changed, cc_on_config_changed, chand);
+  chand->owning_stack = args->channel_stack;
 
   grpc_connectivity_state_init(&chand->state_tracker, GRPC_CHANNEL_IDLE,
                                "client_channel");
diff --git a/src/core/channel/client_uchannel.c b/src/core/channel/client_uchannel.c
index ff50ebfb60..f8c621b8eb 100644
--- a/src/core/channel/client_uchannel.c
+++ b/src/core/channel/client_uchannel.c
@@ -138,7 +138,7 @@ static int cuc_pick_subchannel(grpc_exec_ctx *exec_ctx, void *arg,
 static void cuc_init_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
                                grpc_call_element_args *args) {
   grpc_subchannel_call_holder_init(elem->call_data, cuc_pick_subchannel,
-                                   elem->channel_data);
+                                   elem->channel_data, args->call_stack);
 }
 
 /* Destructor for call_data */
diff --git a/src/core/channel/subchannel_call_holder.c b/src/core/channel/subchannel_call_holder.c
index 9875dd8080..1455cc0a24 100644
--- a/src/core/channel/subchannel_call_holder.c
+++ b/src/core/channel/subchannel_call_holder.c
@@ -57,7 +57,7 @@ static void retry_waiting_locked(grpc_exec_ctx *exec_ctx,
 void grpc_subchannel_call_holder_init(
     grpc_subchannel_call_holder *holder,
     grpc_subchannel_call_holder_pick_subchannel pick_subchannel,
-    void *pick_subchannel_arg) {
+    void *pick_subchannel_arg, grpc_call_stack *owning_call) {
   gpr_atm_rel_store(&holder->subchannel_call, 0);
   holder->pick_subchannel = pick_subchannel;
   holder->pick_subchannel_arg = pick_subchannel_arg;
@@ -67,6 +67,7 @@ void grpc_subchannel_call_holder_init(
   holder->waiting_ops_count = 0;
   holder->waiting_ops_capacity = 0;
   holder->creation_phase = GRPC_SUBCHANNEL_CALL_HOLDER_NOT_CREATING;
+  holder->owning_call = owning_call;
 }
 
 void grpc_subchannel_call_holder_destroy(grpc_exec_ctx *exec_ctx,
@@ -141,10 +142,12 @@ retry:
       op->send_initial_metadata != NULL) {
     holder->creation_phase = GRPC_SUBCHANNEL_CALL_HOLDER_PICKING_SUBCHANNEL;
     grpc_closure_init(&holder->next_step, subchannel_ready, holder);
+    GRPC_CALL_STACK_REF(holder->owning_call, "pick_subchannel");
     if (holder->pick_subchannel(
             exec_ctx, holder->pick_subchannel_arg, op->send_initial_metadata,
             &holder->connected_subchannel, &holder->next_step)) {
       holder->creation_phase = GRPC_SUBCHANNEL_CALL_HOLDER_NOT_CREATING;
+      GRPC_CALL_STACK_UNREF(exec_ctx, holder->owning_call, "pick_subchannel");
     }
   }
   /* if we've got a subchannel, then let's ask it to create a call */
@@ -171,8 +174,8 @@ static void subchannel_ready(grpc_exec_ctx *exec_ctx, void *arg, int success) {
              GRPC_SUBCHANNEL_CALL_HOLDER_PICKING_SUBCHANNEL);
   call = GET_CALL(holder);
   GPR_ASSERT(call == NULL || call == CANCELLED_CALL);
+  holder->creation_phase = GRPC_SUBCHANNEL_CALL_HOLDER_NOT_CREATING;
   if (holder->connected_subchannel == NULL) {
-    holder->creation_phase = GRPC_SUBCHANNEL_CALL_HOLDER_NOT_CREATING;
     fail_locked(exec_ctx, holder);
   } else {
     gpr_atm_rel_store(
@@ -182,6 +185,7 @@ static void subchannel_ready(grpc_exec_ctx *exec_ctx, void *arg, int success) {
     retry_waiting_locked(exec_ctx, holder);
   }
   gpr_mu_unlock(&holder->mu);
+  GRPC_CALL_STACK_UNREF(exec_ctx, holder->owning_call, "pick_subchannel");
 }
 
 typedef struct {
diff --git a/src/core/channel/subchannel_call_holder.h b/src/core/channel/subchannel_call_holder.h
index 692c5e5316..9cf72c6cf7 100644
--- a/src/core/channel/subchannel_call_holder.h
+++ b/src/core/channel/subchannel_call_holder.h
@@ -78,12 +78,14 @@ typedef struct grpc_subchannel_call_holder {
   size_t waiting_ops_capacity;
 
   grpc_closure next_step;
+
+  grpc_call_stack *owning_call;
 } grpc_subchannel_call_holder;
 
 void grpc_subchannel_call_holder_init(
     grpc_subchannel_call_holder *holder,
     grpc_subchannel_call_holder_pick_subchannel pick_subchannel,
-    void *pick_subchannel_arg);
+    void *pick_subchannel_arg, grpc_call_stack *owning_call);
 void grpc_subchannel_call_holder_destroy(grpc_exec_ctx *exec_ctx,
                                          grpc_subchannel_call_holder *holder);