1 files changed, 31 insertions, 64 deletions

diff --git a/src/core/ext/filters/client_channel/lb_policy.cc b/src/core/ext/filters/client_channel/lb_policy.cc
index 7a5a8dec34..cc4fe7ec62 100644
--- a/src/core/ext/filters/client_channel/lb_policy.cc
+++ b/src/core/ext/filters/client_channel/lb_policy.cc
@@ -19,8 +19,6 @@
 #include "src/core/ext/filters/client_channel/lb_policy.h"
 #include "src/core/lib/iomgr/combiner.h"
 
-#define WEAK_REF_BITS 16
-
 grpc_core::DebugOnlyTraceFlag grpc_trace_lb_policy_refcount(
     false, "lb_policy_refcount");
 
@@ -28,91 +26,60 @@ void grpc_lb_policy_init(grpc_lb_policy* policy,
                          const grpc_lb_policy_vtable* vtable,
                          grpc_combiner* combiner) {
   policy->vtable = vtable;
-  gpr_atm_no_barrier_store(&policy->ref_pair, 1 << WEAK_REF_BITS);
+  gpr_ref_init(&policy->refs, 1);
   policy->interested_parties = grpc_pollset_set_create();
   policy->combiner = GRPC_COMBINER_REF(combiner, "lb_policy");
 }
 
 #ifndef NDEBUG
-#define REF_FUNC_EXTRA_ARGS , const char *file, int line, const char *reason
-#define REF_MUTATE_EXTRA_ARGS REF_FUNC_EXTRA_ARGS, const char* purpose
-#define REF_FUNC_PASS_ARGS(new_reason) , file, line, new_reason
-#define REF_MUTATE_PASS_ARGS(purpose) , file, line, reason, purpose
+void grpc_lb_policy_ref(grpc_lb_policy* lb_policy, const char* file, int line,
+                        const char* reason) {
+  if (grpc_trace_lb_policy_refcount.enabled()) {
+    gpr_atm old_refs = gpr_atm_no_barrier_load(&lb_policy->refs.count);
+    gpr_log(file, line, GPR_LOG_SEVERITY_DEBUG,
+            "LB_POLICY:%p   ref %" PRIdPTR " -> %" PRIdPTR " %s", lb_policy,
+            old_refs, old_refs + 1, reason);
+  }
 #else
-#define REF_FUNC_EXTRA_ARGS
-#define REF_MUTATE_EXTRA_ARGS
-#define REF_FUNC_PASS_ARGS(new_reason)
-#define REF_MUTATE_PASS_ARGS(x)
+void grpc_lb_policy_ref(grpc_lb_policy* lb_policy) {
 #endif
+  gpr_ref(&lb_policy->refs);
+}
 
-static gpr_atm ref_mutate(grpc_lb_policy* c, gpr_atm delta,
-                          int barrier REF_MUTATE_EXTRA_ARGS) {
-  gpr_atm old_val = barrier ? gpr_atm_full_fetch_add(&c->ref_pair, delta)
-                            : gpr_atm_no_barrier_fetch_add(&c->ref_pair, delta);
 #ifndef NDEBUG
+void grpc_lb_policy_unref(grpc_lb_policy* lb_policy, const char* file, int line,
+                          const char* reason) {
   if (grpc_trace_lb_policy_refcount.enabled()) {
+    gpr_atm old_refs = gpr_atm_no_barrier_load(&lb_policy->refs.count);
     gpr_log(file, line, GPR_LOG_SEVERITY_DEBUG,
-            "LB_POLICY: %p %12s 0x%" PRIxPTR " -> 0x%" PRIxPTR " [%s]", c,
-            purpose, old_val, old_val + delta, reason);
+            "LB_POLICY:%p unref %" PRIdPTR " -> %" PRIdPTR " %s", lb_policy,
+            old_refs, old_refs - 1, reason);
   }
+#else
+void grpc_lb_policy_unref(grpc_lb_policy* lb_policy) {
 #endif
-  return old_val;
-}
-
-void grpc_lb_policy_ref(grpc_lb_policy* policy REF_FUNC_EXTRA_ARGS) {
-  ref_mutate(policy, 1 << WEAK_REF_BITS, 0 REF_MUTATE_PASS_ARGS("STRONG_REF"));
-}
-
-static void shutdown_locked(void* arg, grpc_error* error) {
-  grpc_lb_policy* policy = (grpc_lb_policy*)arg;
-  policy->vtable->shutdown_locked(policy);
-  GRPC_LB_POLICY_WEAK_UNREF(policy, "strong-unref");
-}
-
-void grpc_lb_policy_unref(grpc_lb_policy* policy REF_FUNC_EXTRA_ARGS) {
-  gpr_atm old_val =
-      ref_mutate(policy, (gpr_atm)1 - (gpr_atm)(1 << WEAK_REF_BITS),
-                 1 REF_MUTATE_PASS_ARGS("STRONG_UNREF"));
-  gpr_atm mask = ~(gpr_atm)((1 << WEAK_REF_BITS) - 1);
-  gpr_atm check = 1 << WEAK_REF_BITS;
-  if ((old_val & mask) == check) {
-    GRPC_CLOSURE_SCHED(
-        GRPC_CLOSURE_CREATE(shutdown_locked, policy,
-                            grpc_combiner_scheduler(policy->combiner)),
-        GRPC_ERROR_NONE);
-  } else {
-    grpc_lb_policy_weak_unref(policy REF_FUNC_PASS_ARGS("strong-unref"));
+  if (gpr_unref(&lb_policy->refs)) {
+    grpc_pollset_set_destroy(lb_policy->interested_parties);
+    grpc_combiner* combiner = lb_policy->combiner;
+    lb_policy->vtable->destroy(lb_policy);
+    GRPC_COMBINER_UNREF(combiner, "lb_policy");
   }
 }
 
-void grpc_lb_policy_weak_ref(grpc_lb_policy* policy REF_FUNC_EXTRA_ARGS) {
-  ref_mutate(policy, 1, 0 REF_MUTATE_PASS_ARGS("WEAK_REF"));
-}
-
-void grpc_lb_policy_weak_unref(grpc_lb_policy* policy REF_FUNC_EXTRA_ARGS) {
-  gpr_atm old_val =
-      ref_mutate(policy, -(gpr_atm)1, 1 REF_MUTATE_PASS_ARGS("WEAK_UNREF"));
-  if (old_val == 1) {
-    grpc_pollset_set_destroy(policy->interested_parties);
-    grpc_combiner* combiner = policy->combiner;
-    policy->vtable->destroy(policy);
-    GRPC_COMBINER_UNREF(combiner, "lb_policy");
-  }
+void grpc_lb_policy_shutdown_locked(grpc_lb_policy* policy,
+                                    grpc_lb_policy* new_policy) {
+  policy->vtable->shutdown_locked(policy, new_policy);
 }
 
 int grpc_lb_policy_pick_locked(grpc_lb_policy* policy,
-                               const grpc_lb_policy_pick_args* pick_args,
-                               grpc_connected_subchannel** target,
-                               grpc_call_context_element* context,
-                               void** user_data, grpc_closure* on_complete) {
-  return policy->vtable->pick_locked(policy, pick_args, target, context,
-                                     user_data, on_complete);
+                               grpc_lb_policy_pick_state* pick) {
+  return policy->vtable->pick_locked(policy, pick);
 }
 
 void grpc_lb_policy_cancel_pick_locked(grpc_lb_policy* policy,
-                                       grpc_connected_subchannel** target,
+                                       grpc_lb_policy_pick_state* pick,
                                        grpc_error* error) {
-  policy->vtable->cancel_pick_locked(policy, target, error);
+  policy->vtable->cancel_pick_locked(policy, pick, error);
 }
 
 void grpc_lb_policy_cancel_picks_locked(grpc_lb_policy* policy,