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