Merge pull request #11865 from yang-g/avl_interface

Adding user_data to avl interface so that an exec_ctx can be passed in.

include/grpc/support/avl.h

@@ -31,18 +31,23 @@ typedef struct gpr_avl_node {
   long height;
 } gpr_avl_node;
 
+/** vtable for the AVL tree
+ * The optional user_data is propagated from the top level gpr_avl_XXX API.
+ * From the same API call, multiple vtable functions may be called multiple
+ * times.
+ */
 typedef struct gpr_avl_vtable {
   /** destroy a key */
-  void (*destroy_key)(void *key);
+  void (*destroy_key)(void *key, void *user_data);
   /** copy a key, returning new value */
-  void *(*copy_key)(void *key);
+  void *(*copy_key)(void *key, void *user_data);
   /** compare key1, key2; return <0 if key1 < key2,
       >0 if key1 > key2, 0 if key1 == key2 */
-  long (*compare_keys)(void *key1, void *key2);
+  long (*compare_keys)(void *key1, void *key2, void *user_data);
   /** destroy a value */
-  void (*destroy_value)(void *value);
+  void (*destroy_value)(void *value, void *user_data);
   /** copy a value */
-  void *(*copy_value)(void *value);
+  void *(*copy_value)(void *value, void *user_data);
 } gpr_avl_vtable;
 
 /** "pointer" to an AVL tree - this is a reference
@@ -53,29 +58,36 @@ typedef struct gpr_avl {
   gpr_avl_node *root;
 } gpr_avl;
 
-/** create an immutable AVL tree */
+/** Create an immutable AVL tree. */
 GPRAPI gpr_avl gpr_avl_create(const gpr_avl_vtable *vtable);
-/** add a reference to an existing tree - returns
-    the tree as a convenience */
-GPRAPI gpr_avl gpr_avl_ref(gpr_avl avl);
-/** remove a reference to a tree - destroying it if there
-    are no references left */
-GPRAPI void gpr_avl_unref(gpr_avl avl);
-/** return a new tree with (key, value) added to avl.
+/** Add a reference to an existing tree - returns
+    the tree as a convenience. The optional user_data will be passed to vtable
+    functions. */
+GPRAPI gpr_avl gpr_avl_ref(gpr_avl avl, void *user_data);
+/** Remove a reference to a tree - destroying it if there
+    are no references left. The optional user_data will be passed to vtable
+    functions. */
+GPRAPI void gpr_avl_unref(gpr_avl avl, void *user_data);
+/** Return a new tree with (key, value) added to avl.
     implicitly unrefs avl to allow easy chaining.
     if key exists in avl, the new tree's key entry updated
-    (i.e. a duplicate is not created) */
-GPRAPI gpr_avl gpr_avl_add(gpr_avl avl, void *key, void *value);
-/** return a new tree with key deleted
-    implicitly unrefs avl to allow easy chaining. */
-GPRAPI gpr_avl gpr_avl_remove(gpr_avl avl, void *key);
-/** lookup key, and return the associated value.
-    does not mutate avl.
-    returns NULL if key is not found. */
-GPRAPI void *gpr_avl_get(gpr_avl avl, void *key);
+    (i.e. a duplicate is not created). The optional user_data will be passed to
+    vtable functions. */
+GPRAPI gpr_avl gpr_avl_add(gpr_avl avl, void *key, void *value,
+                           void *user_data);
+/** Return a new tree with key deleted
+    implicitly unrefs avl to allow easy chaining. The optional user_data will be
+    passed to vtable functions. */
+GPRAPI gpr_avl gpr_avl_remove(gpr_avl avl, void *key, void *user_data);
+/** Lookup key, and return the associated value.
+    Does not mutate avl.
+    Returns NULL if key is not found. The optional user_data will be passed to
+    vtable functions.*/
+GPRAPI void *gpr_avl_get(gpr_avl avl, void *key, void *user_data);
 /** Return 1 if avl contains key, 0 otherwise; if it has the key, sets *value to
-    its value*/
-GPRAPI int gpr_avl_maybe_get(gpr_avl avl, void *key, void **value);
+    its value. THe optional user_data will be passed to vtable functions. */
+GPRAPI int gpr_avl_maybe_get(gpr_avl avl, void *key, void **value,
+                             void *user_data);
 /** Return 1 if avl is empty, 0 otherwise */
 GPRAPI int gpr_avl_is_empty(gpr_avl avl);
 

src/core/ext/filters/client_channel/retry_throttle.c

@@ -130,24 +130,28 @@ static grpc_server_retry_throttle_data* grpc_server_retry_throttle_data_create(
 // avl vtable for string -> server_retry_throttle_data map
 //
 
-static void* copy_server_name(void* key) { return gpr_strdup(key); }
+static void* copy_server_name(void* key, void* unused) {
+  return gpr_strdup(key);
+}
 
-static long compare_server_name(void* key1, void* key2) {
+static long compare_server_name(void* key1, void* key2, void* unused) {
   return strcmp(key1, key2);
 }
 
-static void destroy_server_retry_throttle_data(void* value) {
+static void destroy_server_retry_throttle_data(void* value, void* unused) {
   grpc_server_retry_throttle_data* throttle_data = value;
   grpc_server_retry_throttle_data_unref(throttle_data);
 }
 
-static void* copy_server_retry_throttle_data(void* value) {
+static void* copy_server_retry_throttle_data(void* value, void* unused) {
   grpc_server_retry_throttle_data* throttle_data = value;
   return grpc_server_retry_throttle_data_ref(throttle_data);
 }
 
+static void destroy_server_name(void* key, void* unused) { gpr_free(key); }
+
 static const gpr_avl_vtable avl_vtable = {
-    gpr_free /* destroy_key */, copy_server_name, compare_server_name,
+    destroy_server_name, copy_server_name, compare_server_name,
     destroy_server_retry_throttle_data, copy_server_retry_throttle_data};
 
 //
@@ -164,19 +168,19 @@ void grpc_retry_throttle_map_init() {
 
 void grpc_retry_throttle_map_shutdown() {
   gpr_mu_destroy(&g_mu);
-  gpr_avl_unref(g_avl);
+  gpr_avl_unref(g_avl, NULL);
 }
 
 grpc_server_retry_throttle_data* grpc_retry_throttle_map_get_data_for_server(
     const char* server_name, int max_milli_tokens, int milli_token_ratio) {
   gpr_mu_lock(&g_mu);
   grpc_server_retry_throttle_data* throttle_data =
-      gpr_avl_get(g_avl, (char*)server_name);
+      gpr_avl_get(g_avl, (char*)server_name, NULL);
   if (throttle_data == NULL) {
     // Entry not found.  Create a new one.
     throttle_data = grpc_server_retry_throttle_data_create(
         max_milli_tokens, milli_token_ratio, NULL);
-    g_avl = gpr_avl_add(g_avl, (char*)server_name, throttle_data);
+    g_avl = gpr_avl_add(g_avl, (char*)server_name, throttle_data, NULL);
   } else {
     if (throttle_data->max_milli_tokens != max_milli_tokens ||
         throttle_data->milli_token_ratio != milli_token_ratio) {
@@ -184,7 +188,7 @@ grpc_server_retry_throttle_data* grpc_retry_throttle_map_get_data_for_server(
       // the original one.
       throttle_data = grpc_server_retry_throttle_data_create(
           max_milli_tokens, milli_token_ratio, throttle_data);
-      g_avl = gpr_avl_add(g_avl, (char*)server_name, throttle_data);
+      g_avl = gpr_avl_add(g_avl, (char*)server_name, throttle_data, NULL);
     } else {
       // Entry found.  Increase refcount.
       grpc_server_retry_throttle_data_ref(throttle_data);

src/core/ext/filters/client_channel/subchannel_index.c

@@ -38,26 +38,8 @@ struct grpc_subchannel_key {
   grpc_subchannel_args args;
 };
 
-GPR_TLS_DECL(subchannel_index_exec_ctx);
-
 static bool g_force_creation = false;
 
-static void enter_ctx(grpc_exec_ctx *exec_ctx) {
-  GPR_ASSERT(gpr_tls_get(&subchannel_index_exec_ctx) == 0);
-  gpr_tls_set(&subchannel_index_exec_ctx, (intptr_t)exec_ctx);
-}
-
-static void leave_ctx(grpc_exec_ctx *exec_ctx) {
-  GPR_ASSERT(gpr_tls_get(&subchannel_index_exec_ctx) == (intptr_t)exec_ctx);
-  gpr_tls_set(&subchannel_index_exec_ctx, 0);
-}
-
-static grpc_exec_ctx *current_ctx() {
-  grpc_exec_ctx *c = (grpc_exec_ctx *)gpr_tls_get(&subchannel_index_exec_ctx);
-  GPR_ASSERT(c != NULL);
-  return c;
-}
-
 static grpc_subchannel_key *create_key(
     const grpc_subchannel_args *args,
     grpc_channel_args *(*copy_channel_args)(const grpc_channel_args *args)) {
@@ -104,21 +86,25 @@ void grpc_subchannel_key_destroy(grpc_exec_ctx *exec_ctx,
   gpr_free(k);
 }
 
-static void sck_avl_destroy(void *p) {
-  grpc_subchannel_key_destroy(current_ctx(), p);
+static void sck_avl_destroy(void *p, void *user_data) {
+  grpc_exec_ctx *exec_ctx = (grpc_exec_ctx *)user_data;
+  grpc_subchannel_key_destroy(exec_ctx, p);
 }
 
-static void *sck_avl_copy(void *p) { return subchannel_key_copy(p); }
+static void *sck_avl_copy(void *p, void *unused) {
+  return subchannel_key_copy(p);
+}
 
-static long sck_avl_compare(void *a, void *b) {
+static long sck_avl_compare(void *a, void *b, void *unused) {
   return grpc_subchannel_key_compare(a, b);
 }
 
-static void scv_avl_destroy(void *p) {
-  GRPC_SUBCHANNEL_WEAK_UNREF(current_ctx(), p, "subchannel_index");
+static void scv_avl_destroy(void *p, void *user_data) {
+  grpc_exec_ctx *exec_ctx = (grpc_exec_ctx *)user_data;
+  GRPC_SUBCHANNEL_WEAK_UNREF(exec_ctx, p, "subchannel_index");
 }
 
-static void *scv_avl_copy(void *p) {
+static void *scv_avl_copy(void *p, void *unused) {
   GRPC_SUBCHANNEL_WEAK_REF(p, "subchannel_index");
   return p;
 }
@@ -133,38 +119,33 @@ static const gpr_avl_vtable subchannel_avl_vtable = {
 void grpc_subchannel_index_init(void) {
   g_subchannel_index = gpr_avl_create(&subchannel_avl_vtable);
   gpr_mu_init(&g_mu);
-  gpr_tls_init(&subchannel_index_exec_ctx);
 }
 
 void grpc_subchannel_index_shutdown(void) {
+  grpc_exec_ctx exec_ctx = GRPC_EXEC_CTX_INIT;
   gpr_mu_destroy(&g_mu);
-  gpr_avl_unref(g_subchannel_index);
-  gpr_tls_destroy(&subchannel_index_exec_ctx);
+  gpr_avl_unref(g_subchannel_index, &exec_ctx);
+  grpc_exec_ctx_finish(&exec_ctx);
 }
 
 grpc_subchannel *grpc_subchannel_index_find(grpc_exec_ctx *exec_ctx,
                                             grpc_subchannel_key *key) {
-  enter_ctx(exec_ctx);
-
   // Lock, and take a reference to the subchannel index.
   // We don't need to do the search under a lock as avl's are immutable.
   gpr_mu_lock(&g_mu);
-  gpr_avl index = gpr_avl_ref(g_subchannel_index);
+  gpr_avl index = gpr_avl_ref(g_subchannel_index, exec_ctx);
   gpr_mu_unlock(&g_mu);
 
-  grpc_subchannel *c =
-      GRPC_SUBCHANNEL_REF_FROM_WEAK_REF(gpr_avl_get(index, key), "index_find");
-  gpr_avl_unref(index);
+  grpc_subchannel *c = GRPC_SUBCHANNEL_REF_FROM_WEAK_REF(
+      gpr_avl_get(index, key, exec_ctx), "index_find");
+  gpr_avl_unref(index, exec_ctx);
 
-  leave_ctx(exec_ctx);
   return c;
 }
 
 grpc_subchannel *grpc_subchannel_index_register(grpc_exec_ctx *exec_ctx,
                                                 grpc_subchannel_key *key,
                                                 grpc_subchannel *constructed) {
-  enter_ctx(exec_ctx);
-
   grpc_subchannel *c = NULL;
   bool need_to_unref_constructed;
 
@@ -174,11 +155,11 @@ grpc_subchannel *grpc_subchannel_index_register(grpc_exec_ctx *exec_ctx,
     // Compare and swap loop:
     // - take a reference to the current index
     gpr_mu_lock(&g_mu);
-    gpr_avl index = gpr_avl_ref(g_subchannel_index);
+    gpr_avl index = gpr_avl_ref(g_subchannel_index, exec_ctx);
     gpr_mu_unlock(&g_mu);
 
     // - Check to see if a subchannel already exists
-    c = gpr_avl_get(index, key);
+    c = gpr_avl_get(index, key, exec_ctx);
     if (c != NULL) {
       c = GRPC_SUBCHANNEL_REF_FROM_WEAK_REF(c, "index_register");
     }
@@ -187,9 +168,9 @@ grpc_subchannel *grpc_subchannel_index_register(grpc_exec_ctx *exec_ctx,
       need_to_unref_constructed = true;
     } else {
       // no -> update the avl and compare/swap
-      gpr_avl updated =
-          gpr_avl_add(gpr_avl_ref(index), subchannel_key_copy(key),
-                      GRPC_SUBCHANNEL_WEAK_REF(constructed, "index_register"));
+      gpr_avl updated = gpr_avl_add(
+          gpr_avl_ref(index, exec_ctx), subchannel_key_copy(key),
+          GRPC_SUBCHANNEL_WEAK_REF(constructed, "index_register"), exec_ctx);
 
       // it may happen (but it's expected to be unlikely)
       // that some other thread has changed the index:
@@ -201,13 +182,11 @@ grpc_subchannel *grpc_subchannel_index_register(grpc_exec_ctx *exec_ctx,
       }
       gpr_mu_unlock(&g_mu);
 
-      gpr_avl_unref(updated);
+      gpr_avl_unref(updated, exec_ctx);
     }
-    gpr_avl_unref(index);
+    gpr_avl_unref(index, exec_ctx);
   }
 
-  leave_ctx(exec_ctx);
-
   if (need_to_unref_constructed) {
     GRPC_SUBCHANNEL_UNREF(exec_ctx, constructed, "index_register");
   }
@@ -218,27 +197,26 @@ grpc_subchannel *grpc_subchannel_index_register(grpc_exec_ctx *exec_ctx,
 void grpc_subchannel_index_unregister(grpc_exec_ctx *exec_ctx,
                                       grpc_subchannel_key *key,
                                       grpc_subchannel *constructed) {
-  enter_ctx(exec_ctx);
-
   bool done = false;
   while (!done) {
     // Compare and swap loop:
     // - take a reference to the current index
     gpr_mu_lock(&g_mu);
-    gpr_avl index = gpr_avl_ref(g_subchannel_index);
+    gpr_avl index = gpr_avl_ref(g_subchannel_index, exec_ctx);
     gpr_mu_unlock(&g_mu);
 
     // Check to see if this key still refers to the previously
     // registered subchannel
-    grpc_subchannel *c = gpr_avl_get(index, key);
+    grpc_subchannel *c = gpr_avl_get(index, key, exec_ctx);
     if (c != constructed) {
-      gpr_avl_unref(index);
+      gpr_avl_unref(index, exec_ctx);
       break;
     }
 
     // compare and swap the update (some other thread may have
     // mutated the index behind us)
-    gpr_avl updated = gpr_avl_remove(gpr_avl_ref(index), key);
+    gpr_avl updated =
+        gpr_avl_remove(gpr_avl_ref(index, exec_ctx), key, exec_ctx);
 
     gpr_mu_lock(&g_mu);
     if (index.root == g_subchannel_index.root) {
@@ -247,11 +225,9 @@ void grpc_subchannel_index_unregister(grpc_exec_ctx *exec_ctx,
     }
     gpr_mu_unlock(&g_mu);
 
-    gpr_avl_unref(updated);
-    gpr_avl_unref(index);
+    gpr_avl_unref(updated, exec_ctx);
+    gpr_avl_unref(index, exec_ctx);
   }
-
-  leave_ctx(exec_ctx);
 }
 
 void grpc_subchannel_index_test_only_set_force_creation(bool force_creation) {

src/core/lib/support/avl.c

@@ -39,15 +39,16 @@ static gpr_avl_node *ref_node(gpr_avl_node *node) {
   return node;
 }
 
-static void unref_node(const gpr_avl_vtable *vtable, gpr_avl_node *node) {
+static void unref_node(const gpr_avl_vtable *vtable, gpr_avl_node *node,
+                       void *user_data) {
   if (node == NULL) {
     return;
   }
   if (gpr_unref(&node->refs)) {
-    vtable->destroy_key(node->key);
-    vtable->destroy_value(node->value);
-    unref_node(vtable, node->left);
-    unref_node(vtable, node->right);
+    vtable->destroy_key(node->key, user_data);
+    vtable->destroy_value(node->value, user_data);
+    unref_node(vtable, node->left, user_data);
+    unref_node(vtable, node->right, user_data);
     gpr_free(node);
   }
 }
@@ -87,30 +88,30 @@ gpr_avl_node *new_node(void *key, void *value, gpr_avl_node *left,
 }
 
 static gpr_avl_node *get(const gpr_avl_vtable *vtable, gpr_avl_node *node,
-                         void *key) {
+                         void *key, void *user_data) {
   long cmp;
 
   if (node == NULL) {
     return NULL;
   }
 
-  cmp = vtable->compare_keys(node->key, key);
+  cmp = vtable->compare_keys(node->key, key, user_data);
   if (cmp == 0) {
     return node;
   } else if (cmp > 0) {
-    return get(vtable, node->left, key);
+    return get(vtable, node->left, key, user_data);
   } else {
-    return get(vtable, node->right, key);
+    return get(vtable, node->right, key, user_data);
   }
 }
 
-void *gpr_avl_get(gpr_avl avl, void *key) {
-  gpr_avl_node *node = get(avl.vtable, avl.root, key);
+void *gpr_avl_get(gpr_avl avl, void *key, void *user_data) {
+  gpr_avl_node *node = get(avl.vtable, avl.root, key, user_data);
   return node ? node->value : NULL;
 }
 
-int gpr_avl_maybe_get(gpr_avl avl, void *key, void **value) {
-  gpr_avl_node *node = get(avl.vtable, avl.root, key);
+int gpr_avl_maybe_get(gpr_avl avl, void *key, void **value, void *user_data) {
+  gpr_avl_node *node = get(avl.vtable, avl.root, key, user_data);
   if (node != NULL) {
     *value = node->value;
     return 1;
@@ -120,70 +121,75 @@ int gpr_avl_maybe_get(gpr_avl avl, void *key, void **value) {
 
 static gpr_avl_node *rotate_left(const gpr_avl_vtable *vtable, void *key,
                                  void *value, gpr_avl_node *left,
-                                 gpr_avl_node *right) {
-  gpr_avl_node *n =
-      new_node(vtable->copy_key(right->key), vtable->copy_value(right->value),
-               new_node(key, value, left, ref_node(right->left)),
-               ref_node(right->right));
-  unref_node(vtable, right);
+                                 gpr_avl_node *right, void *user_data) {
+  gpr_avl_node *n = new_node(vtable->copy_key(right->key, user_data),
+                             vtable->copy_value(right->value, user_data),
+                             new_node(key, value, left, ref_node(right->left)),
+                             ref_node(right->right));
+  unref_node(vtable, right, user_data);
   return n;
 }
 
 static gpr_avl_node *rotate_right(const gpr_avl_vtable *vtable, void *key,
                                   void *value, gpr_avl_node *left,
-                                  gpr_avl_node *right) {
-  gpr_avl_node *n = new_node(
-      vtable->copy_key(left->key), vtable->copy_value(left->value),
-      ref_node(left->left), new_node(key, value, ref_node(left->right), right));
-  unref_node(vtable, left);
+                                  gpr_avl_node *right, void *user_data) {
+  gpr_avl_node *n =
+      new_node(vtable->copy_key(left->key, user_data),
+               vtable->copy_value(left->value, user_data), ref_node(left->left),
+               new_node(key, value, ref_node(left->right), right));
+  unref_node(vtable, left, user_data);
   return n;
 }
 
 static gpr_avl_node *rotate_left_right(const gpr_avl_vtable *vtable, void *key,
                                        void *value, gpr_avl_node *left,
-                                       gpr_avl_node *right) {
+                                       gpr_avl_node *right, void *user_data) {
   /* rotate_right(..., rotate_left(left), right) */
-  gpr_avl_node *n = new_node(
-      vtable->copy_key(left->right->key),
-      vtable->copy_value(left->right->value),
-      new_node(vtable->copy_key(left->key), vtable->copy_value(left->value),
-               ref_node(left->left), ref_node(left->right->left)),
-      new_node(key, value, ref_node(left->right->right), right));
-  unref_node(vtable, left);
+  gpr_avl_node *n =
+      new_node(vtable->copy_key(left->right->key, user_data),
+               vtable->copy_value(left->right->value, user_data),
+               new_node(vtable->copy_key(left->key, user_data),
+                        vtable->copy_value(left->value, user_data),
+                        ref_node(left->left), ref_node(left->right->left)),
+               new_node(key, value, ref_node(left->right->right), right));
+  unref_node(vtable, left, user_data);
   return n;
 }
 
 static gpr_avl_node *rotate_right_left(const gpr_avl_vtable *vtable, void *key,
                                        void *value, gpr_avl_node *left,
-                                       gpr_avl_node *right) {
+                                       gpr_avl_node *right, void *user_data) {
   /* rotate_left(..., left, rotate_right(right)) */
-  gpr_avl_node *n = new_node(
-      vtable->copy_key(right->left->key),
-      vtable->copy_value(right->left->value),
-      new_node(key, value, left, ref_node(right->left->left)),
-      new_node(vtable->copy_key(right->key), vtable->copy_value(right->value),
-               ref_node(right->left->right), ref_node(right->right)));
-  unref_node(vtable, right);
+  gpr_avl_node *n =
+      new_node(vtable->copy_key(right->left->key, user_data),
+               vtable->copy_value(right->left->value, user_data),
+               new_node(key, value, left, ref_node(right->left->left)),
+               new_node(vtable->copy_key(right->key, user_data),
+                        vtable->copy_value(right->value, user_data),
+                        ref_node(right->left->right), ref_node(right->right)));
+  unref_node(vtable, right, user_data);
   return n;
 }
 
 static gpr_avl_node *rebalance(const gpr_avl_vtable *vtable, void *key,
                                void *value, gpr_avl_node *left,
-                               gpr_avl_node *right) {
+                               gpr_avl_node *right, void *user_data) {
   switch (node_height(left) - node_height(right)) {
     case 2:
       if (node_height(left->left) - node_height(left->right) == -1) {
         return assert_invariants(
-            rotate_left_right(vtable, key, value, left, right));
+            rotate_left_right(vtable, key, value, left, right, user_data));
       } else {
-        return assert_invariants(rotate_right(vtable, key, value, left, right));
+        return assert_invariants(
+            rotate_right(vtable, key, value, left, right, user_data));
       }
     case -2:
       if (node_height(right->left) - node_height(right->right) == 1) {
         return assert_invariants(
-            rotate_right_left(vtable, key, value, left, right));
+            rotate_right_left(vtable, key, value, left, right, user_data));
       } else {
-        return assert_invariants(rotate_left(vtable, key, value, left, right));
+        return assert_invariants(
+            rotate_left(vtable, key, value, left, right, user_data));
       }
     default:
       return assert_invariants(new_node(key, value, left, right));
@@ -191,30 +197,32 @@ static gpr_avl_node *rebalance(const gpr_avl_vtable *vtable, void *key,
 }
 
 static gpr_avl_node *add_key(const gpr_avl_vtable *vtable, gpr_avl_node *node,
-                             void *key, void *value) {
+                             void *key, void *value, void *user_data) {
   long cmp;
   if (node == NULL) {
     return new_node(key, value, NULL, NULL);
   }
-  cmp = vtable->compare_keys(node->key, key);
+  cmp = vtable->compare_keys(node->key, key, user_data);
   if (cmp == 0) {
     return new_node(key, value, ref_node(node->left), ref_node(node->right));
   } else if (cmp > 0) {
-    return rebalance(
-        vtable, vtable->copy_key(node->key), vtable->copy_value(node->value),
-        add_key(vtable, node->left, key, value), ref_node(node->right));
+    return rebalance(vtable, vtable->copy_key(node->key, user_data),
+                     vtable->copy_value(node->value, user_data),
+                     add_key(vtable, node->left, key, value, user_data),
+                     ref_node(node->right), user_data);
   } else {
-    return rebalance(vtable, vtable->copy_key(node->key),
-                     vtable->copy_value(node->value), ref_node(node->left),
-                     add_key(vtable, node->right, key, value));
+    return rebalance(
+        vtable, vtable->copy_key(node->key, user_data),
+        vtable->copy_value(node->value, user_data), ref_node(node->left),
+        add_key(vtable, node->right, key, value, user_data), user_data);
   }
 }
 
-gpr_avl gpr_avl_add(gpr_avl avl, void *key, void *value) {
+gpr_avl gpr_avl_add(gpr_avl avl, void *key, void *value, void *user_data) {
   gpr_avl_node *old_root = avl.root;
-  avl.root = add_key(avl.vtable, avl.root, key, value);
+  avl.root = add_key(avl.vtable, avl.root, key, value, user_data);
   assert_invariants(avl.root);
-  unref_node(avl.vtable, old_root);
+  unref_node(avl.vtable, old_root, user_data);
   return avl;
 }
 
@@ -233,12 +241,13 @@ static gpr_avl_node *in_order_tail(gpr_avl_node *node) {
 }
 
 static gpr_avl_node *remove_key(const gpr_avl_vtable *vtable,
-                                gpr_avl_node *node, void *key) {
+                                gpr_avl_node *node, void *key,
+                                void *user_data) {
   long cmp;
   if (node == NULL) {
     return NULL;
   }
-  cmp = vtable->compare_keys(node->key, key);
+  cmp = vtable->compare_keys(node->key, key, user_data);
   if (cmp == 0) {
     if (node->left == NULL) {
       return ref_node(node->right);
@@ -246,39 +255,45 @@ static gpr_avl_node *remove_key(const gpr_avl_vtable *vtable,
       return ref_node(node->left);
     } else if (node->left->height < node->right->height) {
       gpr_avl_node *h = in_order_head(node->right);
-      return rebalance(vtable, vtable->copy_key(h->key),
-                       vtable->copy_value(h->value), ref_node(node->left),
-                       remove_key(vtable, node->right, h->key));
+      return rebalance(
+          vtable, vtable->copy_key(h->key, user_data),
+          vtable->copy_value(h->value, user_data), ref_node(node->left),
+          remove_key(vtable, node->right, h->key, user_data), user_data);
     } else {
       gpr_avl_node *h = in_order_tail(node->left);
-      return rebalance(
-          vtable, vtable->copy_key(h->key), vtable->copy_value(h->value),
-          remove_key(vtable, node->left, h->key), ref_node(node->right));
+      return rebalance(vtable, vtable->copy_key(h->key, user_data),
+                       vtable->copy_value(h->value, user_data),
+                       remove_key(vtable, node->left, h->key, user_data),
+                       ref_node(node->right), user_data);
     }
   } else if (cmp > 0) {
-    return rebalance(
-        vtable, vtable->copy_key(node->key), vtable->copy_value(node->value),
-        remove_key(vtable, node->left, key), ref_node(node->right));
+    return rebalance(vtable, vtable->copy_key(node->key, user_data),
+                     vtable->copy_value(node->value, user_data),
+                     remove_key(vtable, node->left, key, user_data),
+                     ref_node(node->right), user_data);
   } else {
-    return rebalance(vtable, vtable->copy_key(node->key),
-                     vtable->copy_value(node->value), ref_node(node->left),
-                     remove_key(vtable, node->right, key));
+    return rebalance(
+        vtable, vtable->copy_key(node->key, user_data),
+        vtable->copy_value(node->value, user_data), ref_node(node->left),
+        remove_key(vtable, node->right, key, user_data), user_data);
   }
 }
 
-gpr_avl gpr_avl_remove(gpr_avl avl, void *key) {
+gpr_avl gpr_avl_remove(gpr_avl avl, void *key, void *user_data) {
   gpr_avl_node *old_root = avl.root;
-  avl.root = remove_key(avl.vtable, avl.root, key);
+  avl.root = remove_key(avl.vtable, avl.root, key, user_data);
   assert_invariants(avl.root);
-  unref_node(avl.vtable, old_root);
+  unref_node(avl.vtable, old_root, user_data);
   return avl;
 }
 
-gpr_avl gpr_avl_ref(gpr_avl avl) {
+gpr_avl gpr_avl_ref(gpr_avl avl, void *user_data) {
   ref_node(avl.root);
   return avl;
 }
 
-void gpr_avl_unref(gpr_avl avl) { unref_node(avl.vtable, avl.root); }
+void gpr_avl_unref(gpr_avl avl, void *user_data) {
+  unref_node(avl.vtable, avl.root, user_data);
+}
 
 int gpr_avl_is_empty(gpr_avl avl) { return avl.root == NULL; }

src/ruby/ext/grpc/rb_grpc_imports.generated.h

@@ -644,22 +644,22 @@ extern gpr_get_allocation_functions_type gpr_get_allocation_functions_import;
 typedef gpr_avl(*gpr_avl_create_type)(const gpr_avl_vtable *vtable);
 extern gpr_avl_create_type gpr_avl_create_import;
 #define gpr_avl_create gpr_avl_create_import
-typedef gpr_avl(*gpr_avl_ref_type)(gpr_avl avl);
+typedef gpr_avl(*gpr_avl_ref_type)(gpr_avl avl, void *user_data);
 extern gpr_avl_ref_type gpr_avl_ref_import;
 #define gpr_avl_ref gpr_avl_ref_import
-typedef void(*gpr_avl_unref_type)(gpr_avl avl);
+typedef void(*gpr_avl_unref_type)(gpr_avl avl, void *user_data);
 extern gpr_avl_unref_type gpr_avl_unref_import;
 #define gpr_avl_unref gpr_avl_unref_import
-typedef gpr_avl(*gpr_avl_add_type)(gpr_avl avl, void *key, void *value);
+typedef gpr_avl(*gpr_avl_add_type)(gpr_avl avl, void *key, void *value, void *user_data);
 extern gpr_avl_add_type gpr_avl_add_import;
 #define gpr_avl_add gpr_avl_add_import
-typedef gpr_avl(*gpr_avl_remove_type)(gpr_avl avl, void *key);
+typedef gpr_avl(*gpr_avl_remove_type)(gpr_avl avl, void *key, void *user_data);
 extern gpr_avl_remove_type gpr_avl_remove_import;
 #define gpr_avl_remove gpr_avl_remove_import
-typedef void *(*gpr_avl_get_type)(gpr_avl avl, void *key);
+typedef void *(*gpr_avl_get_type)(gpr_avl avl, void *key, void *user_data);
 extern gpr_avl_get_type gpr_avl_get_import;
 #define gpr_avl_get gpr_avl_get_import
-typedef int(*gpr_avl_maybe_get_type)(gpr_avl avl, void *key, void **value);
+typedef int(*gpr_avl_maybe_get_type)(gpr_avl avl, void *key, void **value, void *user_data);
 extern gpr_avl_maybe_get_type gpr_avl_maybe_get_import;
 #define gpr_avl_maybe_get gpr_avl_maybe_get_import
 typedef int(*gpr_avl_is_empty_type)(gpr_avl avl);