2 files changed, 49 insertions, 39 deletions

diff --git a/src/core/ext/lb_policy/grpclb/load_balancer_api.c b/src/core/ext/lb_policy/grpclb/load_balancer_api.c
index 6c77612e0d..8b3f2ac3ef 100644
--- a/src/core/ext/lb_policy/grpclb/load_balancer_api.c
+++ b/src/core/ext/lb_policy/grpclb/load_balancer_api.c
@@ -38,9 +38,15 @@
 #include <grpc/support/alloc.h>
 
 typedef struct decode_serverlist_arg {
-  int first_pass;
-  int i;
+  /* The first pass counts the number of servers in the server list. The second
+   * one allocates and decodes. */
+  bool first_pass;
+  /* The decoding callback is invoked once per server in serverlist. Remember
+   * which index of the serverlist are we currently decoding */
+  size_t decoding_idx;
+  /* Populated after the first pass. Number of server in the input serverlist */
   size_t num_servers;
+  /* The decoded serverlist */
   grpc_grpclb_server **servers;
 } decode_serverlist_arg;
 
@@ -48,24 +54,24 @@ typedef struct decode_serverlist_arg {
 static bool decode_serverlist(pb_istream_t *stream, const pb_field_t *field,
                               void **arg) {
   decode_serverlist_arg *dec_arg = *arg;
-  if (dec_arg->first_pass != 0) { /* first pass */
+  if (dec_arg->first_pass) { /* count how many server do we have */
     grpc_grpclb_server server;
     if (!pb_decode(stream, grpc_lb_v1_Server_fields, &server)) {
       return false;
     }
     dec_arg->num_servers++;
-  } else { /* second pass */
+  } else { /* second pass. Actually decode. */
     grpc_grpclb_server *server = gpr_malloc(sizeof(grpc_grpclb_server));
     memset(server, 0, sizeof(grpc_grpclb_server));
     GPR_ASSERT(dec_arg->num_servers > 0);
-    if (dec_arg->i == 0) { /* first iteration of second pass */
+    if (dec_arg->decoding_idx == 0) { /* first iteration of second pass */
       dec_arg->servers =
           gpr_malloc(sizeof(grpc_grpclb_server *) * dec_arg->num_servers);
     }
     if (!pb_decode(stream, grpc_lb_v1_Server_fields, server)) {
       return false;
     }
-    dec_arg->servers[dec_arg->i++] = server;
+    dec_arg->servers[dec_arg->decoding_idx++] = server;
   }
 
   return true;
@@ -103,19 +109,22 @@ void grpc_grpclb_request_destroy(grpc_grpclb_request *request) {
   gpr_free(request);
 }
 
-grpc_grpclb_response *grpc_grpclb_response_parse(gpr_slice encoded_response) {
-  bool status;
+grpc_grpclb_initial_response *grpc_grpclb_initial_response_parse(
+    gpr_slice encoded_response) {
   pb_istream_t stream =
       pb_istream_from_buffer(GPR_SLICE_START_PTR(encoded_response),
                              GPR_SLICE_LENGTH(encoded_response));
-  grpc_grpclb_response *res = gpr_malloc(sizeof(grpc_grpclb_response));
-  memset(res, 0, sizeof(*res));
-  status = pb_decode(&stream, grpc_lb_v1_LoadBalanceResponse_fields, res);
-  if (!status) {
-    grpc_grpclb_response_destroy(res);
+  grpc_grpclb_response res;
+  memset(&res, 0, sizeof(grpc_grpclb_response));
+  if (!pb_decode(&stream, grpc_lb_v1_LoadBalanceResponse_fields, &res)) {
     return NULL;
   }
-  return res;
+  grpc_grpclb_initial_response *initial_res =
+      gpr_malloc(sizeof(grpc_grpclb_initial_response));
+  memcpy(initial_res, &res.initial_response,
+         sizeof(grpc_grpclb_initial_response));
+
+  return initial_res;
 }
 
 grpc_grpclb_serverlist *grpc_grpclb_response_parse_serverlist(
@@ -126,24 +135,22 @@ grpc_grpclb_serverlist *grpc_grpclb_response_parse_serverlist(
       pb_istream_from_buffer(GPR_SLICE_START_PTR(encoded_response),
                              GPR_SLICE_LENGTH(encoded_response));
   pb_istream_t stream_at_start = stream;
-  grpc_grpclb_response *res = gpr_malloc(sizeof(grpc_grpclb_response));
-  memset(res, 0, sizeof(*res));
+  grpc_grpclb_response res;
+  memset(&res, 0, sizeof(grpc_grpclb_response));
   memset(&arg, 0, sizeof(decode_serverlist_arg));
 
-  res->server_list.servers.funcs.decode = decode_serverlist;
-  res->server_list.servers.arg = &arg;
-  arg.first_pass = 1;
-  status = pb_decode(&stream, grpc_lb_v1_LoadBalanceResponse_fields, res);
+  res.server_list.servers.funcs.decode = decode_serverlist;
+  res.server_list.servers.arg = &arg;
+  arg.first_pass = true;
+  status = pb_decode(&stream, grpc_lb_v1_LoadBalanceResponse_fields, &res);
   if (!status) {
-    grpc_grpclb_response_destroy(res);
     return NULL;
   }
 
-  arg.first_pass = 0;
+  arg.first_pass = false;
   status =
-      pb_decode(&stream_at_start, grpc_lb_v1_LoadBalanceResponse_fields, res);
+      pb_decode(&stream_at_start, grpc_lb_v1_LoadBalanceResponse_fields, &res);
   if (!status) {
-    grpc_grpclb_response_destroy(res);
     return NULL;
   }
 
@@ -151,10 +158,9 @@ grpc_grpclb_serverlist *grpc_grpclb_response_parse_serverlist(
   memset(sl, 0, sizeof(*sl));
   sl->num_servers = arg.num_servers;
   sl->servers = arg.servers;
-  if (res->server_list.has_expiration_interval) {
-    sl->expiration_interval = res->server_list.expiration_interval;
+  if (res.server_list.has_expiration_interval) {
+    sl->expiration_interval = res.server_list.expiration_interval;
   }
-  grpc_grpclb_response_destroy(res);
   return sl;
 }
 
@@ -244,6 +250,7 @@ compare_nanos:
   return rhs->has_nanos ? 1 : -1;
 }
 
-void grpc_grpclb_response_destroy(grpc_grpclb_response *response) {
+void grpc_grpclb_initial_response_destroy(
+    grpc_grpclb_initial_response *response) {
   gpr_free(response);
 }
diff --git a/src/core/ext/lb_policy/grpclb/load_balancer_api.h b/src/core/ext/lb_policy/grpclb/load_balancer_api.h
index cc5c4e9343..5f1dff707e 100644
--- a/src/core/ext/lb_policy/grpclb/load_balancer_api.h
+++ b/src/core/ext/lb_policy/grpclb/load_balancer_api.h
@@ -47,6 +47,7 @@ extern "C" {
 
 typedef grpc_lb_v1_LoadBalanceRequest grpc_grpclb_request;
 typedef grpc_lb_v1_LoadBalanceResponse grpc_grpclb_response;
+typedef grpc_lb_v1_InitialLoadBalanceResponse grpc_grpclb_initial_response;
 typedef grpc_lb_v1_Server grpc_grpclb_server;
 typedef grpc_lb_v1_Duration grpc_grpclb_duration;
 typedef struct grpc_grpclb_serverlist {
@@ -65,34 +66,36 @@ gpr_slice grpc_grpclb_request_encode(const grpc_grpclb_request *request);
 void grpc_grpclb_request_destroy(grpc_grpclb_request *request);
 
 /** Parse (ie, decode) the bytes in \a encoded_response as a \a
- * grpc_grpclb_response */
-grpc_grpclb_response *grpc_grpclb_response_parse(gpr_slice encoded_response);
+ * grpc_grpclb_initial_response */
+grpc_grpclb_initial_response *grpc_grpclb_initial_response_parse(
+    gpr_slice encoded_response);
+
+/** Parse the list of servers from an encoded \a grpc_grpclb_response */
+grpc_grpclb_serverlist *grpc_grpclb_response_parse_serverlist(
+    gpr_slice encoded_response);
 
 /** Return a copy of \a sl. The caller is responsible for calling \a
  * grpc_grpclb_destroy_serverlist on the returned copy. */
 grpc_grpclb_serverlist *grpc_grpclb_serverlist_copy(
     const grpc_grpclb_serverlist *sl);
 
-/** Destroy \a serverlist */
-void grpc_grpclb_destroy_serverlist(grpc_grpclb_serverlist *serverlist);
-
-/** Parse the list of servers from an encoded \a grpc_grpclb_response */
-grpc_grpclb_serverlist *grpc_grpclb_response_parse_serverlist(
-    gpr_slice encoded_response);
-
 bool grpc_grpclb_serverlist_equals(const grpc_grpclb_serverlist *lhs,
                                    const grpc_grpclb_serverlist *rhs);
 
 bool grpc_grpclb_server_equals(const grpc_grpclb_server *lhs,
                                const grpc_grpclb_server *rhs);
 
+/** Destroy \a serverlist */
+void grpc_grpclb_destroy_serverlist(grpc_grpclb_serverlist *serverlist);
+
 /** Compare \a lhs against \a rhs and return 0 if \a lhs and \a rhs are equal,
  * < 0 if \a lhs represents a duration shorter than \a rhs and > 0 otherwise */
 int grpc_grpclb_duration_compare(const grpc_grpclb_duration *lhs,
                                  const grpc_grpclb_duration *rhs);
 
-/** Destroy \a response */
-void grpc_grpclb_response_destroy(grpc_grpclb_response *response);
+/** Destroy \a initial_response */
+void grpc_grpclb_initial_response_destroy(
+    grpc_grpclb_initial_response *response);
 
 #ifdef __cplusplus
 }