xref: /linux/lib/percpu-refcount.c (revision 96ac6d435100450f0565708d9b885ea2a7400e0a)
1 // SPDX-License-Identifier: GPL-2.0-only
2 #define pr_fmt(fmt) "%s: " fmt "\n", __func__
3 
4 #include <linux/kernel.h>
5 #include <linux/sched.h>
6 #include <linux/wait.h>
7 #include <linux/percpu-refcount.h>
8 
9 /*
10  * Initially, a percpu refcount is just a set of percpu counters. Initially, we
11  * don't try to detect the ref hitting 0 - which means that get/put can just
12  * increment or decrement the local counter. Note that the counter on a
13  * particular cpu can (and will) wrap - this is fine, when we go to shutdown the
14  * percpu counters will all sum to the correct value
15  *
16  * (More precisely: because modular arithmetic is commutative the sum of all the
17  * percpu_count vars will be equal to what it would have been if all the gets
18  * and puts were done to a single integer, even if some of the percpu integers
19  * overflow or underflow).
20  *
21  * The real trick to implementing percpu refcounts is shutdown. We can't detect
22  * the ref hitting 0 on every put - this would require global synchronization
23  * and defeat the whole purpose of using percpu refs.
24  *
25  * What we do is require the user to keep track of the initial refcount; we know
26  * the ref can't hit 0 before the user drops the initial ref, so as long as we
27  * convert to non percpu mode before the initial ref is dropped everything
28  * works.
29  *
30  * Converting to non percpu mode is done with some RCUish stuff in
31  * percpu_ref_kill. Additionally, we need a bias value so that the
32  * atomic_long_t can't hit 0 before we've added up all the percpu refs.
33  */
34 
35 #define PERCPU_COUNT_BIAS	(1LU << (BITS_PER_LONG - 1))
36 
37 static DEFINE_SPINLOCK(percpu_ref_switch_lock);
38 static DECLARE_WAIT_QUEUE_HEAD(percpu_ref_switch_waitq);
39 
40 static unsigned long __percpu *percpu_count_ptr(struct percpu_ref *ref)
41 {
42 	return (unsigned long __percpu *)
43 		(ref->percpu_count_ptr & ~__PERCPU_REF_ATOMIC_DEAD);
44 }
45 
46 /**
47  * percpu_ref_init - initialize a percpu refcount
48  * @ref: percpu_ref to initialize
49  * @release: function which will be called when refcount hits 0
50  * @flags: PERCPU_REF_INIT_* flags
51  * @gfp: allocation mask to use
52  *
53  * Initializes @ref.  If @flags is zero, @ref starts in percpu mode with a
54  * refcount of 1; analagous to atomic_long_set(ref, 1).  See the
55  * definitions of PERCPU_REF_INIT_* flags for flag behaviors.
56  *
57  * Note that @release must not sleep - it may potentially be called from RCU
58  * callback context by percpu_ref_kill().
59  */
60 int percpu_ref_init(struct percpu_ref *ref, percpu_ref_func_t *release,
61 		    unsigned int flags, gfp_t gfp)
62 {
63 	size_t align = max_t(size_t, 1 << __PERCPU_REF_FLAG_BITS,
64 			     __alignof__(unsigned long));
65 	unsigned long start_count = 0;
66 
67 	ref->percpu_count_ptr = (unsigned long)
68 		__alloc_percpu_gfp(sizeof(unsigned long), align, gfp);
69 	if (!ref->percpu_count_ptr)
70 		return -ENOMEM;
71 
72 	ref->force_atomic = flags & PERCPU_REF_INIT_ATOMIC;
73 
74 	if (flags & (PERCPU_REF_INIT_ATOMIC | PERCPU_REF_INIT_DEAD))
75 		ref->percpu_count_ptr |= __PERCPU_REF_ATOMIC;
76 	else
77 		start_count += PERCPU_COUNT_BIAS;
78 
79 	if (flags & PERCPU_REF_INIT_DEAD)
80 		ref->percpu_count_ptr |= __PERCPU_REF_DEAD;
81 	else
82 		start_count++;
83 
84 	atomic_long_set(&ref->count, start_count);
85 
86 	ref->release = release;
87 	ref->confirm_switch = NULL;
88 	return 0;
89 }
90 EXPORT_SYMBOL_GPL(percpu_ref_init);
91 
92 /**
93  * percpu_ref_exit - undo percpu_ref_init()
94  * @ref: percpu_ref to exit
95  *
96  * This function exits @ref.  The caller is responsible for ensuring that
97  * @ref is no longer in active use.  The usual places to invoke this
98  * function from are the @ref->release() callback or in init failure path
99  * where percpu_ref_init() succeeded but other parts of the initialization
100  * of the embedding object failed.
101  */
102 void percpu_ref_exit(struct percpu_ref *ref)
103 {
104 	unsigned long __percpu *percpu_count = percpu_count_ptr(ref);
105 
106 	if (percpu_count) {
107 		/* non-NULL confirm_switch indicates switching in progress */
108 		WARN_ON_ONCE(ref->confirm_switch);
109 		free_percpu(percpu_count);
110 		ref->percpu_count_ptr = __PERCPU_REF_ATOMIC_DEAD;
111 	}
112 }
113 EXPORT_SYMBOL_GPL(percpu_ref_exit);
114 
115 static void percpu_ref_call_confirm_rcu(struct rcu_head *rcu)
116 {
117 	struct percpu_ref *ref = container_of(rcu, struct percpu_ref, rcu);
118 
119 	ref->confirm_switch(ref);
120 	ref->confirm_switch = NULL;
121 	wake_up_all(&percpu_ref_switch_waitq);
122 
123 	/* drop ref from percpu_ref_switch_to_atomic() */
124 	percpu_ref_put(ref);
125 }
126 
127 static void percpu_ref_switch_to_atomic_rcu(struct rcu_head *rcu)
128 {
129 	struct percpu_ref *ref = container_of(rcu, struct percpu_ref, rcu);
130 	unsigned long __percpu *percpu_count = percpu_count_ptr(ref);
131 	unsigned long count = 0;
132 	int cpu;
133 
134 	for_each_possible_cpu(cpu)
135 		count += *per_cpu_ptr(percpu_count, cpu);
136 
137 	pr_debug("global %ld percpu %ld",
138 		 atomic_long_read(&ref->count), (long)count);
139 
140 	/*
141 	 * It's crucial that we sum the percpu counters _before_ adding the sum
142 	 * to &ref->count; since gets could be happening on one cpu while puts
143 	 * happen on another, adding a single cpu's count could cause
144 	 * @ref->count to hit 0 before we've got a consistent value - but the
145 	 * sum of all the counts will be consistent and correct.
146 	 *
147 	 * Subtracting the bias value then has to happen _after_ adding count to
148 	 * &ref->count; we need the bias value to prevent &ref->count from
149 	 * reaching 0 before we add the percpu counts. But doing it at the same
150 	 * time is equivalent and saves us atomic operations:
151 	 */
152 	atomic_long_add((long)count - PERCPU_COUNT_BIAS, &ref->count);
153 
154 	WARN_ONCE(atomic_long_read(&ref->count) <= 0,
155 		  "percpu ref (%ps) <= 0 (%ld) after switching to atomic",
156 		  ref->release, atomic_long_read(&ref->count));
157 
158 	/* @ref is viewed as dead on all CPUs, send out switch confirmation */
159 	percpu_ref_call_confirm_rcu(rcu);
160 }
161 
162 static void percpu_ref_noop_confirm_switch(struct percpu_ref *ref)
163 {
164 }
165 
166 static void __percpu_ref_switch_to_atomic(struct percpu_ref *ref,
167 					  percpu_ref_func_t *confirm_switch)
168 {
169 	if (ref->percpu_count_ptr & __PERCPU_REF_ATOMIC) {
170 		if (confirm_switch)
171 			confirm_switch(ref);
172 		return;
173 	}
174 
175 	/* switching from percpu to atomic */
176 	ref->percpu_count_ptr |= __PERCPU_REF_ATOMIC;
177 
178 	/*
179 	 * Non-NULL ->confirm_switch is used to indicate that switching is
180 	 * in progress.  Use noop one if unspecified.
181 	 */
182 	ref->confirm_switch = confirm_switch ?: percpu_ref_noop_confirm_switch;
183 
184 	percpu_ref_get(ref);	/* put after confirmation */
185 	call_rcu(&ref->rcu, percpu_ref_switch_to_atomic_rcu);
186 }
187 
188 static void __percpu_ref_switch_to_percpu(struct percpu_ref *ref)
189 {
190 	unsigned long __percpu *percpu_count = percpu_count_ptr(ref);
191 	int cpu;
192 
193 	BUG_ON(!percpu_count);
194 
195 	if (!(ref->percpu_count_ptr & __PERCPU_REF_ATOMIC))
196 		return;
197 
198 	atomic_long_add(PERCPU_COUNT_BIAS, &ref->count);
199 
200 	/*
201 	 * Restore per-cpu operation.  smp_store_release() is paired
202 	 * with READ_ONCE() in __ref_is_percpu() and guarantees that the
203 	 * zeroing is visible to all percpu accesses which can see the
204 	 * following __PERCPU_REF_ATOMIC clearing.
205 	 */
206 	for_each_possible_cpu(cpu)
207 		*per_cpu_ptr(percpu_count, cpu) = 0;
208 
209 	smp_store_release(&ref->percpu_count_ptr,
210 			  ref->percpu_count_ptr & ~__PERCPU_REF_ATOMIC);
211 }
212 
213 static void __percpu_ref_switch_mode(struct percpu_ref *ref,
214 				     percpu_ref_func_t *confirm_switch)
215 {
216 	lockdep_assert_held(&percpu_ref_switch_lock);
217 
218 	/*
219 	 * If the previous ATOMIC switching hasn't finished yet, wait for
220 	 * its completion.  If the caller ensures that ATOMIC switching
221 	 * isn't in progress, this function can be called from any context.
222 	 */
223 	wait_event_lock_irq(percpu_ref_switch_waitq, !ref->confirm_switch,
224 			    percpu_ref_switch_lock);
225 
226 	if (ref->force_atomic || (ref->percpu_count_ptr & __PERCPU_REF_DEAD))
227 		__percpu_ref_switch_to_atomic(ref, confirm_switch);
228 	else
229 		__percpu_ref_switch_to_percpu(ref);
230 }
231 
232 /**
233  * percpu_ref_switch_to_atomic - switch a percpu_ref to atomic mode
234  * @ref: percpu_ref to switch to atomic mode
235  * @confirm_switch: optional confirmation callback
236  *
237  * There's no reason to use this function for the usual reference counting.
238  * Use percpu_ref_kill[_and_confirm]().
239  *
240  * Schedule switching of @ref to atomic mode.  All its percpu counts will
241  * be collected to the main atomic counter.  On completion, when all CPUs
242  * are guaraneed to be in atomic mode, @confirm_switch, which may not
243  * block, is invoked.  This function may be invoked concurrently with all
244  * the get/put operations and can safely be mixed with kill and reinit
245  * operations.  Note that @ref will stay in atomic mode across kill/reinit
246  * cycles until percpu_ref_switch_to_percpu() is called.
247  *
248  * This function may block if @ref is in the process of switching to atomic
249  * mode.  If the caller ensures that @ref is not in the process of
250  * switching to atomic mode, this function can be called from any context.
251  */
252 void percpu_ref_switch_to_atomic(struct percpu_ref *ref,
253 				 percpu_ref_func_t *confirm_switch)
254 {
255 	unsigned long flags;
256 
257 	spin_lock_irqsave(&percpu_ref_switch_lock, flags);
258 
259 	ref->force_atomic = true;
260 	__percpu_ref_switch_mode(ref, confirm_switch);
261 
262 	spin_unlock_irqrestore(&percpu_ref_switch_lock, flags);
263 }
264 EXPORT_SYMBOL_GPL(percpu_ref_switch_to_atomic);
265 
266 /**
267  * percpu_ref_switch_to_atomic_sync - switch a percpu_ref to atomic mode
268  * @ref: percpu_ref to switch to atomic mode
269  *
270  * Schedule switching the ref to atomic mode, and wait for the
271  * switch to complete.  Caller must ensure that no other thread
272  * will switch back to percpu mode.
273  */
274 void percpu_ref_switch_to_atomic_sync(struct percpu_ref *ref)
275 {
276 	percpu_ref_switch_to_atomic(ref, NULL);
277 	wait_event(percpu_ref_switch_waitq, !ref->confirm_switch);
278 }
279 EXPORT_SYMBOL_GPL(percpu_ref_switch_to_atomic_sync);
280 
281 /**
282  * percpu_ref_switch_to_percpu - switch a percpu_ref to percpu mode
283  * @ref: percpu_ref to switch to percpu mode
284  *
285  * There's no reason to use this function for the usual reference counting.
286  * To re-use an expired ref, use percpu_ref_reinit().
287  *
288  * Switch @ref to percpu mode.  This function may be invoked concurrently
289  * with all the get/put operations and can safely be mixed with kill and
290  * reinit operations.  This function reverses the sticky atomic state set
291  * by PERCPU_REF_INIT_ATOMIC or percpu_ref_switch_to_atomic().  If @ref is
292  * dying or dead, the actual switching takes place on the following
293  * percpu_ref_reinit().
294  *
295  * This function may block if @ref is in the process of switching to atomic
296  * mode.  If the caller ensures that @ref is not in the process of
297  * switching to atomic mode, this function can be called from any context.
298  */
299 void percpu_ref_switch_to_percpu(struct percpu_ref *ref)
300 {
301 	unsigned long flags;
302 
303 	spin_lock_irqsave(&percpu_ref_switch_lock, flags);
304 
305 	ref->force_atomic = false;
306 	__percpu_ref_switch_mode(ref, NULL);
307 
308 	spin_unlock_irqrestore(&percpu_ref_switch_lock, flags);
309 }
310 EXPORT_SYMBOL_GPL(percpu_ref_switch_to_percpu);
311 
312 /**
313  * percpu_ref_kill_and_confirm - drop the initial ref and schedule confirmation
314  * @ref: percpu_ref to kill
315  * @confirm_kill: optional confirmation callback
316  *
317  * Equivalent to percpu_ref_kill() but also schedules kill confirmation if
318  * @confirm_kill is not NULL.  @confirm_kill, which may not block, will be
319  * called after @ref is seen as dead from all CPUs at which point all
320  * further invocations of percpu_ref_tryget_live() will fail.  See
321  * percpu_ref_tryget_live() for details.
322  *
323  * This function normally doesn't block and can be called from any context
324  * but it may block if @confirm_kill is specified and @ref is in the
325  * process of switching to atomic mode by percpu_ref_switch_to_atomic().
326  *
327  * There are no implied RCU grace periods between kill and release.
328  */
329 void percpu_ref_kill_and_confirm(struct percpu_ref *ref,
330 				 percpu_ref_func_t *confirm_kill)
331 {
332 	unsigned long flags;
333 
334 	spin_lock_irqsave(&percpu_ref_switch_lock, flags);
335 
336 	WARN_ONCE(ref->percpu_count_ptr & __PERCPU_REF_DEAD,
337 		  "%s called more than once on %ps!", __func__, ref->release);
338 
339 	ref->percpu_count_ptr |= __PERCPU_REF_DEAD;
340 	__percpu_ref_switch_mode(ref, confirm_kill);
341 	percpu_ref_put(ref);
342 
343 	spin_unlock_irqrestore(&percpu_ref_switch_lock, flags);
344 }
345 EXPORT_SYMBOL_GPL(percpu_ref_kill_and_confirm);
346 
347 /**
348  * percpu_ref_reinit - re-initialize a percpu refcount
349  * @ref: perpcu_ref to re-initialize
350  *
351  * Re-initialize @ref so that it's in the same state as when it finished
352  * percpu_ref_init() ignoring %PERCPU_REF_INIT_DEAD.  @ref must have been
353  * initialized successfully and reached 0 but not exited.
354  *
355  * Note that percpu_ref_tryget[_live]() are safe to perform on @ref while
356  * this function is in progress.
357  */
358 void percpu_ref_reinit(struct percpu_ref *ref)
359 {
360 	WARN_ON_ONCE(!percpu_ref_is_zero(ref));
361 
362 	percpu_ref_resurrect(ref);
363 }
364 EXPORT_SYMBOL_GPL(percpu_ref_reinit);
365 
366 /**
367  * percpu_ref_resurrect - modify a percpu refcount from dead to live
368  * @ref: perpcu_ref to resurrect
369  *
370  * Modify @ref so that it's in the same state as before percpu_ref_kill() was
371  * called. @ref must be dead but must not yet have exited.
372  *
373  * If @ref->release() frees @ref then the caller is responsible for
374  * guaranteeing that @ref->release() does not get called while this
375  * function is in progress.
376  *
377  * Note that percpu_ref_tryget[_live]() are safe to perform on @ref while
378  * this function is in progress.
379  */
380 void percpu_ref_resurrect(struct percpu_ref *ref)
381 {
382 	unsigned long __percpu *percpu_count;
383 	unsigned long flags;
384 
385 	spin_lock_irqsave(&percpu_ref_switch_lock, flags);
386 
387 	WARN_ON_ONCE(!(ref->percpu_count_ptr & __PERCPU_REF_DEAD));
388 	WARN_ON_ONCE(__ref_is_percpu(ref, &percpu_count));
389 
390 	ref->percpu_count_ptr &= ~__PERCPU_REF_DEAD;
391 	percpu_ref_get(ref);
392 	__percpu_ref_switch_mode(ref, NULL);
393 
394 	spin_unlock_irqrestore(&percpu_ref_switch_lock, flags);
395 }
396 EXPORT_SYMBOL_GPL(percpu_ref_resurrect);
397