Use call context to propagare LR costs

3 files changed, 24 insertions, 31 deletions

diff --git a/src/core/ext/load_reporting/load_reporting.c b/src/core/ext/load_reporting/load_reporting.c
index 37b06a737f..942aea4fd1 100644
--- a/src/core/ext/load_reporting/load_reporting.c
+++ b/src/core/ext/load_reporting/load_reporting.c
@@ -34,14 +34,34 @@
 #include <limits.h>
 #include <string.h>
 
+#include <grpc/load_reporting.h>
 #include <grpc/support/alloc.h>
 #include <grpc/support/sync.h>
 
 #include "src/core/ext/load_reporting/load_reporting.h"
 #include "src/core/ext/load_reporting/load_reporting_filter.h"
 #include "src/core/lib/channel/channel_stack_builder.h"
+#include "src/core/lib/slice/slice_internal.h"
+#include "src/core/lib/surface/call.h"
 #include "src/core/lib/surface/channel_init.h"
 
+static void destroy_lr_cost_context(void *c) {
+  grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
+  grpc_load_reporting_cost_context *cost_ctx = c;
+  for (size_t i = 0; i < cost_ctx->values_count; ++i) {
+    grpc_slice_unref_internal(&exec_ctx, cost_ctx->values[i]);
+  }
+  grpc_exec_ctx_finish(&exec_ctx);
+  gpr_free(cost_ctx->values);
+  gpr_free(cost_ctx);
+}
+
+void grpc_call_set_load_reporting_cost_context(
+    grpc_call *call, grpc_load_reporting_cost_context *ctx) {
+  grpc_call_context_set(call, GRPC_CONTEXT_LR_COST, ctx,
+                        destroy_lr_cost_context);
+}
+
 static bool is_load_reporting_enabled(const grpc_channel_args *a) {
   if (a == NULL) return false;
   for (size_t i = 0; i < a->num_args; i++) {
diff --git a/src/core/ext/load_reporting/load_reporting.h b/src/core/ext/load_reporting/load_reporting.h
index a157844734..22859a599a 100644
--- a/src/core/ext/load_reporting/load_reporting.h
+++ b/src/core/ext/load_reporting/load_reporting.h
@@ -35,23 +35,8 @@
 #define GRPC_CORE_EXT_LOAD_REPORTING_LOAD_REPORTING_H
 
 #include <grpc/impl/codegen/grpc_types.h>
-#include "src/core/lib/channel/channel_stack.h"
-
-/** Metadata key for the gRPC LB load balancer token.
- *
- * The value corresponding to this key is an opaque token that is given to the
- * frontend as part of each pick; the frontend sends this token to the backend
- * in each request it sends when using that pick. The token is used by the
- * backend to verify the request and to allow the backend to report load to the
- * gRPC LB system. */
-#define GRPC_LB_TOKEN_MD_KEY "lb-token"
 
-/** Metadata key for gRPC LB cost reporting.
- *
- * The value corresponding to this key is an opaque binary blob reported by the
- * backend as part of its trailing metadata containing cost information for the
- * call. */
-#define GRPC_LB_COST_MD_KEY "lb-cost-bin"
+#include "src/core/lib/channel/channel_stack.h"
 
 /** Identifiers for the invocation point of the users LR callback */
 typedef enum grpc_load_reporting_source {
diff --git a/src/core/ext/load_reporting/load_reporting_filter.c b/src/core/ext/load_reporting/load_reporting_filter.c
index 8af6191c3b..7fc88db603 100644
--- a/src/core/ext/load_reporting/load_reporting_filter.c
+++ b/src/core/ext/load_reporting/load_reporting_filter.c
@@ -31,11 +31,13 @@
  *
  */
 
+#include <string.h>
+
+#include <grpc/load_reporting.h>
 #include <grpc/support/alloc.h>
 #include <grpc/support/log.h>
 #include <grpc/support/string_util.h>
 #include <grpc/support/sync.h>
-#include <string.h>
 
 #include "src/core/ext/load_reporting/load_reporting.h"
 #include "src/core/ext/load_reporting/load_reporting_filter.h"
@@ -46,8 +48,6 @@
 
 typedef struct call_data {
   intptr_t id; /**< an id unique to the call */
-  bool have_trailing_md_string;
-  grpc_slice trailing_md_string;
   bool have_initial_md_string;
   grpc_slice initial_md_string;
   bool have_service_method;
@@ -142,9 +142,6 @@ static void destroy_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
   if (calld->have_initial_md_string) {
     grpc_slice_unref_internal(exec_ctx, calld->initial_md_string);
   }
-  if (calld->have_trailing_md_string) {
-    grpc_slice_unref_internal(exec_ctx, calld->trailing_md_string);
-  }
   if (calld->have_service_method) {
     grpc_slice_unref_internal(exec_ctx, calld->service_method);
   }
@@ -201,15 +198,6 @@ static void lr_start_transport_stream_op(grpc_exec_ctx *exec_ctx,
     /* substitute our callback for the higher callback */
     calld->ops_recv_initial_metadata_ready = op->recv_initial_metadata_ready;
     op->recv_initial_metadata_ready = &calld->on_initial_md_ready;
-  } else if (op->send_trailing_metadata) {
-    if (op->send_trailing_metadata->idx.named.lb_cost_bin != NULL) {
-      calld->trailing_md_string = grpc_slice_ref_internal(
-          GRPC_MDVALUE(op->send_trailing_metadata->idx.named.lb_cost_bin->md));
-      calld->have_trailing_md_string = true;
-      grpc_metadata_batch_remove(
-          exec_ctx, op->send_trailing_metadata,
-          op->send_trailing_metadata->idx.named.lb_cost_bin);
-    }
   }
   grpc_call_next_op(exec_ctx, elem, op);