Fix metadata batch removal ref counting

8 files changed, 36 insertions, 22 deletions

diff --git a/src/core/lib/channel/compress_filter.c b/src/core/lib/channel/compress_filter.c
index 706c8df90c..330688c29c 100644
--- a/src/core/lib/channel/compress_filter.c
+++ b/src/core/lib/channel/compress_filter.c
@@ -134,7 +134,7 @@ static grpc_error *process_send_initial_metadata(
     calld->has_compression_algorithm = 1;
 
     grpc_metadata_batch_remove(
-        initial_metadata,
+        exec_ctx, initial_metadata,
         initial_metadata->idx.named.grpc_internal_encoding_request);
   } else {
     /* If no algorithm was found in the metadata and we aren't
diff --git a/src/core/lib/channel/http_client_filter.c b/src/core/lib/channel/http_client_filter.c
index 4593a9cb6d..426377ccc4 100644
--- a/src/core/lib/channel/http_client_filter.c
+++ b/src/core/lib/channel/http_client_filter.c
@@ -99,7 +99,7 @@ static grpc_error *client_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
                                                    grpc_metadata_batch *b) {
   if (b->idx.named.status != NULL) {
     if (grpc_mdelem_eq(b->idx.named.status->md, GRPC_MDELEM_STATUS_200)) {
-      grpc_metadata_batch_remove(b, b->idx.named.status);
+      grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.status);
     } else {
       char *val = grpc_dump_slice(GRPC_MDVALUE(b->idx.named.status->md),
                                   GPR_DUMP_ASCII);
@@ -150,7 +150,7 @@ static grpc_error *client_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
         gpr_free(val);
       }
     }
-    grpc_metadata_batch_remove(b, b->idx.named.content_type);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.content_type);
   }
 
   return GRPC_ERROR_NONE;
@@ -201,10 +201,11 @@ static void send_done(grpc_exec_ctx *exec_ctx, void *elemp, grpc_error *error) {
   calld->post_send->cb(exec_ctx, calld->post_send->cb_arg, error);
 }
 
-static void remove_if_present(grpc_metadata_batch *batch,
+static void remove_if_present(grpc_exec_ctx *exec_ctx,
+                              grpc_metadata_batch *batch,
                               grpc_metadata_batch_callouts_index idx) {
   if (batch->idx.array[idx] != NULL) {
-    grpc_metadata_batch_remove(batch, batch->idx.array[idx]);
+    grpc_metadata_batch_remove(exec_ctx, batch, batch->idx.array[idx]);
   }
 }
 
@@ -303,11 +304,13 @@ static grpc_error *hc_mutate_op(grpc_exec_ctx *exec_ctx,
       }
     }
 
-    remove_if_present(op->send_initial_metadata, GRPC_BATCH_METHOD);
-    remove_if_present(op->send_initial_metadata, GRPC_BATCH_SCHEME);
-    remove_if_present(op->send_initial_metadata, GRPC_BATCH_TE);
-    remove_if_present(op->send_initial_metadata, GRPC_BATCH_CONTENT_TYPE);
-    remove_if_present(op->send_initial_metadata, GRPC_BATCH_USER_AGENT);
+    remove_if_present(exec_ctx, op->send_initial_metadata, GRPC_BATCH_METHOD);
+    remove_if_present(exec_ctx, op->send_initial_metadata, GRPC_BATCH_SCHEME);
+    remove_if_present(exec_ctx, op->send_initial_metadata, GRPC_BATCH_TE);
+    remove_if_present(exec_ctx, op->send_initial_metadata,
+                      GRPC_BATCH_CONTENT_TYPE);
+    remove_if_present(exec_ctx, op->send_initial_metadata,
+                      GRPC_BATCH_USER_AGENT);
 
     /* Send : prefixed headers, which have to be before any application
        layer headers. */
diff --git a/src/core/lib/channel/http_server_filter.c b/src/core/lib/channel/http_server_filter.c
index 07c67e54c1..911589cc51 100644
--- a/src/core/lib/channel/http_server_filter.c
+++ b/src/core/lib/channel/http_server_filter.c
@@ -128,7 +128,7 @@ static grpc_error *server_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
                 grpc_attach_md_to_error(GRPC_ERROR_CREATE("Bad header"),
                                         b->idx.named.method->md));
     }
-    grpc_metadata_batch_remove(b, b->idx.named.method);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.method);
   } else {
     add_error(error_name, &error,
               grpc_error_set_str(GRPC_ERROR_CREATE("Missing header"),
@@ -141,7 +141,7 @@ static grpc_error *server_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
                 grpc_attach_md_to_error(GRPC_ERROR_CREATE("Bad header"),
                                         b->idx.named.te->md));
     }
-    grpc_metadata_batch_remove(b, b->idx.named.te);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.te);
   } else {
     add_error(error_name, &error,
               grpc_error_set_str(GRPC_ERROR_CREATE("Missing header"),
@@ -156,7 +156,7 @@ static grpc_error *server_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
                 grpc_attach_md_to_error(GRPC_ERROR_CREATE("Bad header"),
                                         b->idx.named.scheme->md));
     }
-    grpc_metadata_batch_remove(b, b->idx.named.scheme);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.scheme);
   } else {
     add_error(error_name, &error,
               grpc_error_set_str(GRPC_ERROR_CREATE("Missing header"),
@@ -189,7 +189,7 @@ static grpc_error *server_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
         gpr_free(val);
       }
     }
-    grpc_metadata_batch_remove(b, b->idx.named.content_type);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.content_type);
   }
 
   if (b->idx.named.path == NULL) {
@@ -220,7 +220,7 @@ static grpc_error *server_filter_incoming_metadata(grpc_exec_ctx *exec_ctx,
                               GRPC_MDVALUE(b->idx.named.grpc_payload_bin->md)));
     grpc_slice_buffer_stream_init(&calld->read_stream,
                                   &calld->read_slice_buffer, 0);
-    grpc_metadata_batch_remove(b, b->idx.named.grpc_payload_bin);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.grpc_payload_bin);
   }
 
   return error;
diff --git a/src/core/lib/surface/call.c b/src/core/lib/surface/call.c
index d5e90ccdda..7a7c19ad8b 100644
--- a/src/core/lib/surface/call.c
+++ b/src/core/lib/surface/call.c
@@ -866,7 +866,7 @@ static void recv_common_filter(grpc_exec_ctx *exec_ctx, grpc_call *call,
     GPR_TIMER_BEGIN("status", 0);
     set_status_code(call, STATUS_FROM_WIRE,
                     decode_status(b->idx.named.grpc_status->md));
-    grpc_metadata_batch_remove(b, b->idx.named.grpc_status);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.grpc_status);
     GPR_TIMER_END("status", 0);
   }
 
@@ -875,7 +875,7 @@ static void recv_common_filter(grpc_exec_ctx *exec_ctx, grpc_call *call,
     set_status_details(
         exec_ctx, call, STATUS_FROM_WIRE,
         grpc_slice_ref_internal(GRPC_MDVALUE(b->idx.named.grpc_message->md)));
-    grpc_metadata_batch_remove(b, b->idx.named.grpc_message);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.grpc_message);
     GPR_TIMER_END("status-details", 0);
   }
 }
@@ -910,14 +910,14 @@ static void recv_initial_filter(grpc_exec_ctx *exec_ctx, grpc_call *call,
     set_incoming_compression_algorithm(
         call, decode_compression(b->idx.named.grpc_encoding->md));
     GPR_TIMER_END("incoming_compression_algorithm", 0);
-    grpc_metadata_batch_remove(b, b->idx.named.grpc_encoding);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.grpc_encoding);
   }
 
   if (b->idx.named.grpc_accept_encoding != NULL) {
     GPR_TIMER_BEGIN("encodings_accepted_by_peer", 0);
     set_encodings_accepted_by_peer(exec_ctx, call,
                                    b->idx.named.grpc_accept_encoding->md);
-    grpc_metadata_batch_remove(b, b->idx.named.grpc_accept_encoding);
+    grpc_metadata_batch_remove(exec_ctx, b, b->idx.named.grpc_accept_encoding);
     GPR_TIMER_END("encodings_accepted_by_peer", 0);
   }
 
diff --git a/src/core/lib/transport/metadata.c b/src/core/lib/transport/metadata.c
index 537b60f739..9f13c7ded6 100644
--- a/src/core/lib/transport/metadata.c
+++ b/src/core/lib/transport/metadata.c
@@ -260,6 +260,14 @@ grpc_mdelem grpc_mdelem_create(
     allocated->key = grpc_slice_ref_internal(key);
     allocated->value = grpc_slice_ref_internal(value);
     gpr_atm_rel_store(&allocated->refcnt, 1);
+#ifdef GRPC_METADATA_REFCOUNT_DEBUG
+    char *key_str = grpc_dump_slice(allocated->key, GPR_DUMP_ASCII);
+    char *value_str = grpc_dump_slice(allocated->value, GPR_DUMP_ASCII);
+    gpr_log(GPR_DEBUG, "ELM ALLOC:%p:%zu: '%s' = '%s'", (void *)allocated,
+            gpr_atm_no_barrier_load(&allocated->refcnt), key_str, value_str);
+    gpr_free(key_str);
+    gpr_free(value_str);
+#endif
     return GRPC_MAKE_MDELEM(allocated, GRPC_MDELEM_STORAGE_ALLOCATED);
   }
 
diff --git a/src/core/lib/transport/metadata.h b/src/core/lib/transport/metadata.h
index f4ba86c854..aad944c973 100644
--- a/src/core/lib/transport/metadata.h
+++ b/src/core/lib/transport/metadata.h
@@ -148,7 +148,7 @@ void *grpc_mdelem_set_user_data(grpc_mdelem md, void (*destroy_func)(void *),
                                 void *user_data);
 
 /* Reference counting */
-//#define GRPC_METADATA_REFCOUNT_DEBUG
+#define GRPC_METADATA_REFCOUNT_DEBUG
 #ifdef GRPC_METADATA_REFCOUNT_DEBUG
 #define GRPC_MDELEM_REF(s) grpc_mdelem_ref((s), __FILE__, __LINE__)
 #define GRPC_MDELEM_UNREF(exec_ctx, s) \
diff --git a/src/core/lib/transport/metadata_batch.c b/src/core/lib/transport/metadata_batch.c
index 25cc65d5f1..a83f0e0f34 100644
--- a/src/core/lib/transport/metadata_batch.c
+++ b/src/core/lib/transport/metadata_batch.c
@@ -228,11 +228,13 @@ static void unlink_storage(grpc_mdelem_list *list,
   assert_valid_list(list);
 }
 
-void grpc_metadata_batch_remove(grpc_metadata_batch *batch,
+void grpc_metadata_batch_remove(grpc_exec_ctx *exec_ctx,
+                                grpc_metadata_batch *batch,
                                 grpc_linked_mdelem *storage) {
   assert_valid_callouts(batch);
   maybe_unlink_callout(batch, storage);
   unlink_storage(&batch->list, storage);
+  GRPC_MDELEM_UNREF(exec_ctx, storage->md);
   assert_valid_callouts(batch);
 }
 
diff --git a/src/core/lib/transport/metadata_batch.h b/src/core/lib/transport/metadata_batch.h
index a95a8887f4..d7f6485565 100644
--- a/src/core/lib/transport/metadata_batch.h
+++ b/src/core/lib/transport/metadata_batch.h
@@ -81,7 +81,8 @@ bool grpc_metadata_batch_is_empty(grpc_metadata_batch *batch);
 size_t grpc_metadata_batch_size(grpc_metadata_batch *batch);
 
 /** Remove \a storage from the batch, unreffing the mdelem contained */
-void grpc_metadata_batch_remove(grpc_metadata_batch *batch,
+void grpc_metadata_batch_remove(grpc_exec_ctx *exec_ctx,
+                                grpc_metadata_batch *batch,
                                 grpc_linked_mdelem *storage);
 
 /** Substitute a new mdelem for an old value */