After addition of grpc_call_stats

src/core/ext/load_reporting/load_reporting_filter.c

@@ -31,53 +31,32 @@
  *
  */
 
-#include <string.h>
-
 #include <grpc/support/log.h>
 
 #include "src/core/ext/load_reporting/load_reporting_filter.h"
 #include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/load_reporting/load_reporting.h"
 #include "src/core/lib/profiling/timers.h"
-#include "src/core/lib/support/string.h"
-
-typedef struct call_data { load_reporting_data lr_data; } call_data;
 
+typedef struct call_data { void *dummy; } call_data;
 typedef struct channel_data { void *dummy; } channel_data;
 
-static void load_reporting_start_transport_stream_op(
-    grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
-    grpc_transport_stream_op *op) {
-  call_data *calld = elem->call_data;
-
-  GPR_TIMER_BEGIN("load_reporting_start_transport_stream_op", 0);
-  grpc_load_reporting_call(&calld->lr_data);
-  grpc_call_next_op(exec_ctx, elem, op);
-  GPR_TIMER_END("load_reporting_start_transport_stream_op", 0);
-}
-
 /* Constructor for call_data */
 static void init_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
-                           grpc_call_element_args *args) {
-  /* grab pointers to our data from the call element */
-  call_data *calld = elem->call_data;
-
-  /* initialize members */
-  memset(&calld->lr_data, 0, sizeof(load_reporting_data));
-}
+                           grpc_call_element_args *args) {}
 
 /* Destructor for call_data */
-static void destroy_call_elem(grpc_exec_ctx *exec_ctx,
-                              grpc_call_element *elem) {
-  /* grab pointers to our data from the call element */
-  /*call_data *calld = elem->call_data;*/
+static void destroy_call_elem(grpc_exec_ctx *exec_ctx, grpc_call_element *elem,
+                              const grpc_call_stats *stats) {
+  GPR_TIMER_BEGIN("load_reporting_filter", 0);
+  grpc_load_reporting_call(stats);
+  GPR_TIMER_END("load_reporting_filter", 0);
 }
 
 /* Constructor for channel_data */
 static void init_channel_elem(grpc_exec_ctx *exec_ctx,
                               grpc_channel_element *elem,
                               grpc_channel_element_args *args) {
-  /*channel_data *channeld = elem->channel_data;*/
   GPR_ASSERT(!args->is_last);
 }
 
@@ -86,7 +65,7 @@ static void destroy_channel_elem(grpc_exec_ctx *exec_ctx,
                                  grpc_channel_element *elem) {}
 
 const grpc_channel_filter grpc_load_reporting_filter = {
-    load_reporting_start_transport_stream_op,
+    grpc_call_next_op,
     grpc_channel_next_op,
     sizeof(call_data),
     init_call_elem,

src/core/lib/load_reporting/load_reporting.c

@@ -31,18 +31,42 @@
  *
  */
 
+#include <grpc/support/alloc.h>
+#include <grpc/support/sync.h>
+
 #include "src/core/lib/load_reporting/load_reporting.h"
 
-#include <grpc/grpc.h>
+typedef struct load_reporting {
+  gpr_mu mu;
+  load_reporting_fn fn;
+  void *data;
+} load_reporting;
+
+static load_reporting g_load_reporting;
 
-static load_reporting_fn g_load_reporting_fn;
+void grpc_load_reporting_init(load_reporting_fn fn, void *data) {
+  gpr_mu_init(&g_load_reporting.mu);
+  g_load_reporting.fn = fn;
+  g_load_reporting.data = data;
+}
 
-void grpc_load_reporting_init(load_reporting_fn fn) {
-  g_load_reporting_fn = fn;
+void grpc_load_reporting_destroy() {
+  gpr_free(g_load_reporting.data);
+  g_load_reporting.data = NULL;
+  gpr_mu_destroy(&g_load_reporting.mu);
 }
 
-void grpc_load_reporting_call(load_reporting_data *lr_data) {
-  if (g_load_reporting_fn != NULL) {
-    g_load_reporting_fn(lr_data);
+void grpc_load_reporting_call(const grpc_call_stats *stats) {
+  if (g_load_reporting.fn != NULL) {
+    gpr_mu_lock(&g_load_reporting.mu);
+    g_load_reporting.fn(g_load_reporting.data, stats);
+    gpr_mu_unlock(&g_load_reporting.mu);
   }
 }
+
+void *grpc_load_reporting_data() {
+  gpr_mu_lock(&g_load_reporting.mu);
+  void *data = g_load_reporting.data;
+  gpr_mu_unlock(&g_load_reporting.mu);
+  return data;
+}

src/core/lib/load_reporting/load_reporting.h

@@ -34,13 +34,32 @@
 #ifndef GRPC_CORE_LIB_LOAD_REPORTING_LOAD_REPORTING_H
 #define GRPC_CORE_LIB_LOAD_REPORTING_LOAD_REPORTING_H
 
-typedef struct load_reporting_data { void *data; } load_reporting_data;
+#include "src/core/lib/iomgr/closure.h"
+#include "src/core/lib/surface/call.h"
 
-typedef void (*load_reporting_fn)(load_reporting_data *lr_data);
+/** Custom function to be called by the load reporting filter.
+ *
+ * The \a data pointer is the same as the one passed to \a
+ * grpc_load_reporting_init. \a stats are the final per-call statistics gathered
+ * by the gRPC runtime. */
+typedef void (*load_reporting_fn)(void *data, const grpc_call_stats *stats);
+
+/** Register \a fn as the function to be invoked by the load reporting filter,
+ * passing \a data as its namesake argument. To be called only from a plugin
+ * init function. */
+void grpc_load_reporting_init(load_reporting_fn fn, void *data);
+
+/** Takes care of freeing the memory allocated for \a data (see \a
+ * grpc_load_reporting_init), if any. To be called only from a plugin destroy
+ * function. */
+void grpc_load_reporting_destroy();
 
-/** To be called only from a plugin init function */
-void grpc_load_reporting_init(load_reporting_fn fn);
+/** Invoke the function registered by \a grpc_load_reporting_init, passing it \a
+ * stats as one of the arguments (see \a load_reporting_fn). */
+void grpc_load_reporting_call(const grpc_call_stats *stats);
 
-void grpc_load_reporting_call(load_reporting_data *lr_data);
+/** Returns the custom load reporting data, as registered in \a
+ * grpc_load_reporting_init. */
+void *grpc_load_reporting_data();
 
 #endif /* GRPC_CORE_LIB_LOAD_REPORTING_LOAD_REPORTING_H */

test/core/load_reporting/load_reporting_registration.c

@@ -40,27 +40,35 @@
 #include "src/core/lib/surface/api_trace.h"
 #include "test/core/util/test_config.h"
 
-static void noop(load_reporting_data *lr_data) {
-  uint32_t *d = (uint32_t *)(lr_data->data);
-  *d = 0xdeadbeef;
+typedef struct { uint64_t total_bytes; } aggregated_bw_stats;
+
+static void sample_fn(void *lr_data, const grpc_call_stats *stats) {
+  aggregated_bw_stats *custom_stats = (aggregated_bw_stats *)lr_data;
+  custom_stats->total_bytes =
+      stats->transport_stream_stats.outgoing.data_bytes +
+      stats->transport_stream_stats.incoming.data_bytes;
 }
 
-static void lr_plugin_init(void) { grpc_load_reporting_init(noop); }
+static void lr_plugin_init(void) {
+  aggregated_bw_stats *data = gpr_malloc(sizeof(aggregated_bw_stats));
+  grpc_load_reporting_init(sample_fn, data);
+}
 
-static void lr_plugin_destroy(void) {}
+static void lr_plugin_destroy(void) { grpc_load_reporting_destroy(); }
 
 static void load_reporting_register() {
   grpc_register_plugin(lr_plugin_init, lr_plugin_destroy);
 }
 
 static void test_load_reporter_registration(void) {
-  load_reporting_data lr_data;
-  lr_data.data = gpr_malloc(sizeof(uint32_t));
-  grpc_load_reporting_call(&lr_data);
+  grpc_call_stats stats;
+  stats.transport_stream_stats.outgoing.data_bytes = 123;
+  stats.transport_stream_stats.incoming.data_bytes = 456;
 
-  GPR_ASSERT(*((uint32_t *)lr_data.data) == 0xdeadbeef);
+  grpc_load_reporting_call(&stats);
 
-  gpr_free(lr_data.data);
+  GPR_ASSERT(((aggregated_bw_stats *)grpc_load_reporting_data())->total_bytes ==
+             123 + 456);
 }
 
 int main(int argc, char **argv) {