Change LB policies to get their input from channel args.

3 files changed, 90 insertions, 32 deletions

diff --git a/src/core/ext/lb_policy/grpclb/grpclb.c b/src/core/ext/lb_policy/grpclb/grpclb.c
index 1467508d6d..fdc0bec996 100644
--- a/src/core/ext/lb_policy/grpclb/grpclb.c
+++ b/src/core/ext/lb_policy/grpclb/grpclb.c
@@ -470,9 +470,17 @@ static grpc_lb_policy *create_rr_locked(
   args.server_name = glb_policy->server_name;
   args.client_channel_factory = glb_policy->cc_factory;
   args.addresses = process_serverlist(serverlist);
-  args.args = glb_policy->args;
+
+  // Replace the LB addresses in the channel args that we pass down to
+  // the subchannel.
+  static const char* keys_to_remove[] = {GRPC_ARG_LB_ADDRESSES};
+  const grpc_arg arg = grpc_lb_addresses_create_channel_arg(args.addresses);
+  args.args = grpc_channel_args_copy_and_add_and_remove(
+      glb_policy->args, keys_to_remove, GPR_ARRAY_SIZE(keys_to_remove), &arg,
+      1);
 
   grpc_lb_policy *rr = grpc_lb_policy_create(exec_ctx, "round_robin", &args);
+  grpc_channel_args_destroy(args.args);
 
   if (glb_policy->addresses != NULL) {
     /* dispose of the previous version */
@@ -574,6 +582,13 @@ static void glb_rr_connectivity_changed(grpc_exec_ctx *exec_ctx, void *arg,
 static grpc_lb_policy *glb_create(grpc_exec_ctx *exec_ctx,
                                   grpc_lb_policy_factory *factory,
                                   grpc_lb_policy_args *args) {
+  /* Get server name. */
+  const grpc_arg* arg =
+      grpc_channel_args_find(args->args, GRPC_ARG_SERVER_NAME);
+  const char* server_name =
+      arg != NULL && arg->type == GRPC_ARG_STRING
+          ? arg->value.string : NULL;
+
   /* Count the number of gRPC-LB addresses. There must be at least one.
    * TODO(roth): For now, we ignore non-balancer addresses, but in the
    * future, we may change the behavior such that we fall back to using
@@ -581,22 +596,25 @@ static grpc_lb_policy *glb_create(grpc_exec_ctx *exec_ctx,
    * time, this should be changed to allow a list with no balancer addresses,
    * since the resolver might fail to return a balancer address even when
    * this is the right LB policy to use. */
+  arg = grpc_channel_args_find(args->args, GRPC_ARG_LB_ADDRESSES);
+  GPR_ASSERT(arg != NULL && arg->type == GRPC_ARG_POINTER);
+  grpc_lb_addresses* addresses = arg->value.pointer.p;
   size_t num_grpclb_addrs = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; ++i) {
-    if (args->addresses->addresses[i].is_balancer) ++num_grpclb_addrs;
+  for (size_t i = 0; i < addresses->num_addresses; ++i) {
+    if (addresses->addresses[i].is_balancer) ++num_grpclb_addrs;
   }
   if (num_grpclb_addrs == 0) return NULL;
 
   glb_lb_policy *glb_policy = gpr_malloc(sizeof(*glb_policy));
   memset(glb_policy, 0, sizeof(*glb_policy));
 
-  /* All input addresses in args->addresses come from a resolver that claims
+  /* All input addresses in addresses come from a resolver that claims
    * they are LB services. It's the resolver's responsibility to make sure
    * this
    * policy is only instantiated and used in that case.
    *
    * Create a client channel over them to communicate with a LB service */
-  glb_policy->server_name = gpr_strdup(args->server_name);
+  glb_policy->server_name = gpr_strdup(server_name);
   glb_policy->cc_factory = args->client_channel_factory;
   glb_policy->args = grpc_channel_args_copy(args->args);
   GPR_ASSERT(glb_policy->cc_factory != NULL);
@@ -606,20 +624,20 @@ static grpc_lb_policy *glb_create(grpc_exec_ctx *exec_ctx,
    * TODO(dgq): support mixed ip version */
   char **addr_strs = gpr_malloc(sizeof(char *) * num_grpclb_addrs);
   size_t addr_index = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; i++) {
-    if (args->addresses->addresses[i].user_data != NULL) {
+  for (size_t i = 0; i < addresses->num_addresses; i++) {
+    if (addresses->addresses[i].user_data != NULL) {
       gpr_log(GPR_ERROR,
               "This LB policy doesn't support user data. It will be ignored");
     }
-    if (args->addresses->addresses[i].is_balancer) {
+    if (addresses->addresses[i].is_balancer) {
       if (addr_index == 0) {
         addr_strs[addr_index++] = grpc_sockaddr_to_uri(
-            (const struct sockaddr *)&args->addresses->addresses[i]
+            (const struct sockaddr *)&addresses->addresses[i]
                 .address.addr);
       } else {
         GPR_ASSERT(grpc_sockaddr_to_string(
                        &addr_strs[addr_index++],
-                       (const struct sockaddr *)&args->addresses->addresses[i]
+                       (const struct sockaddr *)&addresses->addresses[i]
                            .address.addr,
                        true) > 0);
       }
@@ -629,10 +647,29 @@ static grpc_lb_policy *glb_create(grpc_exec_ctx *exec_ctx,
   char *target_uri_str = gpr_strjoin_sep((const char **)addr_strs,
                                          num_grpclb_addrs, ",", &uri_path_len);
 
-  /* will pick using pick_first */
+  /* Create a channel to talk to the LBs.
+   *
+   * We strip out the channel arg for the LB policy name, since we want
+   * to use the default (pick_first) in this case.
+   *
+   * We also strip out the channel arg for the resolved addresses, since
+   * that will be generated by the name resolver used in the LB channel.
+   * Note that the LB channel will use the sockaddr resolver, so this
+   * won't actually generate a query to DNS (or some other name service).
+   * However, the addresses returned by the sockaddr resolver will have
+   * is_balancer=false, whereas our own addresses have is_balancer=true.
+   * We need the LB channel to return addresses with is_balancer=false
+   * so that it does not wind up recursively using the grpclb LB policy,
+   * as per the special case logic in client_channel.c.
+   */
+  static const char* keys_to_remove[] = {
+      GRPC_ARG_LB_POLICY_NAME, GRPC_ARG_LB_ADDRESSES};
+  grpc_channel_args *new_args = grpc_channel_args_copy_and_remove(
+      args->args, keys_to_remove, GPR_ARRAY_SIZE(keys_to_remove));
   glb_policy->lb_channel = grpc_client_channel_factory_create_channel(
       exec_ctx, glb_policy->cc_factory, target_uri_str,
-      GRPC_CLIENT_CHANNEL_TYPE_LOAD_BALANCING, args->args);
+      GRPC_CLIENT_CHANNEL_TYPE_LOAD_BALANCING, new_args);
+  grpc_channel_args_destroy(new_args);
 
   gpr_free(target_uri_str);
   for (size_t i = 0; i < num_grpclb_addrs; i++) {
diff --git a/src/core/ext/lb_policy/pick_first/pick_first.c b/src/core/ext/lb_policy/pick_first/pick_first.c
index 3683079cf4..d95f0310b6 100644
--- a/src/core/ext/lb_policy/pick_first/pick_first.c
+++ b/src/core/ext/lb_policy/pick_first/pick_first.c
@@ -34,7 +34,9 @@
 #include <string.h>
 
 #include <grpc/support/alloc.h>
+
 #include "src/core/ext/client_config/lb_policy_registry.h"
+#include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/transport/connectivity_state.h"
 
 typedef struct pending_pick {
@@ -432,14 +434,23 @@ static void pick_first_factory_unref(grpc_lb_policy_factory *factory) {}
 static grpc_lb_policy *create_pick_first(grpc_exec_ctx *exec_ctx,
                                          grpc_lb_policy_factory *factory,
                                          grpc_lb_policy_args *args) {
-  GPR_ASSERT(args->addresses != NULL);
   GPR_ASSERT(args->client_channel_factory != NULL);
 
+  /* Get server name. */
+  const grpc_arg* arg =
+      grpc_channel_args_find(args->args, GRPC_ARG_SERVER_NAME);
+  const char* server_name =
+      arg != NULL && arg->type == GRPC_ARG_STRING
+          ? arg->value.string : NULL;
+
   /* Find the number of backend addresses. We ignore balancer
    * addresses, since we don't know how to handle them. */
+  arg = grpc_channel_args_find(args->args, GRPC_ARG_LB_ADDRESSES);
+  GPR_ASSERT(arg != NULL && arg->type == GRPC_ARG_POINTER);
+  grpc_lb_addresses* addresses = arg->value.pointer.p;
   size_t num_addrs = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; i++) {
-    if (!args->addresses->addresses[i].is_balancer) ++num_addrs;
+  for (size_t i = 0; i < addresses->num_addresses; i++) {
+    if (!addresses->addresses[i].is_balancer) ++num_addrs;
   }
   if (num_addrs == 0) return NULL;
 
@@ -450,22 +461,22 @@ static grpc_lb_policy *create_pick_first(grpc_exec_ctx *exec_ctx,
   memset(p->subchannels, 0, sizeof(*p->subchannels) * num_addrs);
   grpc_subchannel_args sc_args;
   size_t subchannel_idx = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; i++) {
+  for (size_t i = 0; i < addresses->num_addresses; i++) {
     /* Skip balancer addresses, since we only know how to handle backends. */
-    if (args->addresses->addresses[i].is_balancer) continue;
+    if (addresses->addresses[i].is_balancer) continue;
 
-    if (args->addresses->addresses[i].user_data != NULL) {
+    if (addresses->addresses[i].user_data != NULL) {
       gpr_log(GPR_ERROR,
               "This LB policy doesn't support user data. It will be ignored");
     }
 
     memset(&sc_args, 0, sizeof(grpc_subchannel_args));
     /* server_name will be copied as part of the subchannel creation. This makes
-     * the copying of args->server_name (a borrowed pointer) OK. */
-    sc_args.server_name = args->server_name;
+     * the copying of server_name (a borrowed pointer) OK. */
+    sc_args.server_name = server_name;
     sc_args.addr =
-        (struct sockaddr *)(&args->addresses->addresses[i].address.addr);
-    sc_args.addr_len = args->addresses->addresses[i].address.len;
+        (struct sockaddr *)(&addresses->addresses[i].address.addr);
+    sc_args.addr_len = addresses->addresses[i].address.len;
     sc_args.args = args->args;
 
     grpc_subchannel *subchannel = grpc_client_channel_factory_create_subchannel(
diff --git a/src/core/ext/lb_policy/round_robin/round_robin.c b/src/core/ext/lb_policy/round_robin/round_robin.c
index 00a18974df..2deb0af99c 100644
--- a/src/core/ext/lb_policy/round_robin/round_robin.c
+++ b/src/core/ext/lb_policy/round_robin/round_robin.c
@@ -64,6 +64,7 @@
 #include <grpc/support/alloc.h>
 
 #include "src/core/ext/client_config/lb_policy_registry.h"
+#include "src/core/lib/channel/channel_args.h"
 #include "src/core/lib/debug/trace.h"
 #include "src/core/lib/transport/connectivity_state.h"
 #include "src/core/lib/transport/static_metadata.h"
@@ -598,14 +599,23 @@ static void round_robin_factory_unref(grpc_lb_policy_factory *factory) {}
 static grpc_lb_policy *round_robin_create(grpc_exec_ctx *exec_ctx,
                                           grpc_lb_policy_factory *factory,
                                           grpc_lb_policy_args *args) {
-  GPR_ASSERT(args->addresses != NULL);
   GPR_ASSERT(args->client_channel_factory != NULL);
 
+  /* Get server name. */
+  const grpc_arg* arg =
+      grpc_channel_args_find(args->args, GRPC_ARG_SERVER_NAME);
+  const char* server_name =
+      arg != NULL && arg->type == GRPC_ARG_STRING
+          ? arg->value.string : NULL;
+
   /* Find the number of backend addresses. We ignore balancer
    * addresses, since we don't know how to handle them. */
+  arg = grpc_channel_args_find(args->args, GRPC_ARG_LB_ADDRESSES);
+  GPR_ASSERT(arg != NULL && arg->type == GRPC_ARG_POINTER);
+  grpc_lb_addresses* addresses = arg->value.pointer.p;
   size_t num_addrs = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; i++) {
-    if (!args->addresses->addresses[i].is_balancer) ++num_addrs;
+  for (size_t i = 0; i < addresses->num_addresses; i++) {
+    if (!addresses->addresses[i].is_balancer) ++num_addrs;
   }
   if (num_addrs == 0) return NULL;
 
@@ -618,17 +628,17 @@ static grpc_lb_policy *round_robin_create(grpc_exec_ctx *exec_ctx,
 
   grpc_subchannel_args sc_args;
   size_t subchannel_idx = 0;
-  for (size_t i = 0; i < args->addresses->num_addresses; i++) {
+  for (size_t i = 0; i < addresses->num_addresses; i++) {
     /* Skip balancer addresses, since we only know how to handle backends. */
-    if (args->addresses->addresses[i].is_balancer) continue;
+    if (addresses->addresses[i].is_balancer) continue;
 
     memset(&sc_args, 0, sizeof(grpc_subchannel_args));
     /* server_name will be copied as part of the subchannel creation. This makes
-     * the copying of args->server_name (a borrowed pointer) OK. */
-    sc_args.server_name = args->server_name;
+     * the copying of server_name (a borrowed pointer) OK. */
+    sc_args.server_name = server_name;
     sc_args.addr =
-        (struct sockaddr *)(&args->addresses->addresses[i].address.addr);
-    sc_args.addr_len = args->addresses->addresses[i].address.len;
+        (struct sockaddr *)(&addresses->addresses[i].address.addr);
+    sc_args.addr_len = addresses->addresses[i].address.len;
     sc_args.args = args->args;
 
     grpc_subchannel *subchannel = grpc_client_channel_factory_create_subchannel(
@@ -641,7 +651,7 @@ static grpc_lb_policy *round_robin_create(grpc_exec_ctx *exec_ctx,
       sd->policy = p;
       sd->index = subchannel_idx;
       sd->subchannel = subchannel;
-      sd->user_data = args->addresses->addresses[i].user_data;
+      sd->user_data = addresses->addresses[i].user_data;
       ++subchannel_idx;
       grpc_closure_init(&sd->connectivity_changed_closure,
                         rr_connectivity_changed, sd);