1/* @title: Per-CPU Reference Counter */
2#pragma once
3#include <compiler.h>
4#include <kassert.h>
5#include <smp/core.h>
6#include <stdatomic.h>
7#include <stdbool.h>
8#include <stdint.h>
9#include <sync/rcu.h>
10
11/*
12 * Per-CPU Reference Counter
13 *
14 * Fast, no contention reference counting by distributing RC inc/dec across
15 * per-CPU counters in fast path, using large bias and RCU GP to transition
16 * to atomic counter without lost updates */
17
18struct percpu_rc;
19typedef void (*percpu_rc_release_fn)(struct percpu_rc *);
20
21enum percpu_rc_flags {
22 PERCPU_RC_INIT_ATOMIC = 1 << 0,
23 PERCPU_RC_ALLOW_REINIT = 1 << 1,
24};
25
26#define PERCPU_COUNT_BIAS (1LL << 30)
27
28#define PERCPU_RC_DEAD (1UL << 0)
29#define PERCPU_RC_ATOMIC (1UL << 1)
30#define PERCPU_RC_PTR_MASK (~(PERCPU_RC_DEAD | PERCPU_RC_ATOMIC))
31#define PERCPU_RC_PTR(m) (int64_t *) (m & PERCPU_RC_PTR_MASK)
32
33struct percpu_rc {
34 _Atomic int64_t count;
35 _Atomic uintptr_t percpu_count_ptr;
36 percpu_rc_release_fn release;
37 bool allow_reinit;
38 struct rcu_cb rcu;
39} __cache_aligned;
40
41int percpu_rc_init(struct percpu_rc *ref, percpu_rc_release_fn release,
42 enum percpu_rc_flags flags);
43void percpu_rc_destroy(struct percpu_rc *ref);
44void percpu_rc_kill(struct percpu_rc *ref);
45void percpu_rc_reinit(struct percpu_rc *ref);
46void percpu_rc_resurrect(struct percpu_rc *ref);
47int64_t percpu_rc_read(struct percpu_rc *ref);
48
49static inline bool percpu_rc_is_percpu(uintptr_t pcpu) {
50 return (pcpu & (PERCPU_RC_DEAD | PERCPU_RC_ATOMIC)) == 0;
51}
52
53/* Notes on why TOPOC_NONE is fine here:
54 *
55 * percpu_rc is designed with the premise of "getting migrated, dec'ing the
56 * wrong counter, and all that other business will cause no problem".
57 *
58 * the eventual kill will deal with all the +'s and -'s,
59 * meaning that if you end up with
60 *
61 * CPU0: rc == 9
62 * CPU1: rc == -9
63 *
64 * the eventual cleanup goes "9 - 9 = 0" and handles it perfectly ok
65 *
66 */
67static inline void percpu_rc_get(struct percpu_rc *ref) {
68 rcu_read_lock();
69 uintptr_t pcpu =
70 atomic_load_explicit(&ref->percpu_count_ptr, memory_order_relaxed);
71
72 if (likely(percpu_rc_is_percpu(pcpu))) {
73 int64_t *counters = PERCPU_RC_PTR(pcpu);
74 counters[smp_id(cond: TOPC_NONE)]++;
75 } else {
76 atomic_fetch_add_explicit(&ref->count, 1, memory_order_relaxed);
77 }
78 rcu_read_unlock();
79}
80
81static inline void percpu_rc_put(struct percpu_rc *ref) {
82 rcu_read_lock();
83 uintptr_t pcpu =
84 atomic_load_explicit(&ref->percpu_count_ptr, memory_order_relaxed);
85
86 if (likely(percpu_rc_is_percpu(pcpu))) {
87 int64_t *counters = PERCPU_RC_PTR(pcpu);
88 counters[smp_id(cond: TOPC_NONE)]--;
89 rcu_read_unlock();
90 } else {
91 rcu_read_unlock();
92 if (atomic_fetch_sub_explicit(&ref->count, 1, memory_order_acq_rel) ==
93 1) {
94 if (ref->release)
95 ref->release(ref);
96 }
97 }
98}
99
100static inline bool percpu_rc_tryget(struct percpu_rc *ref) {
101 rcu_read_lock();
102 uintptr_t pcpu =
103 atomic_load_explicit(&ref->percpu_count_ptr, memory_order_relaxed);
104
105 if (likely(percpu_rc_is_percpu(pcpu))) {
106 int64_t *counters = PERCPU_RC_PTR(pcpu);
107 counters[smp_id(cond: TOPC_NONE)]++;
108 rcu_read_unlock();
109 return true;
110 }
111
112 rcu_read_unlock();
113 int64_t c = atomic_load_explicit(&ref->count, memory_order_relaxed);
114 do {
115 if (c <= 0)
116 return false;
117 } while (!atomic_compare_exchange_weak_explicit(
118 &ref->count, &c, c + 1, memory_order_relaxed, memory_order_relaxed));
119 return true;
120}
121
122static inline bool percpu_rc_tryget_live(struct percpu_rc *ref) {
123 rcu_read_lock();
124 uintptr_t pcpu =
125 atomic_load_explicit(&ref->percpu_count_ptr, memory_order_relaxed);
126
127 if (likely(percpu_rc_is_percpu(pcpu))) {
128 int64_t *counters = PERCPU_RC_PTR(pcpu);
129 counters[smp_id(cond: TOPC_NONE)]++;
130 rcu_read_unlock();
131 return true;
132 }
133
134 if (pcpu & PERCPU_RC_DEAD) {
135 rcu_read_unlock();
136 return false;
137 }
138
139 rcu_read_unlock();
140 int64_t c = atomic_load_explicit(&ref->count, memory_order_relaxed);
141 do {
142 if (c <= 0)
143 return false;
144 } while (!atomic_compare_exchange_weak_explicit(
145 &ref->count, &c, c + 1, memory_order_relaxed, memory_order_relaxed));
146 return true;
147}
148
149static inline bool percpu_rc_is_zero(struct percpu_rc *ref) {
150 return atomic_load_explicit(&ref->count, memory_order_relaxed) == 0;
151}
152
153static inline bool percpu_rc_is_dying(struct percpu_rc *ref) {
154 return (atomic_load_explicit(&ref->percpu_count_ptr, memory_order_relaxed) &
155 PERCPU_RC_DEAD) != 0;
156}
157