1 // SPDX-License-Identifier: GPL-2.0
3 #include <linux/objpool.h>
4 #include <linux/slab.h>
5 #include <linux/vmalloc.h>
6 #include <linux/atomic.h>
7 #include <linux/irqflags.h>
8 #include <linux/cpumask.h>
9 #include <linux/log2.h>
12 * objpool: ring-array based lockless MPMC/FIFO queues
14 * Copyright: wuqiang.matt@bytedance.com,mhiramat@kernel.org
17 /* initialize percpu objpool_slot */
19 objpool_init_percpu_slot(struct objpool_head
*pool
,
20 struct objpool_slot
*slot
,
21 int nodes
, void *context
,
22 objpool_init_obj_cb objinit
)
24 void *obj
= (void *)&slot
->entries
[pool
->capacity
];
27 /* initialize elements of percpu objpool_slot */
28 slot
->mask
= pool
->capacity
- 1;
30 for (i
= 0; i
< nodes
; i
++) {
32 int rc
= objinit(obj
, context
);
36 slot
->entries
[slot
->tail
& slot
->mask
] = obj
;
37 obj
= obj
+ pool
->obj_size
;
39 slot
->last
= slot
->tail
;
46 /* allocate and initialize percpu slots */
48 objpool_init_percpu_slots(struct objpool_head
*pool
, int nr_objs
,
49 void *context
, objpool_init_obj_cb objinit
)
53 for (i
= 0; i
< pool
->nr_cpus
; i
++) {
55 struct objpool_slot
*slot
;
58 /* skip the cpu node which could never be present */
62 /* compute how many objects to be allocated with this slot */
63 nodes
= nr_objs
/ num_possible_cpus();
64 if (cpu_count
< (nr_objs
% num_possible_cpus()))
68 size
= struct_size(slot
, entries
, pool
->capacity
) +
69 pool
->obj_size
* nodes
;
72 * here we allocate percpu-slot & objs together in a single
73 * allocation to make it more compact, taking advantage of
74 * warm caches and TLB hits. in default vmalloc is used to
75 * reduce the pressure of kernel slab system. as we know,
76 * mimimal size of vmalloc is one page since vmalloc would
77 * always align the requested size to page size
79 if (pool
->gfp
& GFP_ATOMIC
)
80 slot
= kmalloc_node(size
, pool
->gfp
, cpu_to_node(i
));
82 slot
= __vmalloc_node(size
, sizeof(void *), pool
->gfp
,
83 cpu_to_node(i
), __builtin_return_address(0));
86 memset(slot
, 0, size
);
87 pool
->cpu_slots
[i
] = slot
;
89 /* initialize the objpool_slot of cpu node i */
90 rc
= objpool_init_percpu_slot(pool
, slot
, nodes
, context
, objinit
);
98 /* cleanup all percpu slots of the object pool */
99 static void objpool_fini_percpu_slots(struct objpool_head
*pool
)
103 if (!pool
->cpu_slots
)
106 for (i
= 0; i
< pool
->nr_cpus
; i
++)
107 kvfree(pool
->cpu_slots
[i
]);
108 kfree(pool
->cpu_slots
);
111 /* initialize object pool and pre-allocate objects */
112 int objpool_init(struct objpool_head
*pool
, int nr_objs
, int object_size
,
113 gfp_t gfp
, void *context
, objpool_init_obj_cb objinit
,
114 objpool_fini_cb release
)
116 int rc
, capacity
, slot_size
;
118 /* check input parameters */
119 if (nr_objs
<= 0 || nr_objs
> OBJPOOL_NR_OBJECT_MAX
||
120 object_size
<= 0 || object_size
> OBJPOOL_OBJECT_SIZE_MAX
)
123 /* align up to unsigned long size */
124 object_size
= ALIGN(object_size
, sizeof(long));
126 /* calculate capacity of percpu objpool_slot */
127 capacity
= roundup_pow_of_two(nr_objs
);
131 /* initialize objpool pool */
132 memset(pool
, 0, sizeof(struct objpool_head
));
133 pool
->nr_cpus
= nr_cpu_ids
;
134 pool
->obj_size
= object_size
;
135 pool
->capacity
= capacity
;
136 pool
->gfp
= gfp
& ~__GFP_ZERO
;
137 pool
->context
= context
;
138 pool
->release
= release
;
139 slot_size
= pool
->nr_cpus
* sizeof(struct objpool_slot
);
140 pool
->cpu_slots
= kzalloc(slot_size
, pool
->gfp
);
141 if (!pool
->cpu_slots
)
144 /* initialize per-cpu slots */
145 rc
= objpool_init_percpu_slots(pool
, nr_objs
, context
, objinit
);
147 objpool_fini_percpu_slots(pool
);
149 refcount_set(&pool
->ref
, pool
->nr_objs
+ 1);
153 EXPORT_SYMBOL_GPL(objpool_init
);
155 /* adding object to slot, abort if the slot was already full */
157 objpool_try_add_slot(void *obj
, struct objpool_head
*pool
, int cpu
)
159 struct objpool_slot
*slot
= pool
->cpu_slots
[cpu
];
162 /* loading tail and head as a local snapshot, tail first */
163 tail
= READ_ONCE(slot
->tail
);
166 head
= READ_ONCE(slot
->head
);
167 /* fault caught: something must be wrong */
168 WARN_ON_ONCE(tail
- head
> pool
->nr_objs
);
169 } while (!try_cmpxchg_acquire(&slot
->tail
, &tail
, tail
+ 1));
171 /* now the tail position is reserved for the given obj */
172 WRITE_ONCE(slot
->entries
[tail
& slot
->mask
], obj
);
173 /* update sequence to make this obj available for pop() */
174 smp_store_release(&slot
->last
, tail
+ 1);
179 /* reclaim an object to object pool */
180 int objpool_push(void *obj
, struct objpool_head
*pool
)
185 /* disable local irq to avoid preemption & interruption */
186 raw_local_irq_save(flags
);
187 rc
= objpool_try_add_slot(obj
, pool
, raw_smp_processor_id());
188 raw_local_irq_restore(flags
);
192 EXPORT_SYMBOL_GPL(objpool_push
);
194 /* try to retrieve object from slot */
195 static inline void *objpool_try_get_slot(struct objpool_head
*pool
, int cpu
)
197 struct objpool_slot
*slot
= pool
->cpu_slots
[cpu
];
198 /* load head snapshot, other cpus may change it */
199 uint32_t head
= smp_load_acquire(&slot
->head
);
201 while (head
!= READ_ONCE(slot
->last
)) {
205 * data visibility of 'last' and 'head' could be out of
206 * order since memory updating of 'last' and 'head' are
207 * performed in push() and pop() independently
209 * before any retrieving attempts, pop() must guarantee
210 * 'last' is behind 'head', that is to say, there must
211 * be available objects in slot, which could be ensured
212 * by condition 'last != head && last - head <= nr_objs'
213 * that is equivalent to 'last - head - 1 < nr_objs' as
214 * 'last' and 'head' are both unsigned int32
216 if (READ_ONCE(slot
->last
) - head
- 1 >= pool
->nr_objs
) {
217 head
= READ_ONCE(slot
->head
);
221 /* obj must be retrieved before moving forward head */
222 obj
= READ_ONCE(slot
->entries
[head
& slot
->mask
]);
224 /* move head forward to mark it's consumption */
225 if (try_cmpxchg_release(&slot
->head
, &head
, head
+ 1))
232 /* allocate an object from object pool */
233 void *objpool_pop(struct objpool_head
*pool
)
239 /* disable local irq to avoid preemption & interruption */
240 raw_local_irq_save(flags
);
242 cpu
= raw_smp_processor_id();
243 for (i
= 0; i
< num_possible_cpus(); i
++) {
244 obj
= objpool_try_get_slot(pool
, cpu
);
247 cpu
= cpumask_next_wrap(cpu
, cpu_possible_mask
, -1, 1);
249 raw_local_irq_restore(flags
);
253 EXPORT_SYMBOL_GPL(objpool_pop
);
255 /* release whole objpool forcely */
256 void objpool_free(struct objpool_head
*pool
)
258 if (!pool
->cpu_slots
)
261 /* release percpu slots */
262 objpool_fini_percpu_slots(pool
);
264 /* call user's cleanup callback if provided */
266 pool
->release(pool
, pool
->context
);
268 EXPORT_SYMBOL_GPL(objpool_free
);
270 /* drop the allocated object, rather reclaim it to objpool */
271 int objpool_drop(void *obj
, struct objpool_head
*pool
)
276 if (refcount_dec_and_test(&pool
->ref
)) {
283 EXPORT_SYMBOL_GPL(objpool_drop
);
285 /* drop unused objects and defref objpool for releasing */
286 void objpool_fini(struct objpool_head
*pool
)
288 int count
= 1; /* extra ref for objpool itself */
290 /* drop all remained objects from objpool */
291 while (objpool_pop(pool
))
294 if (refcount_sub_and_test(count
, &pool
->ref
))
297 EXPORT_SYMBOL_GPL(objpool_fini
);