percpu-refcount: implement percpu_tryget() along with percpu_ref_kill_and_confirm()

Implement percpu_tryget() which stops giving out references once the percpu_ref is visible as killed. Because the refcnt is per-cpu, different CPUs will start to see a refcnt as killed at different points in time and tryget() may continue to succeed on subset of cpus for a while after percpu_ref_kill() returns. For use cases where it's necessary to know when all CPUs start to see the refcnt as dead, percpu_ref_kill_and_confirm() is added. The new function takes an extra argument @confirm_kill which is invoked when the refcnt is guaranteed to be viewed as killed on all CPUs. While this isn't the prettiest interface, it doesn't force synchronous wait and is much safer than requiring the caller to do its own call_rcu(). v2: Patch description rephrased to emphasize that tryget() may continue to succeed on some CPUs after kill() returns as suggested by Kent. v3: Function comment in percpu_ref_kill_and_confirm() updated warning people to not depend on the implied RCU grace period from the confirm callback as it's an implementation detail. Signed-off-by: Tejun Heo <tj@kernel.org> Slightly-Grumpily-Acked-by: Kent Overstreet <koverstreet@google.com>

percpu-refcount: implement percpu_tryget() along with percpu_ref_kill_and_confirm()
Implement percpu_tryget() which stops giving out references once the percpu_ref is visible as killed. Because the refcnt is per-cpu, different CPUs will start to see a refcnt as killed at different points in time and tryget() may continue to succeed on subset of cpus for a while after percpu_ref_kill() returns. For use cases where it's necessary to know when all CPUs start to see the refcnt as dead, percpu_ref_kill_and_confirm() is added. The new function takes an extra argument @confirm_kill which is invoked when the refcnt is guaranteed to be viewed as killed on all CPUs. While this isn't the prettiest interface, it doesn't force synchronous wait and is much safer than requiring the caller to do its own call_rcu(). v2: Patch description rephrased to emphasize that tryget() may continue to succeed on some CPUs after kill() returns as suggested by Kent. v3: Function comment in percpu_ref_kill_and_confirm() updated warning people to not depend on the implied RCU grace period from the confirm callback as it's an implementation detail. Signed-off-by: Tejun Heo <tj@kernel.org> Slightly-Grumpily-Acked-by: Kent Overstreet <koverstreet@google.com>
Tejun Heo
1 parent bc497bd33b
Showing 2 changed files with 66 additions and 7 deletions Side-by-side Diff
include/linux/percpu-refcount.h
lib/percpu-refcount.c
@@ -63,14 +63,31 @@
 	 */
 	unsigned __percpu	*pcpu_count;
 	percpu_ref_func_t	*release;
+	percpu_ref_func_t	*confirm_kill;
 	struct rcu_head		rcu;
 };
  
 int __must_check percpu_ref_init(struct percpu_ref *ref,
 				 percpu_ref_func_t *release);
 void percpu_ref_cancel_init(struct percpu_ref *ref);
-void percpu_ref_kill(struct percpu_ref *ref);
+void percpu_ref_kill_and_confirm(struct percpu_ref *ref,
+				 percpu_ref_func_t *confirm_kill);
  
+/**
+ * percpu_ref_kill - drop the initial ref
+ * @ref: percpu_ref to kill
+ *
+ * Must be used to drop the initial ref on a percpu refcount; must be called
+ * precisely once before shutdown.
+ *
+ * Puts @ref in non percpu mode, then does a call_rcu() before gathering up the
+ * percpu counters and dropping the initial ref.
+ */
+static inline void percpu_ref_kill(struct percpu_ref *ref)
+{
+	return percpu_ref_kill_and_confirm(ref, NULL);
+}
+
 #define PCPU_STATUS_BITS	2
 #define PCPU_STATUS_MASK	((1 << PCPU_STATUS_BITS) - 1)
 #define PCPU_REF_PTR		0
@@ -98,6 +115,37 @@
 		atomic_inc(&ref->count);
  
 	rcu_read_unlock();
+}
+
+/**
+ * percpu_ref_tryget - try to increment a percpu refcount
+ * @ref: percpu_ref to try-get
+ *
+ * Increment a percpu refcount unless it has already been killed.  Returns
+ * %true on success; %false on failure.
+ *
+ * Completion of percpu_ref_kill() in itself doesn't guarantee that tryget
+ * will fail.  For such guarantee, percpu_ref_kill_and_confirm() should be
+ * used.  After the confirm_kill callback is invoked, it's guaranteed that
+ * no new reference will be given out by percpu_ref_tryget().
+ */
+static inline bool percpu_ref_tryget(struct percpu_ref *ref)
+{
+	unsigned __percpu *pcpu_count;
+	int ret = false;
+
+	rcu_read_lock();
+
+	pcpu_count = ACCESS_ONCE(ref->pcpu_count);
+
+	if (likely(REF_STATUS(pcpu_count) == PCPU_REF_PTR)) {
+		__this_cpu_inc(*pcpu_count);
+		ret = true;
+	}
+
+	rcu_read_unlock();
+
+	return ret;
 }
  
 /**
@@ -118,6 +118,10 @@
  
 	atomic_add((int) count - PCPU_COUNT_BIAS, &ref->count);
  
+	/* @ref is viewed as dead on all CPUs, send out kill confirmation */
+	if (ref->confirm_kill)
+		ref->confirm_kill(ref);
+
 	/*
 	 * Now we're in single atomic_t mode with a consistent refcount, so it's
 	 * safe to drop our initial ref:
  
  
  
  
  
@@ -126,22 +130,29 @@
 }
  
 /**
- * percpu_ref_kill - safely drop initial ref
+ * percpu_ref_kill_and_confirm - drop the initial ref and schedule confirmation
  * @ref: percpu_ref to kill
+ * @confirm_kill: optional confirmation callback
  *
- * Must be used to drop the initial ref on a percpu refcount; must be called
- * precisely once before shutdown.
+ * Equivalent to percpu_ref_kill() but also schedules kill confirmation if
+ * @confirm_kill is not NULL.  @confirm_kill, which may not block, will be
+ * called after @ref is seen as dead from all CPUs - all further
+ * invocations of percpu_ref_tryget() will fail.  See percpu_ref_tryget()
+ * for more details.
  *
- * Puts @ref in non percpu mode, then does a call_rcu() before gathering up the
- * percpu counters and dropping the initial ref.
+ * Due to the way percpu_ref is implemented, @confirm_kill will be called
+ * after at least one full RCU grace period has passed but this is an
+ * implementation detail and callers must not depend on it.
  */
-void percpu_ref_kill(struct percpu_ref *ref)
+void percpu_ref_kill_and_confirm(struct percpu_ref *ref,
+				 percpu_ref_func_t *confirm_kill)
 {
 	WARN_ONCE(REF_STATUS(ref->pcpu_count) == PCPU_REF_DEAD,
 		  "percpu_ref_kill() called more than once!\n");
  
 	ref->pcpu_count = (unsigned __percpu *)
 		(((unsigned long) ref->pcpu_count)|PCPU_REF_DEAD);
+	ref->confirm_kill = confirm_kill;
  
 	call_rcu(&ref->rcu, percpu_ref_kill_rcu);
 }
...	...	@@ -63,14 +63,31 @@
63	63	*/
64	64	unsigned __percpu *pcpu_count;
65	65	percpu_ref_func_t *release;
	66	+ percpu_ref_func_t *confirm_kill;
66	67	struct rcu_head rcu;
67	68	};
68	69
69	70	int __must_check percpu_ref_init(struct percpu_ref *ref,
70	71	percpu_ref_func_t *release);
71	72	void percpu_ref_cancel_init(struct percpu_ref *ref);
72		-void percpu_ref_kill(struct percpu_ref *ref);
	73	+void percpu_ref_kill_and_confirm(struct percpu_ref *ref,
	74	+ percpu_ref_func_t *confirm_kill);
73	75
	76	+/**
	77	+ * percpu_ref_kill - drop the initial ref
	78	+ * @ref: percpu_ref to kill
	79	+ *
	80	+ * Must be used to drop the initial ref on a percpu refcount; must be called
	81	+ * precisely once before shutdown.
	82	+ *
	83	+ * Puts @ref in non percpu mode, then does a call_rcu() before gathering up the
	84	+ * percpu counters and dropping the initial ref.
	85	+ */
	86	+static inline void percpu_ref_kill(struct percpu_ref *ref)
	87	+{
	88	+ return percpu_ref_kill_and_confirm(ref, NULL);
	89	+}
	90	+
74	91	#define PCPU_STATUS_BITS 2
75	92	#define PCPU_STATUS_MASK ((1 << PCPU_STATUS_BITS) - 1)
76	93	#define PCPU_REF_PTR 0
...	...	@@ -98,6 +115,37 @@
98	115	atomic_inc(&ref->count);
99	116
100	117	rcu_read_unlock();
	118	+}
	119	+
	120	+/**
	121	+ * percpu_ref_tryget - try to increment a percpu refcount
	122	+ * @ref: percpu_ref to try-get
	123	+ *
	124	+ * Increment a percpu refcount unless it has already been killed. Returns
	125	+ * %true on success; %false on failure.
	126	+ *
	127	+ * Completion of percpu_ref_kill() in itself doesn't guarantee that tryget
	128	+ * will fail. For such guarantee, percpu_ref_kill_and_confirm() should be
	129	+ * used. After the confirm_kill callback is invoked, it's guaranteed that
	130	+ * no new reference will be given out by percpu_ref_tryget().
	131	+ */
	132	+static inline bool percpu_ref_tryget(struct percpu_ref *ref)
	133	+{
	134	+ unsigned __percpu *pcpu_count;
	135	+ int ret = false;
	136	+
	137	+ rcu_read_lock();
	138	+
	139	+ pcpu_count = ACCESS_ONCE(ref->pcpu_count);
	140	+
	141	+ if (likely(REF_STATUS(pcpu_count) == PCPU_REF_PTR)) {
	142	+ __this_cpu_inc(*pcpu_count);
	143	+ ret = true;
	144	+ }
	145	+
	146	+ rcu_read_unlock();
	147	+
	148	+ return ret;
101	149	}
102	150
103	151	/**
...	...	@@ -118,6 +118,10 @@
118	118
119	119	atomic_add((int) count - PCPU_COUNT_BIAS, &ref->count);
120	120
	121	+ /* @ref is viewed as dead on all CPUs, send out kill confirmation */
	122	+ if (ref->confirm_kill)
	123	+ ref->confirm_kill(ref);
	124	+
121	125	/*
122	126	* Now we're in single atomic_t mode with a consistent refcount, so it's
123	127	* safe to drop our initial ref:
124	128
125	129
126	130
127	131
128	132
...	...	@@ -126,22 +130,29 @@
126	130	}
127	131
128	132	/**
129		- * percpu_ref_kill - safely drop initial ref
	133	+ * percpu_ref_kill_and_confirm - drop the initial ref and schedule confirmation
130	134	* @ref: percpu_ref to kill
	135	+ * @confirm_kill: optional confirmation callback
131	136	*
132		- * Must be used to drop the initial ref on a percpu refcount; must be called
133		- * precisely once before shutdown.
	137	+ * Equivalent to percpu_ref_kill() but also schedules kill confirmation if
	138	+ * @confirm_kill is not NULL. @confirm_kill, which may not block, will be
	139	+ * called after @ref is seen as dead from all CPUs - all further
	140	+ * invocations of percpu_ref_tryget() will fail. See percpu_ref_tryget()
	141	+ * for more details.
134	142	*
135		- * Puts @ref in non percpu mode, then does a call_rcu() before gathering up the
136		- * percpu counters and dropping the initial ref.
	143	+ * Due to the way percpu_ref is implemented, @confirm_kill will be called
	144	+ * after at least one full RCU grace period has passed but this is an
	145	+ * implementation detail and callers must not depend on it.
137	146	*/
138		-void percpu_ref_kill(struct percpu_ref *ref)
	147	+void percpu_ref_kill_and_confirm(struct percpu_ref *ref,
	148	+ percpu_ref_func_t *confirm_kill)
139	149	{
140	150	WARN_ONCE(REF_STATUS(ref->pcpu_count) == PCPU_REF_DEAD,
141	151	"percpu_ref_kill() called more than once!\n");
142	152
143	153	ref->pcpu_count = (unsigned __percpu *)
144	154	(((unsigned long) ref->pcpu_count)\|PCPU_REF_DEAD);
	155	+ ref->confirm_kill = confirm_kill;
145	156
146	157	call_rcu(&ref->rcu, percpu_ref_kill_rcu);
147	158	}