Revert "Simplify LB policy refcounting."

1 files changed, 51 insertions, 39 deletions

diff --git a/src/core/ext/filters/client_channel/lb_policy.h b/src/core/ext/filters/client_channel/lb_policy.h
index 1176a05b78..3572c97ed1 100644
--- a/src/core/ext/filters/client_channel/lb_policy.h
+++ b/src/core/ext/filters/client_channel/lb_policy.h
@@ -33,7 +33,7 @@ extern grpc_core::DebugOnlyTraceFlag grpc_trace_lb_policy_refcount;
 
 struct grpc_lb_policy {
   const grpc_lb_policy_vtable* vtable;
-  gpr_refcount refs;
+  gpr_atm ref_pair;
   /* owned pointer to interested parties in load balancing decisions */
   grpc_pollset_set* interested_parties;
   /* combiner under which lb_policy actions take place */
@@ -42,42 +42,32 @@ struct grpc_lb_policy {
   grpc_closure* request_reresolution;
 };
 
-/// State used for an LB pick.
-typedef struct grpc_lb_policy_pick_state {
-  /// Initial metadata associated with the picking call.
+/** Extra arguments for an LB pick */
+typedef struct grpc_lb_policy_pick_args {
+  /** Initial metadata associated with the picking call. */
   grpc_metadata_batch* initial_metadata;
-  /// Bitmask used for selective cancelling. See \a
-  /// grpc_lb_policy_cancel_picks() and \a GRPC_INITIAL_METADATA_* in
-  /// grpc_types.h.
+  /** Bitmask used for selective cancelling. See \a
+   * grpc_lb_policy_cancel_picks() and \a GRPC_INITIAL_METADATA_* in
+   * grpc_types.h */
   uint32_t initial_metadata_flags;
-  /// Storage for LB token in \a initial_metadata, or NULL if not used.
-  grpc_linked_mdelem lb_token_mdelem_storage;
-  /// Closure to run when pick is complete, if not completed synchronously.
-  grpc_closure* on_complete;
-  /// Will be set to the selected subchannel, or NULL on failure or when
-  /// the LB policy decides to drop the call.
-  grpc_connected_subchannel* connected_subchannel;
-  /// Will be populated with context to pass to the subchannel call, if needed.
-  grpc_call_context_element subchannel_call_context[GRPC_CONTEXT_COUNT];
-  /// Upon success, \a *user_data will be set to whatever opaque information
-  /// may need to be propagated from the LB policy, or NULL if not needed.
-  void** user_data;
-  /// Next pointer.  For internal use by LB policy.
-  struct grpc_lb_policy_pick_state* next;
-} grpc_lb_policy_pick_state;
+  /** Storage for LB token in \a initial_metadata, or NULL if not used */
+  grpc_linked_mdelem* lb_token_mdelem_storage;
+} grpc_lb_policy_pick_args;
 
 struct grpc_lb_policy_vtable {
   void (*destroy)(grpc_lb_policy* policy);
-
-  /// \see grpc_lb_policy_shutdown_locked().
-  void (*shutdown_locked)(grpc_lb_policy* policy, grpc_lb_policy* new_policy);
+  void (*shutdown_locked)(grpc_lb_policy* policy);
 
   /** \see grpc_lb_policy_pick */
-  int (*pick_locked)(grpc_lb_policy* policy, grpc_lb_policy_pick_state* pick);
+  int (*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);
 
   /** \see grpc_lb_policy_cancel_pick */
   void (*cancel_pick_locked)(grpc_lb_policy* policy,
-                             grpc_lb_policy_pick_state* pick,
+                             grpc_connected_subchannel** target,
                              grpc_error* error);
 
   /** \see grpc_lb_policy_cancel_picks */
@@ -113,19 +103,37 @@ struct grpc_lb_policy_vtable {
 };
 
 #ifndef NDEBUG
+
+/* Strong references: the policy will shutdown when they reach zero */
 #define GRPC_LB_POLICY_REF(p, r) \
   grpc_lb_policy_ref((p), __FILE__, __LINE__, (r))
 #define GRPC_LB_POLICY_UNREF(p, r) \
   grpc_lb_policy_unref((p), __FILE__, __LINE__, (r))
+
+/* Weak references: they don't prevent the shutdown of the LB policy. When no
+ * strong references are left but there are still weak ones, shutdown is called.
+ * Once the weak reference also reaches zero, the LB policy is destroyed. */
+#define GRPC_LB_POLICY_WEAK_REF(p, r) \
+  grpc_lb_policy_weak_ref((p), __FILE__, __LINE__, (r))
+#define GRPC_LB_POLICY_WEAK_UNREF(p, r) \
+  grpc_lb_policy_weak_unref((p), __FILE__, __LINE__, (r))
 void grpc_lb_policy_ref(grpc_lb_policy* policy, const char* file, int line,
                         const char* reason);
 void grpc_lb_policy_unref(grpc_lb_policy* policy, const char* file, int line,
                           const char* reason);
-#else  // !NDEBUG
+void grpc_lb_policy_weak_ref(grpc_lb_policy* policy, const char* file, int line,
+                             const char* reason);
+void grpc_lb_policy_weak_unref(grpc_lb_policy* policy, const char* file,
+                               int line, const char* reason);
+#else
 #define GRPC_LB_POLICY_REF(p, r) grpc_lb_policy_ref((p))
 #define GRPC_LB_POLICY_UNREF(p, r) grpc_lb_policy_unref((p))
+#define GRPC_LB_POLICY_WEAK_REF(p, r) grpc_lb_policy_weak_ref((p))
+#define GRPC_LB_POLICY_WEAK_UNREF(p, r) grpc_lb_policy_weak_unref((p))
 void grpc_lb_policy_ref(grpc_lb_policy* policy);
 void grpc_lb_policy_unref(grpc_lb_policy* policy);
+void grpc_lb_policy_weak_ref(grpc_lb_policy* policy);
+void grpc_lb_policy_weak_unref(grpc_lb_policy* policy);
 #endif
 
 /** called by concrete implementations to initialize the base struct */
@@ -133,24 +141,28 @@ void grpc_lb_policy_init(grpc_lb_policy* policy,
                          const grpc_lb_policy_vtable* vtable,
                          grpc_combiner* combiner);
 
-/// Shuts down \a policy.
-/// If \a new_policy is non-null, any pending picks will be restarted
-/// on that policy; otherwise, they will be failed.
-void grpc_lb_policy_shutdown_locked(grpc_lb_policy* policy,
-                                    grpc_lb_policy* new_policy);
+/** Finds an appropriate subchannel for a call, based on \a pick_args.
+
+    \a target will be set to the selected subchannel, or NULL on failure
+    or when the LB policy decides to drop the call.
 
-/** Finds an appropriate subchannel for a call, based on data in \a pick.
-    \a pick must remain alive until the pick is complete.
+    Upon success, \a user_data will be set to whatever opaque information
+    may need to be propagated from the LB policy, or NULL if not needed.
+    \a context will be populated with context to pass to the subchannel
+    call, if needed.
 
     If the pick succeeds and a result is known immediately, a non-zero
-    value will be returned.  Otherwise, \a pick->on_complete will be invoked
+    value will be returned.  Otherwise, \a on_complete will be invoked
     once the pick is complete with its error argument set to indicate
     success or failure.
 
     Any IO should be done under the \a interested_parties \a grpc_pollset_set
     in the \a grpc_lb_policy struct. */
 int grpc_lb_policy_pick_locked(grpc_lb_policy* policy,
-                               grpc_lb_policy_pick_state* pick);
+                               const grpc_lb_policy_pick_args* pick_args,
+                               grpc_connected_subchannel** target,
+                               grpc_call_context_element* context,
+                               void** user_data, grpc_closure* on_complete);
 
 /** Perform a connected subchannel ping (see \a grpc_connected_subchannel_ping)
     against one of the connected subchannels managed by \a policy. */
@@ -158,11 +170,11 @@ void grpc_lb_policy_ping_one_locked(grpc_lb_policy* policy,
                                     grpc_closure* on_initiate,
                                     grpc_closure* on_ack);
 
-/** Cancel picks for \a pick.
+/** Cancel picks for \a target.
     The \a on_complete callback of the pending picks will be invoked with \a
     *target set to NULL. */
 void grpc_lb_policy_cancel_pick_locked(grpc_lb_policy* policy,
-                                       grpc_lb_policy_pick_state* pick,
+                                       grpc_connected_subchannel** target,
                                        grpc_error* error);
 
 /** Cancel all pending picks for which their \a initial_metadata_flags (as given