Change LB policies to get their input from channel args.

2 files changed, 57 insertions, 9 deletions

diff --git a/src/core/lib/channel/channel_args.c b/src/core/lib/channel/channel_args.c
index 2957d2c818..0270d6f224 100644
--- a/src/core/lib/channel/channel_args.c
+++ b/src/core/lib/channel/channel_args.c
@@ -66,22 +66,59 @@ static grpc_arg copy_arg(const grpc_arg *src) {
 grpc_channel_args *grpc_channel_args_copy_and_add(const grpc_channel_args *src,
                                                   const grpc_arg *to_add,
                                                   size_t num_to_add) {
+  return grpc_channel_args_copy_and_add_and_remove(src, NULL, 0, to_add,
+                                                   num_to_add);
+}
+
+grpc_channel_args *grpc_channel_args_copy_and_remove(
+    const grpc_channel_args *src, const char** to_remove,
+    size_t num_to_remove) {
+  return grpc_channel_args_copy_and_add_and_remove(src, to_remove,
+                                                   num_to_remove, NULL, 0);
+}
+
+static bool should_remove_arg(const grpc_arg* arg, const char** to_remove,
+                              size_t num_to_remove) {
+  for (size_t i = 0; i < num_to_remove; ++i) {
+    if (strcmp(arg->key, to_remove[i]) == 0) return true;
+  }
+  return false;
+}
+
+grpc_channel_args *grpc_channel_args_copy_and_add_and_remove(
+    const grpc_channel_args *src, const char** to_remove, size_t num_to_remove,
+    const grpc_arg *to_add, size_t num_to_add) {
+  // Figure out how many args we'll be copying.
+  size_t num_args_to_copy = 0;
+  if (src != NULL) {
+    for (size_t i = 0; i < src->num_args; ++i) {
+      if (!should_remove_arg(&src->args[i], to_remove, num_to_remove)) {
+        ++num_args_to_copy;
+      }
+    }
+  }
+  // Create result.
   grpc_channel_args *dst = gpr_malloc(sizeof(grpc_channel_args));
-  size_t i;
-  size_t src_num_args = (src == NULL) ? 0 : src->num_args;
-  if (!src && !to_add) {
-    dst->num_args = 0;
+  dst->num_args = num_args_to_copy + num_to_add;
+  if (dst->num_args == 0) {
     dst->args = NULL;
     return dst;
   }
-  dst->num_args = src_num_args + num_to_add;
   dst->args = gpr_malloc(sizeof(grpc_arg) * dst->num_args);
-  for (i = 0; i < src_num_args; i++) {
-    dst->args[i] = copy_arg(&src->args[i]);
+  // Copy args from src that are not being removed.
+  size_t dst_idx = 0;
+  if (src != NULL) {
+    for (size_t i = 0; i < src->num_args; ++i) {
+      if (!should_remove_arg(&src->args[i], to_remove, num_to_remove)) {
+        dst->args[dst_idx++] = copy_arg(&src->args[i]);
+      }
+    }
   }
-  for (i = 0; i < num_to_add; i++) {
-    dst->args[i + src_num_args] = copy_arg(&to_add[i]);
+  // Add args from to_add.
+  for (size_t i = 0; i < num_to_add; ++i) {
+    dst->args[dst_idx++] = copy_arg(&to_add[i]);
   }
+  GPR_ASSERT(dst_idx == dst->num_args);
   return dst;
 }
 
diff --git a/src/core/lib/channel/channel_args.h b/src/core/lib/channel/channel_args.h
index a80340c0fa..1508d23748 100644
--- a/src/core/lib/channel/channel_args.h
+++ b/src/core/lib/channel/channel_args.h
@@ -51,6 +51,17 @@ grpc_channel_args *grpc_channel_args_copy_and_add(const grpc_channel_args *src,
                                                   const grpc_arg *to_add,
                                                   size_t num_to_add);
 
+/** Copies the arguments in \a src except for those whose keys are in
+    \a to_remove. */
+grpc_channel_args *grpc_channel_args_copy_and_remove(
+    const grpc_channel_args *src, const char** to_remove, size_t num_to_remove);
+
+/** Copies the arguments from \a src except for those whose keys are in
+    \a to_remove and appends the arguments in \a to_add. */
+grpc_channel_args *grpc_channel_args_copy_and_add_and_remove(
+    const grpc_channel_args *src, const char** to_remove, size_t num_to_remove,
+    const grpc_arg *to_add, size_t num_to_add);
+
 /** Concatenate args from \a a and \a b into a new instance */
 grpc_channel_args *grpc_channel_args_merge(const grpc_channel_args *a,
                                            const grpc_channel_args *b);