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