Improvements to the grpc_channel_args_compression_algorithm_set_state api

3 files changed, 28 insertions, 19 deletions

diff --git a/src/core/channel/channel_args.c b/src/core/channel/channel_args.c
index dc66de7dd6..54ee75af28 100644
--- a/src/core/channel/channel_args.c
+++ b/src/core/channel/channel_args.c
@@ -167,13 +167,13 @@ static int find_compression_algorithm_states_bitset(
 }
 
 grpc_channel_args *grpc_channel_args_compression_algorithm_set_state(
-    grpc_channel_args *a,
+    grpc_channel_args **a,
     grpc_compression_algorithm algorithm,
     int state) {
   int *states_arg;
-  grpc_channel_args *result = a;
+  grpc_channel_args *result = *a;
   const int states_arg_found =
-      find_compression_algorithm_states_bitset(a, &states_arg);
+      find_compression_algorithm_states_bitset(*a, &states_arg);
 
   if (states_arg_found) {
     if (state != 0) {
@@ -193,8 +193,9 @@ grpc_channel_args *grpc_channel_args_compression_algorithm_set_state(
     } else {
       GPR_BITCLEAR(&tmp.value.integer, algorithm);
     }
-    result = grpc_channel_args_copy_and_add(a, &tmp, 1);
-    grpc_channel_args_destroy(a);
+    result = grpc_channel_args_copy_and_add(*a, &tmp, 1);
+    grpc_channel_args_destroy(*a);
+    *a = result;
   }
   return result;
 }
diff --git a/src/core/channel/channel_args.h b/src/core/channel/channel_args.h
index e557f9a9d9..06a6012dee 100644
--- a/src/core/channel/channel_args.h
+++ b/src/core/channel/channel_args.h
@@ -70,9 +70,12 @@ grpc_channel_args *grpc_channel_args_set_compression_algorithm(
 /** Sets the support for the given compression algorithm. By default, all
  * compression algorithms are enabled. It's an error to disable an algorithm set
  * by grpc_channel_args_set_compression_algorithm.
- * */
+ *
+ * Returns an instance will the updated algorithm states. The \a a pointer is
+ * modified to point to the returned instance (which may be different from the
+ * input value of \a a). */
 grpc_channel_args *grpc_channel_args_compression_algorithm_set_state(
-    grpc_channel_args *a,
+    grpc_channel_args **a,
     grpc_compression_algorithm algorithm,
     int enabled);
 
diff --git a/test/core/channel/channel_args_test.c b/test/core/channel/channel_args_test.c
index 227cc1f415..87f006acde 100644
--- a/test/core/channel/channel_args_test.c
+++ b/test/core/channel/channel_args_test.c
@@ -45,7 +45,7 @@ static void test_create(void) {
   grpc_arg arg_string;
   grpc_arg to_add[2];
   grpc_channel_args *ch_args;
-  
+
   arg_int.key = "int_arg";
   arg_int.type = GRPC_ARG_INTEGER;
   arg_int.value.integer = 123;
@@ -57,7 +57,7 @@ static void test_create(void) {
   to_add[0] = arg_int;
   to_add[1] = arg_string;
   ch_args = grpc_channel_args_copy_and_add(NULL, to_add, 2);
-  
+
   GPR_ASSERT(ch_args->num_args == 2);
   GPR_ASSERT(strcmp(ch_args->args[0].key, arg_int.key) == 0);
   GPR_ASSERT(ch_args->args[0].type == arg_int.type);
@@ -84,7 +84,7 @@ static void test_set_compression_algorithm(void) {
 }
 
 static void test_compression_algorithm_states(void) {
-  grpc_channel_args *ch_args;
+  grpc_channel_args *ch_args, *ch_args_wo_gzip, *ch_args_wo_gzip_deflate;
   int states_bitset;
   size_t i;
 
@@ -97,12 +97,15 @@ static void test_compression_algorithm_states(void) {
   }
 
   /* disable gzip and deflate */
-  ch_args = grpc_channel_args_compression_algorithm_set_state(
-      ch_args, GRPC_COMPRESS_GZIP, 0);
-  ch_args = grpc_channel_args_compression_algorithm_set_state(
-      ch_args, GRPC_COMPRESS_DEFLATE, 0);
-
-  states_bitset = grpc_channel_args_compression_algorithm_get_states(ch_args);
+  ch_args_wo_gzip = grpc_channel_args_compression_algorithm_set_state(
+      &ch_args, GRPC_COMPRESS_GZIP, 0);
+  GPR_ASSERT(ch_args == ch_args_wo_gzip);
+  ch_args_wo_gzip_deflate = grpc_channel_args_compression_algorithm_set_state(
+      &ch_args_wo_gzip, GRPC_COMPRESS_DEFLATE, 0);
+  GPR_ASSERT(ch_args_wo_gzip == ch_args_wo_gzip_deflate);
+
+  states_bitset = grpc_channel_args_compression_algorithm_get_states(
+      ch_args_wo_gzip_deflate);
   for (i = 0; i < GRPC_COMPRESS_ALGORITHMS_COUNT; i++) {
     if (i == GRPC_COMPRESS_GZIP || i == GRPC_COMPRESS_DEFLATE) {
       GPR_ASSERT(GPR_BITGET(states_bitset, i) == 0);
@@ -112,10 +115,12 @@ static void test_compression_algorithm_states(void) {
   }
 
   /* re-enabled gzip only */
-  ch_args = grpc_channel_args_compression_algorithm_set_state(
-      ch_args, GRPC_COMPRESS_GZIP, 1);
+  ch_args_wo_gzip = grpc_channel_args_compression_algorithm_set_state(
+      &ch_args_wo_gzip_deflate, GRPC_COMPRESS_GZIP, 1);
+  GPR_ASSERT(ch_args_wo_gzip == ch_args_wo_gzip_deflate);
 
-  states_bitset = grpc_channel_args_compression_algorithm_get_states(ch_args);
+  states_bitset =
+      grpc_channel_args_compression_algorithm_get_states(ch_args_wo_gzip);
   for (i = 0; i < GRPC_COMPRESS_ALGORITHMS_COUNT; i++) {
     if (i == GRPC_COMPRESS_DEFLATE) {
       GPR_ASSERT(GPR_BITGET(states_bitset, i) == 0);