17c8199e2SAlexei Starovoitov // SPDX-License-Identifier: GPL-2.0-only
27c8199e2SAlexei Starovoitov /* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
37c8199e2SAlexei Starovoitov #include <linux/mm.h>
47c8199e2SAlexei Starovoitov #include <linux/llist.h>
57c8199e2SAlexei Starovoitov #include <linux/bpf.h>
67c8199e2SAlexei Starovoitov #include <linux/irq_work.h>
77c8199e2SAlexei Starovoitov #include <linux/bpf_mem_alloc.h>
87c8199e2SAlexei Starovoitov #include <linux/memcontrol.h>
97c8199e2SAlexei Starovoitov #include <asm/local.h>
107c8199e2SAlexei Starovoitov
117c8199e2SAlexei Starovoitov /* Any context (including NMI) BPF specific memory allocator.
127c8199e2SAlexei Starovoitov *
137c8199e2SAlexei Starovoitov * Tracing BPF programs can attach to kprobe and fentry. Hence they
147c8199e2SAlexei Starovoitov * run in unknown context where calling plain kmalloc() might not be safe.
157c8199e2SAlexei Starovoitov *
167c8199e2SAlexei Starovoitov * Front-end kmalloc() with per-cpu per-bucket cache of free elements.
177c8199e2SAlexei Starovoitov * Refill this cache asynchronously from irq_work.
187c8199e2SAlexei Starovoitov *
197c8199e2SAlexei Starovoitov * CPU_0 buckets
207c8199e2SAlexei Starovoitov * 16 32 64 96 128 196 256 512 1024 2048 4096
217c8199e2SAlexei Starovoitov * ...
227c8199e2SAlexei Starovoitov * CPU_N buckets
237c8199e2SAlexei Starovoitov * 16 32 64 96 128 196 256 512 1024 2048 4096
247c8199e2SAlexei Starovoitov *
257c8199e2SAlexei Starovoitov * The buckets are prefilled at the start.
267c8199e2SAlexei Starovoitov * BPF programs always run with migration disabled.
277c8199e2SAlexei Starovoitov * It's safe to allocate from cache of the current cpu with irqs disabled.
287c8199e2SAlexei Starovoitov * Free-ing is always done into bucket of the current cpu as well.
297c8199e2SAlexei Starovoitov * irq_work trims extra free elements from buckets with kfree
307c8199e2SAlexei Starovoitov * and refills them with kmalloc, so global kmalloc logic takes care
317c8199e2SAlexei Starovoitov * of freeing objects allocated by one cpu and freed on another.
327c8199e2SAlexei Starovoitov *
337c8199e2SAlexei Starovoitov * Every allocated objected is padded with extra 8 bytes that contains
347c8199e2SAlexei Starovoitov * struct llist_node.
357c8199e2SAlexei Starovoitov */
367c8199e2SAlexei Starovoitov #define LLIST_NODE_SZ sizeof(struct llist_node)
377c8199e2SAlexei Starovoitov
387c8199e2SAlexei Starovoitov /* similar to kmalloc, but sizeof == 8 bucket is gone */
397c8199e2SAlexei Starovoitov static u8 size_index[24] __ro_after_init = {
407c8199e2SAlexei Starovoitov 3, /* 8 */
417c8199e2SAlexei Starovoitov 3, /* 16 */
427c8199e2SAlexei Starovoitov 4, /* 24 */
437c8199e2SAlexei Starovoitov 4, /* 32 */
447c8199e2SAlexei Starovoitov 5, /* 40 */
457c8199e2SAlexei Starovoitov 5, /* 48 */
467c8199e2SAlexei Starovoitov 5, /* 56 */
477c8199e2SAlexei Starovoitov 5, /* 64 */
487c8199e2SAlexei Starovoitov 1, /* 72 */
497c8199e2SAlexei Starovoitov 1, /* 80 */
507c8199e2SAlexei Starovoitov 1, /* 88 */
517c8199e2SAlexei Starovoitov 1, /* 96 */
527c8199e2SAlexei Starovoitov 6, /* 104 */
537c8199e2SAlexei Starovoitov 6, /* 112 */
547c8199e2SAlexei Starovoitov 6, /* 120 */
557c8199e2SAlexei Starovoitov 6, /* 128 */
567c8199e2SAlexei Starovoitov 2, /* 136 */
577c8199e2SAlexei Starovoitov 2, /* 144 */
587c8199e2SAlexei Starovoitov 2, /* 152 */
597c8199e2SAlexei Starovoitov 2, /* 160 */
607c8199e2SAlexei Starovoitov 2, /* 168 */
617c8199e2SAlexei Starovoitov 2, /* 176 */
627c8199e2SAlexei Starovoitov 2, /* 184 */
637c8199e2SAlexei Starovoitov 2 /* 192 */
647c8199e2SAlexei Starovoitov };
657c8199e2SAlexei Starovoitov
bpf_mem_cache_idx(size_t size)667c8199e2SAlexei Starovoitov static int bpf_mem_cache_idx(size_t size)
677c8199e2SAlexei Starovoitov {
687c8199e2SAlexei Starovoitov if (!size || size > 4096)
697c8199e2SAlexei Starovoitov return -1;
707c8199e2SAlexei Starovoitov
717c8199e2SAlexei Starovoitov if (size <= 192)
727c8199e2SAlexei Starovoitov return size_index[(size - 1) / 8] - 1;
737c8199e2SAlexei Starovoitov
7436024d02SHou Tao return fls(size - 1) - 2;
757c8199e2SAlexei Starovoitov }
767c8199e2SAlexei Starovoitov
777c8199e2SAlexei Starovoitov #define NUM_CACHES 11
787c8199e2SAlexei Starovoitov
797c8199e2SAlexei Starovoitov struct bpf_mem_cache {
807c8199e2SAlexei Starovoitov /* per-cpu list of free objects of size 'unit_size'.
817c8199e2SAlexei Starovoitov * All accesses are done with interrupts disabled and 'active' counter
827c8199e2SAlexei Starovoitov * protection with __llist_add() and __llist_del_first().
837c8199e2SAlexei Starovoitov */
847c8199e2SAlexei Starovoitov struct llist_head free_llist;
857c8199e2SAlexei Starovoitov local_t active;
867c8199e2SAlexei Starovoitov
877c8199e2SAlexei Starovoitov /* Operations on the free_list from unit_alloc/unit_free/bpf_mem_refill
887c8199e2SAlexei Starovoitov * are sequenced by per-cpu 'active' counter. But unit_free() cannot
897c8199e2SAlexei Starovoitov * fail. When 'active' is busy the unit_free() will add an object to
907c8199e2SAlexei Starovoitov * free_llist_extra.
917c8199e2SAlexei Starovoitov */
927c8199e2SAlexei Starovoitov struct llist_head free_llist_extra;
937c8199e2SAlexei Starovoitov
947c8199e2SAlexei Starovoitov struct irq_work refill_work;
957c8199e2SAlexei Starovoitov struct obj_cgroup *objcg;
967c8199e2SAlexei Starovoitov int unit_size;
977c8199e2SAlexei Starovoitov /* count of objects in free_llist */
987c8199e2SAlexei Starovoitov int free_cnt;
997c266178SAlexei Starovoitov int low_watermark, high_watermark, batch;
100bfc03c15SAlexei Starovoitov int percpu_size;
101d114dde2SAlexei Starovoitov bool draining;
102822fb26bSAlexei Starovoitov struct bpf_mem_cache *tgt;
1038d5a8011SAlexei Starovoitov
1045af6807bSAlexei Starovoitov /* list of objects to be freed after RCU GP */
1055af6807bSAlexei Starovoitov struct llist_head free_by_rcu;
1065af6807bSAlexei Starovoitov struct llist_node *free_by_rcu_tail;
1075af6807bSAlexei Starovoitov struct llist_head waiting_for_gp;
1085af6807bSAlexei Starovoitov struct llist_node *waiting_for_gp_tail;
1095af6807bSAlexei Starovoitov struct rcu_head rcu;
1105af6807bSAlexei Starovoitov atomic_t call_rcu_in_progress;
1115af6807bSAlexei Starovoitov struct llist_head free_llist_extra_rcu;
1125af6807bSAlexei Starovoitov
11312c8d0f4SAlexei Starovoitov /* list of objects to be freed after RCU tasks trace GP */
11412c8d0f4SAlexei Starovoitov struct llist_head free_by_rcu_ttrace;
11512c8d0f4SAlexei Starovoitov struct llist_head waiting_for_gp_ttrace;
11612c8d0f4SAlexei Starovoitov struct rcu_head rcu_ttrace;
11712c8d0f4SAlexei Starovoitov atomic_t call_rcu_ttrace_in_progress;
1187c8199e2SAlexei Starovoitov };
1197c8199e2SAlexei Starovoitov
1207c8199e2SAlexei Starovoitov struct bpf_mem_caches {
1217c8199e2SAlexei Starovoitov struct bpf_mem_cache cache[NUM_CACHES];
1227c8199e2SAlexei Starovoitov };
1237c8199e2SAlexei Starovoitov
__llist_del_first(struct llist_head * head)1247c8199e2SAlexei Starovoitov static struct llist_node notrace *__llist_del_first(struct llist_head *head)
1257c8199e2SAlexei Starovoitov {
1267c8199e2SAlexei Starovoitov struct llist_node *entry, *next;
1277c8199e2SAlexei Starovoitov
1287c8199e2SAlexei Starovoitov entry = head->first;
1297c8199e2SAlexei Starovoitov if (!entry)
1307c8199e2SAlexei Starovoitov return NULL;
1317c8199e2SAlexei Starovoitov next = entry->next;
1327c8199e2SAlexei Starovoitov head->first = next;
1337c8199e2SAlexei Starovoitov return entry;
1347c8199e2SAlexei Starovoitov }
1357c8199e2SAlexei Starovoitov
__alloc(struct bpf_mem_cache * c,int node,gfp_t flags)136e65a5c6eSMartin KaFai Lau static void *__alloc(struct bpf_mem_cache *c, int node, gfp_t flags)
1377c8199e2SAlexei Starovoitov {
138bfc03c15SAlexei Starovoitov if (c->percpu_size) {
139bfc03c15SAlexei Starovoitov void **obj = kmalloc_node(c->percpu_size, flags, node);
1404ab67149SAlexei Starovoitov void *pptr = __alloc_percpu_gfp(c->unit_size, 8, flags);
1414ab67149SAlexei Starovoitov
1424ab67149SAlexei Starovoitov if (!obj || !pptr) {
1434ab67149SAlexei Starovoitov free_percpu(pptr);
1444ab67149SAlexei Starovoitov kfree(obj);
1454ab67149SAlexei Starovoitov return NULL;
1464ab67149SAlexei Starovoitov }
1474ab67149SAlexei Starovoitov obj[1] = pptr;
1484ab67149SAlexei Starovoitov return obj;
1494ab67149SAlexei Starovoitov }
1504ab67149SAlexei Starovoitov
151997849c4SHou Tao return kmalloc_node(c->unit_size, flags | __GFP_ZERO, node);
1527c8199e2SAlexei Starovoitov }
1537c8199e2SAlexei Starovoitov
get_memcg(const struct bpf_mem_cache * c)1547c8199e2SAlexei Starovoitov static struct mem_cgroup *get_memcg(const struct bpf_mem_cache *c)
1557c8199e2SAlexei Starovoitov {
1567c8199e2SAlexei Starovoitov #ifdef CONFIG_MEMCG_KMEM
1577c8199e2SAlexei Starovoitov if (c->objcg)
1587c8199e2SAlexei Starovoitov return get_mem_cgroup_from_objcg(c->objcg);
1597c8199e2SAlexei Starovoitov #endif
1607c8199e2SAlexei Starovoitov
1617c8199e2SAlexei Starovoitov #ifdef CONFIG_MEMCG
1627c8199e2SAlexei Starovoitov return root_mem_cgroup;
1637c8199e2SAlexei Starovoitov #else
1647c8199e2SAlexei Starovoitov return NULL;
1657c8199e2SAlexei Starovoitov #endif
1667c8199e2SAlexei Starovoitov }
1677c8199e2SAlexei Starovoitov
inc_active(struct bpf_mem_cache * c,unsigned long * flags)16818e027b1SAlexei Starovoitov static void inc_active(struct bpf_mem_cache *c, unsigned long *flags)
16905ae6865SAlexei Starovoitov {
17005ae6865SAlexei Starovoitov if (IS_ENABLED(CONFIG_PREEMPT_RT))
17105ae6865SAlexei Starovoitov /* In RT irq_work runs in per-cpu kthread, so disable
17205ae6865SAlexei Starovoitov * interrupts to avoid preemption and interrupts and
17305ae6865SAlexei Starovoitov * reduce the chance of bpf prog executing on this cpu
17405ae6865SAlexei Starovoitov * when active counter is busy.
17505ae6865SAlexei Starovoitov */
17618e027b1SAlexei Starovoitov local_irq_save(*flags);
17705ae6865SAlexei Starovoitov /* alloc_bulk runs from irq_work which will not preempt a bpf
17805ae6865SAlexei Starovoitov * program that does unit_alloc/unit_free since IRQs are
17905ae6865SAlexei Starovoitov * disabled there. There is no race to increment 'active'
18005ae6865SAlexei Starovoitov * counter. It protects free_llist from corruption in case NMI
18105ae6865SAlexei Starovoitov * bpf prog preempted this loop.
18205ae6865SAlexei Starovoitov */
18305ae6865SAlexei Starovoitov WARN_ON_ONCE(local_inc_return(&c->active) != 1);
18418e027b1SAlexei Starovoitov }
18518e027b1SAlexei Starovoitov
dec_active(struct bpf_mem_cache * c,unsigned long * flags)18663e2da3bSArnd Bergmann static void dec_active(struct bpf_mem_cache *c, unsigned long *flags)
18718e027b1SAlexei Starovoitov {
18805ae6865SAlexei Starovoitov local_dec(&c->active);
18905ae6865SAlexei Starovoitov if (IS_ENABLED(CONFIG_PREEMPT_RT))
19063e2da3bSArnd Bergmann local_irq_restore(*flags);
19105ae6865SAlexei Starovoitov }
19205ae6865SAlexei Starovoitov
add_obj_to_free_list(struct bpf_mem_cache * c,void * obj)19318e027b1SAlexei Starovoitov static void add_obj_to_free_list(struct bpf_mem_cache *c, void *obj)
19418e027b1SAlexei Starovoitov {
19518e027b1SAlexei Starovoitov unsigned long flags;
19618e027b1SAlexei Starovoitov
19718e027b1SAlexei Starovoitov inc_active(c, &flags);
19818e027b1SAlexei Starovoitov __llist_add(obj, &c->free_llist);
19918e027b1SAlexei Starovoitov c->free_cnt++;
20063e2da3bSArnd Bergmann dec_active(c, &flags);
20118e027b1SAlexei Starovoitov }
20218e027b1SAlexei Starovoitov
2037c8199e2SAlexei Starovoitov /* Mostly runs from irq_work except __init phase. */
alloc_bulk(struct bpf_mem_cache * c,int cnt,int node,bool atomic)204d1a02358SYiFei Zhu static void alloc_bulk(struct bpf_mem_cache *c, int cnt, int node, bool atomic)
2057c8199e2SAlexei Starovoitov {
2067c8199e2SAlexei Starovoitov struct mem_cgroup *memcg = NULL, *old_memcg;
207d1a02358SYiFei Zhu gfp_t gfp;
2087c8199e2SAlexei Starovoitov void *obj;
2097c8199e2SAlexei Starovoitov int i;
2107c8199e2SAlexei Starovoitov
211d1a02358SYiFei Zhu gfp = __GFP_NOWARN | __GFP_ACCOUNT;
212d1a02358SYiFei Zhu gfp |= atomic ? GFP_NOWAIT : GFP_KERNEL;
213d1a02358SYiFei Zhu
2147c8199e2SAlexei Starovoitov for (i = 0; i < cnt; i++) {
2150893d600SHou Tao /*
216822fb26bSAlexei Starovoitov * For every 'c' llist_del_first(&c->free_by_rcu_ttrace); is
217822fb26bSAlexei Starovoitov * done only by one CPU == current CPU. Other CPUs might
218822fb26bSAlexei Starovoitov * llist_add() and llist_del_all() in parallel.
2190893d600SHou Tao */
220822fb26bSAlexei Starovoitov obj = llist_del_first(&c->free_by_rcu_ttrace);
22174680482SAlexei Starovoitov if (!obj)
22274680482SAlexei Starovoitov break;
22374680482SAlexei Starovoitov add_obj_to_free_list(c, obj);
22474680482SAlexei Starovoitov }
22574680482SAlexei Starovoitov if (i >= cnt)
22674680482SAlexei Starovoitov return;
22774680482SAlexei Starovoitov
22804fabf00SAlexei Starovoitov for (; i < cnt; i++) {
22904fabf00SAlexei Starovoitov obj = llist_del_first(&c->waiting_for_gp_ttrace);
23004fabf00SAlexei Starovoitov if (!obj)
23104fabf00SAlexei Starovoitov break;
23204fabf00SAlexei Starovoitov add_obj_to_free_list(c, obj);
23304fabf00SAlexei Starovoitov }
23404fabf00SAlexei Starovoitov if (i >= cnt)
23504fabf00SAlexei Starovoitov return;
23604fabf00SAlexei Starovoitov
23774680482SAlexei Starovoitov memcg = get_memcg(c);
23874680482SAlexei Starovoitov old_memcg = set_active_memcg(memcg);
23974680482SAlexei Starovoitov for (; i < cnt; i++) {
240e65a5c6eSMartin KaFai Lau /* Allocate, but don't deplete atomic reserves that typical
241e65a5c6eSMartin KaFai Lau * GFP_ATOMIC would do. irq_work runs on this cpu and kmalloc
242e65a5c6eSMartin KaFai Lau * will allocate from the current numa node which is what we
243e65a5c6eSMartin KaFai Lau * want here.
244e65a5c6eSMartin KaFai Lau */
245d1a02358SYiFei Zhu obj = __alloc(c, node, gfp);
2467c8199e2SAlexei Starovoitov if (!obj)
2477c8199e2SAlexei Starovoitov break;
24805ae6865SAlexei Starovoitov add_obj_to_free_list(c, obj);
2497c8199e2SAlexei Starovoitov }
2507c8199e2SAlexei Starovoitov set_active_memcg(old_memcg);
2517c8199e2SAlexei Starovoitov mem_cgroup_put(memcg);
2527c8199e2SAlexei Starovoitov }
2537c8199e2SAlexei Starovoitov
free_one(void * obj,bool percpu)254aa7881fcSHou Tao static void free_one(void *obj, bool percpu)
2557c8199e2SAlexei Starovoitov {
256aa7881fcSHou Tao if (percpu) {
2574ab67149SAlexei Starovoitov free_percpu(((void **)obj)[1]);
258bfc03c15SAlexei Starovoitov kfree(obj);
2594ab67149SAlexei Starovoitov return;
2604ab67149SAlexei Starovoitov }
2614ab67149SAlexei Starovoitov
2627c8199e2SAlexei Starovoitov kfree(obj);
2637c8199e2SAlexei Starovoitov }
2647c8199e2SAlexei Starovoitov
free_all(struct llist_node * llnode,bool percpu)2659de3e815SAlexei Starovoitov static int free_all(struct llist_node *llnode, bool percpu)
2668d5a8011SAlexei Starovoitov {
2678d5a8011SAlexei Starovoitov struct llist_node *pos, *t;
2689de3e815SAlexei Starovoitov int cnt = 0;
2698d5a8011SAlexei Starovoitov
2709de3e815SAlexei Starovoitov llist_for_each_safe(pos, t, llnode) {
271aa7881fcSHou Tao free_one(pos, percpu);
2729de3e815SAlexei Starovoitov cnt++;
2739de3e815SAlexei Starovoitov }
2749de3e815SAlexei Starovoitov return cnt;
275aa7881fcSHou Tao }
276aa7881fcSHou Tao
__free_rcu(struct rcu_head * head)277aa7881fcSHou Tao static void __free_rcu(struct rcu_head *head)
278aa7881fcSHou Tao {
27912c8d0f4SAlexei Starovoitov struct bpf_mem_cache *c = container_of(head, struct bpf_mem_cache, rcu_ttrace);
280aa7881fcSHou Tao
28112c8d0f4SAlexei Starovoitov free_all(llist_del_all(&c->waiting_for_gp_ttrace), !!c->percpu_size);
28212c8d0f4SAlexei Starovoitov atomic_set(&c->call_rcu_ttrace_in_progress, 0);
2838d5a8011SAlexei Starovoitov }
2848d5a8011SAlexei Starovoitov
__free_rcu_tasks_trace(struct rcu_head * head)285dccb4a90SAlexei Starovoitov static void __free_rcu_tasks_trace(struct rcu_head *head)
286dccb4a90SAlexei Starovoitov {
28759be91e5SHou Tao /* If RCU Tasks Trace grace period implies RCU grace period,
28859be91e5SHou Tao * there is no need to invoke call_rcu().
28959be91e5SHou Tao */
29059be91e5SHou Tao if (rcu_trace_implies_rcu_gp())
29159be91e5SHou Tao __free_rcu(head);
29259be91e5SHou Tao else
29359be91e5SHou Tao call_rcu(head, __free_rcu);
294dccb4a90SAlexei Starovoitov }
295dccb4a90SAlexei Starovoitov
enque_to_free(struct bpf_mem_cache * c,void * obj)2968d5a8011SAlexei Starovoitov static void enque_to_free(struct bpf_mem_cache *c, void *obj)
2978d5a8011SAlexei Starovoitov {
2988d5a8011SAlexei Starovoitov struct llist_node *llnode = obj;
2998d5a8011SAlexei Starovoitov
3008d5a8011SAlexei Starovoitov /* bpf_mem_cache is a per-cpu object. Freeing happens in irq_work.
30112c8d0f4SAlexei Starovoitov * Nothing races to add to free_by_rcu_ttrace list.
3028d5a8011SAlexei Starovoitov */
303822fb26bSAlexei Starovoitov llist_add(llnode, &c->free_by_rcu_ttrace);
3048d5a8011SAlexei Starovoitov }
3058d5a8011SAlexei Starovoitov
do_call_rcu_ttrace(struct bpf_mem_cache * c)30612c8d0f4SAlexei Starovoitov static void do_call_rcu_ttrace(struct bpf_mem_cache *c)
3078d5a8011SAlexei Starovoitov {
3088d5a8011SAlexei Starovoitov struct llist_node *llnode, *t;
3098d5a8011SAlexei Starovoitov
310822fb26bSAlexei Starovoitov if (atomic_xchg(&c->call_rcu_ttrace_in_progress, 1)) {
311822fb26bSAlexei Starovoitov if (unlikely(READ_ONCE(c->draining))) {
312822fb26bSAlexei Starovoitov llnode = llist_del_all(&c->free_by_rcu_ttrace);
313822fb26bSAlexei Starovoitov free_all(llnode, !!c->percpu_size);
314822fb26bSAlexei Starovoitov }
3158d5a8011SAlexei Starovoitov return;
316822fb26bSAlexei Starovoitov }
3178d5a8011SAlexei Starovoitov
31812c8d0f4SAlexei Starovoitov WARN_ON_ONCE(!llist_empty(&c->waiting_for_gp_ttrace));
319822fb26bSAlexei Starovoitov llist_for_each_safe(llnode, t, llist_del_all(&c->free_by_rcu_ttrace))
32004fabf00SAlexei Starovoitov llist_add(llnode, &c->waiting_for_gp_ttrace);
321d114dde2SAlexei Starovoitov
322d114dde2SAlexei Starovoitov if (unlikely(READ_ONCE(c->draining))) {
323d114dde2SAlexei Starovoitov __free_rcu(&c->rcu_ttrace);
324d114dde2SAlexei Starovoitov return;
325d114dde2SAlexei Starovoitov }
326d114dde2SAlexei Starovoitov
327dccb4a90SAlexei Starovoitov /* Use call_rcu_tasks_trace() to wait for sleepable progs to finish.
32859be91e5SHou Tao * If RCU Tasks Trace grace period implies RCU grace period, free
32959be91e5SHou Tao * these elements directly, else use call_rcu() to wait for normal
33059be91e5SHou Tao * progs to finish and finally do free_one() on each element.
331dccb4a90SAlexei Starovoitov */
33212c8d0f4SAlexei Starovoitov call_rcu_tasks_trace(&c->rcu_ttrace, __free_rcu_tasks_trace);
3338d5a8011SAlexei Starovoitov }
3348d5a8011SAlexei Starovoitov
free_bulk(struct bpf_mem_cache * c)3357c8199e2SAlexei Starovoitov static void free_bulk(struct bpf_mem_cache *c)
3367c8199e2SAlexei Starovoitov {
337822fb26bSAlexei Starovoitov struct bpf_mem_cache *tgt = c->tgt;
3387c8199e2SAlexei Starovoitov struct llist_node *llnode, *t;
3397c8199e2SAlexei Starovoitov unsigned long flags;
3407c8199e2SAlexei Starovoitov int cnt;
3417c8199e2SAlexei Starovoitov
342822fb26bSAlexei Starovoitov WARN_ON_ONCE(tgt->unit_size != c->unit_size);
343822fb26bSAlexei Starovoitov
3447c8199e2SAlexei Starovoitov do {
34518e027b1SAlexei Starovoitov inc_active(c, &flags);
3467c8199e2SAlexei Starovoitov llnode = __llist_del_first(&c->free_llist);
3477c8199e2SAlexei Starovoitov if (llnode)
3487c8199e2SAlexei Starovoitov cnt = --c->free_cnt;
3497c8199e2SAlexei Starovoitov else
3507c8199e2SAlexei Starovoitov cnt = 0;
35163e2da3bSArnd Bergmann dec_active(c, &flags);
352c31b38cbSHou Tao if (llnode)
353822fb26bSAlexei Starovoitov enque_to_free(tgt, llnode);
3547c266178SAlexei Starovoitov } while (cnt > (c->high_watermark + c->low_watermark) / 2);
3557c8199e2SAlexei Starovoitov
3567c8199e2SAlexei Starovoitov /* and drain free_llist_extra */
3577c8199e2SAlexei Starovoitov llist_for_each_safe(llnode, t, llist_del_all(&c->free_llist_extra))
358822fb26bSAlexei Starovoitov enque_to_free(tgt, llnode);
359822fb26bSAlexei Starovoitov do_call_rcu_ttrace(tgt);
3607c8199e2SAlexei Starovoitov }
3617c8199e2SAlexei Starovoitov
__free_by_rcu(struct rcu_head * head)3625af6807bSAlexei Starovoitov static void __free_by_rcu(struct rcu_head *head)
3635af6807bSAlexei Starovoitov {
3645af6807bSAlexei Starovoitov struct bpf_mem_cache *c = container_of(head, struct bpf_mem_cache, rcu);
3655af6807bSAlexei Starovoitov struct bpf_mem_cache *tgt = c->tgt;
3665af6807bSAlexei Starovoitov struct llist_node *llnode;
3675af6807bSAlexei Starovoitov
3685af6807bSAlexei Starovoitov llnode = llist_del_all(&c->waiting_for_gp);
3695af6807bSAlexei Starovoitov if (!llnode)
3705af6807bSAlexei Starovoitov goto out;
3715af6807bSAlexei Starovoitov
3725af6807bSAlexei Starovoitov llist_add_batch(llnode, c->waiting_for_gp_tail, &tgt->free_by_rcu_ttrace);
3735af6807bSAlexei Starovoitov
3745af6807bSAlexei Starovoitov /* Objects went through regular RCU GP. Send them to RCU tasks trace */
3755af6807bSAlexei Starovoitov do_call_rcu_ttrace(tgt);
3765af6807bSAlexei Starovoitov out:
3775af6807bSAlexei Starovoitov atomic_set(&c->call_rcu_in_progress, 0);
3785af6807bSAlexei Starovoitov }
3795af6807bSAlexei Starovoitov
check_free_by_rcu(struct bpf_mem_cache * c)3805af6807bSAlexei Starovoitov static void check_free_by_rcu(struct bpf_mem_cache *c)
3815af6807bSAlexei Starovoitov {
3825af6807bSAlexei Starovoitov struct llist_node *llnode, *t;
3835af6807bSAlexei Starovoitov unsigned long flags;
3845af6807bSAlexei Starovoitov
3855af6807bSAlexei Starovoitov /* drain free_llist_extra_rcu */
3865af6807bSAlexei Starovoitov if (unlikely(!llist_empty(&c->free_llist_extra_rcu))) {
3875af6807bSAlexei Starovoitov inc_active(c, &flags);
3885af6807bSAlexei Starovoitov llist_for_each_safe(llnode, t, llist_del_all(&c->free_llist_extra_rcu))
3895af6807bSAlexei Starovoitov if (__llist_add(llnode, &c->free_by_rcu))
3905af6807bSAlexei Starovoitov c->free_by_rcu_tail = llnode;
39163e2da3bSArnd Bergmann dec_active(c, &flags);
3925af6807bSAlexei Starovoitov }
3935af6807bSAlexei Starovoitov
3945af6807bSAlexei Starovoitov if (llist_empty(&c->free_by_rcu))
3955af6807bSAlexei Starovoitov return;
3965af6807bSAlexei Starovoitov
3975af6807bSAlexei Starovoitov if (atomic_xchg(&c->call_rcu_in_progress, 1)) {
3985af6807bSAlexei Starovoitov /*
3995af6807bSAlexei Starovoitov * Instead of kmalloc-ing new rcu_head and triggering 10k
4005af6807bSAlexei Starovoitov * call_rcu() to hit rcutree.qhimark and force RCU to notice
4015af6807bSAlexei Starovoitov * the overload just ask RCU to hurry up. There could be many
4025af6807bSAlexei Starovoitov * objects in free_by_rcu list.
4035af6807bSAlexei Starovoitov * This hint reduces memory consumption for an artificial
4045af6807bSAlexei Starovoitov * benchmark from 2 Gbyte to 150 Mbyte.
4055af6807bSAlexei Starovoitov */
4065af6807bSAlexei Starovoitov rcu_request_urgent_qs_task(current);
4075af6807bSAlexei Starovoitov return;
4085af6807bSAlexei Starovoitov }
4095af6807bSAlexei Starovoitov
4105af6807bSAlexei Starovoitov WARN_ON_ONCE(!llist_empty(&c->waiting_for_gp));
4115af6807bSAlexei Starovoitov
4125af6807bSAlexei Starovoitov inc_active(c, &flags);
4135af6807bSAlexei Starovoitov WRITE_ONCE(c->waiting_for_gp.first, __llist_del_all(&c->free_by_rcu));
4145af6807bSAlexei Starovoitov c->waiting_for_gp_tail = c->free_by_rcu_tail;
41563e2da3bSArnd Bergmann dec_active(c, &flags);
4165af6807bSAlexei Starovoitov
4175af6807bSAlexei Starovoitov if (unlikely(READ_ONCE(c->draining))) {
4185af6807bSAlexei Starovoitov free_all(llist_del_all(&c->waiting_for_gp), !!c->percpu_size);
4195af6807bSAlexei Starovoitov atomic_set(&c->call_rcu_in_progress, 0);
4205af6807bSAlexei Starovoitov } else {
4215af6807bSAlexei Starovoitov call_rcu_hurry(&c->rcu, __free_by_rcu);
4225af6807bSAlexei Starovoitov }
4235af6807bSAlexei Starovoitov }
4245af6807bSAlexei Starovoitov
bpf_mem_refill(struct irq_work * work)4257c8199e2SAlexei Starovoitov static void bpf_mem_refill(struct irq_work *work)
4267c8199e2SAlexei Starovoitov {
4277c8199e2SAlexei Starovoitov struct bpf_mem_cache *c = container_of(work, struct bpf_mem_cache, refill_work);
4287c8199e2SAlexei Starovoitov int cnt;
4297c8199e2SAlexei Starovoitov
4307c8199e2SAlexei Starovoitov /* Racy access to free_cnt. It doesn't need to be 100% accurate */
4317c8199e2SAlexei Starovoitov cnt = c->free_cnt;
4327c266178SAlexei Starovoitov if (cnt < c->low_watermark)
4337c8199e2SAlexei Starovoitov /* irq_work runs on this cpu and kmalloc will allocate
4347c8199e2SAlexei Starovoitov * from the current numa node which is what we want here.
4357c8199e2SAlexei Starovoitov */
436d1a02358SYiFei Zhu alloc_bulk(c, c->batch, NUMA_NO_NODE, true);
4377c266178SAlexei Starovoitov else if (cnt > c->high_watermark)
4387c8199e2SAlexei Starovoitov free_bulk(c);
4395af6807bSAlexei Starovoitov
4405af6807bSAlexei Starovoitov check_free_by_rcu(c);
4417c8199e2SAlexei Starovoitov }
4427c8199e2SAlexei Starovoitov
irq_work_raise(struct bpf_mem_cache * c)4437c8199e2SAlexei Starovoitov static void notrace irq_work_raise(struct bpf_mem_cache *c)
4447c8199e2SAlexei Starovoitov {
4457c8199e2SAlexei Starovoitov irq_work_queue(&c->refill_work);
4467c8199e2SAlexei Starovoitov }
4477c8199e2SAlexei Starovoitov
4487c266178SAlexei Starovoitov /* For typical bpf map case that uses bpf_mem_cache_alloc and single bucket
4497c266178SAlexei Starovoitov * the freelist cache will be elem_size * 64 (or less) on each cpu.
4507c266178SAlexei Starovoitov *
4517c266178SAlexei Starovoitov * For bpf programs that don't have statically known allocation sizes and
4527c266178SAlexei Starovoitov * assuming (low_mark + high_mark) / 2 as an average number of elements per
4537c266178SAlexei Starovoitov * bucket and all buckets are used the total amount of memory in freelists
4547c266178SAlexei Starovoitov * on each cpu will be:
4557c266178SAlexei Starovoitov * 64*16 + 64*32 + 64*64 + 64*96 + 64*128 + 64*196 + 64*256 + 32*512 + 16*1024 + 8*2048 + 4*4096
4567c266178SAlexei Starovoitov * == ~ 116 Kbyte using below heuristic.
4577c266178SAlexei Starovoitov * Initialized, but unused bpf allocator (not bpf map specific one) will
4587c266178SAlexei Starovoitov * consume ~ 11 Kbyte per cpu.
4597c266178SAlexei Starovoitov * Typical case will be between 11K and 116K closer to 11K.
4607c266178SAlexei Starovoitov * bpf progs can and should share bpf_mem_cache when possible.
4617c266178SAlexei Starovoitov */
init_refill_work(struct bpf_mem_cache * c)462b1d53958SHou Tao static void init_refill_work(struct bpf_mem_cache *c)
4637c8199e2SAlexei Starovoitov {
4647c8199e2SAlexei Starovoitov init_irq_work(&c->refill_work, bpf_mem_refill);
4657c266178SAlexei Starovoitov if (c->unit_size <= 256) {
4667c266178SAlexei Starovoitov c->low_watermark = 32;
4677c266178SAlexei Starovoitov c->high_watermark = 96;
4687c266178SAlexei Starovoitov } else {
4697c266178SAlexei Starovoitov /* When page_size == 4k, order-0 cache will have low_mark == 2
4707c266178SAlexei Starovoitov * and high_mark == 6 with batch alloc of 3 individual pages at
4717c266178SAlexei Starovoitov * a time.
4727c266178SAlexei Starovoitov * 8k allocs and above low == 1, high == 3, batch == 1.
4737c266178SAlexei Starovoitov */
4747c266178SAlexei Starovoitov c->low_watermark = max(32 * 256 / c->unit_size, 1);
4757c266178SAlexei Starovoitov c->high_watermark = max(96 * 256 / c->unit_size, 3);
4767c266178SAlexei Starovoitov }
4777c266178SAlexei Starovoitov c->batch = max((c->high_watermark - c->low_watermark) / 4 * 3, 1);
478b1d53958SHou Tao }
4797c266178SAlexei Starovoitov
prefill_mem_cache(struct bpf_mem_cache * c,int cpu)480b1d53958SHou Tao static void prefill_mem_cache(struct bpf_mem_cache *c, int cpu)
481b1d53958SHou Tao {
4827c8199e2SAlexei Starovoitov /* To avoid consuming memory assume that 1st run of bpf
4837c8199e2SAlexei Starovoitov * prog won't be doing more than 4 map_update_elem from
4847c8199e2SAlexei Starovoitov * irq disabled region
4857c8199e2SAlexei Starovoitov */
486d1a02358SYiFei Zhu alloc_bulk(c, c->unit_size <= 256 ? 4 : 1, cpu_to_node(cpu), false);
4877c8199e2SAlexei Starovoitov }
4887c8199e2SAlexei Starovoitov
489bfc03c15SAlexei Starovoitov /* When size != 0 bpf_mem_cache for each cpu.
4907c8199e2SAlexei Starovoitov * This is typical bpf hash map use case when all elements have equal size.
4917c8199e2SAlexei Starovoitov *
4927c8199e2SAlexei Starovoitov * When size == 0 allocate 11 bpf_mem_cache-s for each cpu, then rely on
4937c8199e2SAlexei Starovoitov * kmalloc/kfree. Max allocation size is 4096 in this case.
4947c8199e2SAlexei Starovoitov * This is bpf_dynptr and bpf_kptr use case.
4957c8199e2SAlexei Starovoitov */
bpf_mem_alloc_init(struct bpf_mem_alloc * ma,int size,bool percpu)4964ab67149SAlexei Starovoitov int bpf_mem_alloc_init(struct bpf_mem_alloc *ma, int size, bool percpu)
4977c8199e2SAlexei Starovoitov {
4987c8199e2SAlexei Starovoitov static u16 sizes[NUM_CACHES] = {96, 192, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096};
4997c8199e2SAlexei Starovoitov struct bpf_mem_caches *cc, __percpu *pcc;
5007c8199e2SAlexei Starovoitov struct bpf_mem_cache *c, __percpu *pc;
5017c8199e2SAlexei Starovoitov struct obj_cgroup *objcg = NULL;
502*3bc29c78SHou Tao int cpu, i, unit_size, percpu_size = 0;
5037c8199e2SAlexei Starovoitov
50463ddf081SHou Tao ma->percpu = percpu;
50563ddf081SHou Tao
5067c8199e2SAlexei Starovoitov if (size) {
5077c8199e2SAlexei Starovoitov pc = __alloc_percpu_gfp(sizeof(*pc), 8, GFP_KERNEL);
5087c8199e2SAlexei Starovoitov if (!pc)
5097c8199e2SAlexei Starovoitov return -ENOMEM;
5104ab67149SAlexei Starovoitov
511bfc03c15SAlexei Starovoitov if (percpu)
5124ab67149SAlexei Starovoitov /* room for llist_node and per-cpu pointer */
513bfc03c15SAlexei Starovoitov percpu_size = LLIST_NODE_SZ + sizeof(void *);
514bfc03c15SAlexei Starovoitov else
5157c8199e2SAlexei Starovoitov size += LLIST_NODE_SZ; /* room for llist_node */
5164ab67149SAlexei Starovoitov unit_size = size;
5174ab67149SAlexei Starovoitov
5187c8199e2SAlexei Starovoitov #ifdef CONFIG_MEMCG_KMEM
519ee53cbfbSYafang Shao if (memcg_bpf_enabled())
5207c8199e2SAlexei Starovoitov objcg = get_obj_cgroup_from_current();
5217c8199e2SAlexei Starovoitov #endif
5227c8199e2SAlexei Starovoitov for_each_possible_cpu(cpu) {
5237c8199e2SAlexei Starovoitov c = per_cpu_ptr(pc, cpu);
5244ab67149SAlexei Starovoitov c->unit_size = unit_size;
5257c8199e2SAlexei Starovoitov c->objcg = objcg;
526bfc03c15SAlexei Starovoitov c->percpu_size = percpu_size;
527822fb26bSAlexei Starovoitov c->tgt = c;
528b1d53958SHou Tao init_refill_work(c);
5297c8199e2SAlexei Starovoitov prefill_mem_cache(c, cpu);
5307c8199e2SAlexei Starovoitov }
5317c8199e2SAlexei Starovoitov ma->cache = pc;
5327c8199e2SAlexei Starovoitov return 0;
5337c8199e2SAlexei Starovoitov }
5347c8199e2SAlexei Starovoitov
5354ab67149SAlexei Starovoitov /* size == 0 && percpu is an invalid combination */
5364ab67149SAlexei Starovoitov if (WARN_ON_ONCE(percpu))
5374ab67149SAlexei Starovoitov return -EINVAL;
5384ab67149SAlexei Starovoitov
5397c8199e2SAlexei Starovoitov pcc = __alloc_percpu_gfp(sizeof(*cc), 8, GFP_KERNEL);
5407c8199e2SAlexei Starovoitov if (!pcc)
5417c8199e2SAlexei Starovoitov return -ENOMEM;
5427c8199e2SAlexei Starovoitov #ifdef CONFIG_MEMCG_KMEM
5437c8199e2SAlexei Starovoitov objcg = get_obj_cgroup_from_current();
5447c8199e2SAlexei Starovoitov #endif
5457c8199e2SAlexei Starovoitov for_each_possible_cpu(cpu) {
5467c8199e2SAlexei Starovoitov cc = per_cpu_ptr(pcc, cpu);
5477c8199e2SAlexei Starovoitov for (i = 0; i < NUM_CACHES; i++) {
5487c8199e2SAlexei Starovoitov c = &cc->cache[i];
5497c8199e2SAlexei Starovoitov c->unit_size = sizes[i];
5507c8199e2SAlexei Starovoitov c->objcg = objcg;
551822fb26bSAlexei Starovoitov c->tgt = c;
552b1d53958SHou Tao
553b1d53958SHou Tao init_refill_work(c);
5547c8199e2SAlexei Starovoitov prefill_mem_cache(c, cpu);
5557c8199e2SAlexei Starovoitov }
5567c8199e2SAlexei Starovoitov }
557c9304725SHou Tao
5587c8199e2SAlexei Starovoitov ma->caches = pcc;
559*3bc29c78SHou Tao return 0;
5607c8199e2SAlexei Starovoitov }
5617c8199e2SAlexei Starovoitov
drain_mem_cache(struct bpf_mem_cache * c)5627c8199e2SAlexei Starovoitov static void drain_mem_cache(struct bpf_mem_cache *c)
5637c8199e2SAlexei Starovoitov {
564aa7881fcSHou Tao bool percpu = !!c->percpu_size;
5657c8199e2SAlexei Starovoitov
5669f2c6e96SAlexei Starovoitov /* No progs are using this bpf_mem_cache, but htab_map_free() called
5679f2c6e96SAlexei Starovoitov * bpf_mem_cache_free() for all remaining elements and they can be in
56812c8d0f4SAlexei Starovoitov * free_by_rcu_ttrace or in waiting_for_gp_ttrace lists, so drain those lists now.
569fa4447cbSHou Tao *
57012c8d0f4SAlexei Starovoitov * Except for waiting_for_gp_ttrace list, there are no concurrent operations
571fa4447cbSHou Tao * on these lists, so it is safe to use __llist_del_all().
5728d5a8011SAlexei Starovoitov */
573822fb26bSAlexei Starovoitov free_all(llist_del_all(&c->free_by_rcu_ttrace), percpu);
57412c8d0f4SAlexei Starovoitov free_all(llist_del_all(&c->waiting_for_gp_ttrace), percpu);
575aa7881fcSHou Tao free_all(__llist_del_all(&c->free_llist), percpu);
576aa7881fcSHou Tao free_all(__llist_del_all(&c->free_llist_extra), percpu);
5775af6807bSAlexei Starovoitov free_all(__llist_del_all(&c->free_by_rcu), percpu);
5785af6807bSAlexei Starovoitov free_all(__llist_del_all(&c->free_llist_extra_rcu), percpu);
5795af6807bSAlexei Starovoitov free_all(llist_del_all(&c->waiting_for_gp), percpu);
5807c8199e2SAlexei Starovoitov }
5817c8199e2SAlexei Starovoitov
check_mem_cache(struct bpf_mem_cache * c)5824ed8b5bcSHou Tao static void check_mem_cache(struct bpf_mem_cache *c)
5834ed8b5bcSHou Tao {
5844ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->free_by_rcu_ttrace));
5854ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->waiting_for_gp_ttrace));
5864ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->free_llist));
5874ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->free_llist_extra));
5884ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->free_by_rcu));
5894ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->free_llist_extra_rcu));
5904ed8b5bcSHou Tao WARN_ON_ONCE(!llist_empty(&c->waiting_for_gp));
5914ed8b5bcSHou Tao }
5924ed8b5bcSHou Tao
check_leaked_objs(struct bpf_mem_alloc * ma)5934ed8b5bcSHou Tao static void check_leaked_objs(struct bpf_mem_alloc *ma)
5944ed8b5bcSHou Tao {
5954ed8b5bcSHou Tao struct bpf_mem_caches *cc;
5964ed8b5bcSHou Tao struct bpf_mem_cache *c;
5974ed8b5bcSHou Tao int cpu, i;
5984ed8b5bcSHou Tao
5994ed8b5bcSHou Tao if (ma->cache) {
6004ed8b5bcSHou Tao for_each_possible_cpu(cpu) {
6014ed8b5bcSHou Tao c = per_cpu_ptr(ma->cache, cpu);
6024ed8b5bcSHou Tao check_mem_cache(c);
6034ed8b5bcSHou Tao }
6044ed8b5bcSHou Tao }
6054ed8b5bcSHou Tao if (ma->caches) {
6064ed8b5bcSHou Tao for_each_possible_cpu(cpu) {
6074ed8b5bcSHou Tao cc = per_cpu_ptr(ma->caches, cpu);
6084ed8b5bcSHou Tao for (i = 0; i < NUM_CACHES; i++) {
6094ed8b5bcSHou Tao c = &cc->cache[i];
6104ed8b5bcSHou Tao check_mem_cache(c);
6114ed8b5bcSHou Tao }
6124ed8b5bcSHou Tao }
6134ed8b5bcSHou Tao }
6144ed8b5bcSHou Tao }
6154ed8b5bcSHou Tao
free_mem_alloc_no_barrier(struct bpf_mem_alloc * ma)6169f2c6e96SAlexei Starovoitov static void free_mem_alloc_no_barrier(struct bpf_mem_alloc *ma)
6179f2c6e96SAlexei Starovoitov {
6184ed8b5bcSHou Tao check_leaked_objs(ma);
6199f2c6e96SAlexei Starovoitov free_percpu(ma->cache);
6209f2c6e96SAlexei Starovoitov free_percpu(ma->caches);
6219f2c6e96SAlexei Starovoitov ma->cache = NULL;
6229f2c6e96SAlexei Starovoitov ma->caches = NULL;
6239f2c6e96SAlexei Starovoitov }
6249f2c6e96SAlexei Starovoitov
free_mem_alloc(struct bpf_mem_alloc * ma)6259f2c6e96SAlexei Starovoitov static void free_mem_alloc(struct bpf_mem_alloc *ma)
6269f2c6e96SAlexei Starovoitov {
6275af6807bSAlexei Starovoitov /* waiting_for_gp[_ttrace] lists were drained, but RCU callbacks
6285af6807bSAlexei Starovoitov * might still execute. Wait for them.
629822ed78fSHou Tao *
630822ed78fSHou Tao * rcu_barrier_tasks_trace() doesn't imply synchronize_rcu_tasks_trace(),
631822ed78fSHou Tao * but rcu_barrier_tasks_trace() and rcu_barrier() below are only used
632822ed78fSHou Tao * to wait for the pending __free_rcu_tasks_trace() and __free_rcu(),
633822ed78fSHou Tao * so if call_rcu(head, __free_rcu) is skipped due to
634822ed78fSHou Tao * rcu_trace_implies_rcu_gp(), it will be OK to skip rcu_barrier() by
635822ed78fSHou Tao * using rcu_trace_implies_rcu_gp() as well.
6369f2c6e96SAlexei Starovoitov */
6375af6807bSAlexei Starovoitov rcu_barrier(); /* wait for __free_by_rcu */
6385af6807bSAlexei Starovoitov rcu_barrier_tasks_trace(); /* wait for __free_rcu */
639822ed78fSHou Tao if (!rcu_trace_implies_rcu_gp())
6409f2c6e96SAlexei Starovoitov rcu_barrier();
6419f2c6e96SAlexei Starovoitov free_mem_alloc_no_barrier(ma);
6429f2c6e96SAlexei Starovoitov }
6439f2c6e96SAlexei Starovoitov
free_mem_alloc_deferred(struct work_struct * work)6449f2c6e96SAlexei Starovoitov static void free_mem_alloc_deferred(struct work_struct *work)
6459f2c6e96SAlexei Starovoitov {
6469f2c6e96SAlexei Starovoitov struct bpf_mem_alloc *ma = container_of(work, struct bpf_mem_alloc, work);
6479f2c6e96SAlexei Starovoitov
6489f2c6e96SAlexei Starovoitov free_mem_alloc(ma);
6499f2c6e96SAlexei Starovoitov kfree(ma);
6509f2c6e96SAlexei Starovoitov }
6519f2c6e96SAlexei Starovoitov
destroy_mem_alloc(struct bpf_mem_alloc * ma,int rcu_in_progress)6529f2c6e96SAlexei Starovoitov static void destroy_mem_alloc(struct bpf_mem_alloc *ma, int rcu_in_progress)
6539f2c6e96SAlexei Starovoitov {
6549f2c6e96SAlexei Starovoitov struct bpf_mem_alloc *copy;
6559f2c6e96SAlexei Starovoitov
6569f2c6e96SAlexei Starovoitov if (!rcu_in_progress) {
6579f2c6e96SAlexei Starovoitov /* Fast path. No callbacks are pending, hence no need to do
6589f2c6e96SAlexei Starovoitov * rcu_barrier-s.
6599f2c6e96SAlexei Starovoitov */
6609f2c6e96SAlexei Starovoitov free_mem_alloc_no_barrier(ma);
6619f2c6e96SAlexei Starovoitov return;
6629f2c6e96SAlexei Starovoitov }
6639f2c6e96SAlexei Starovoitov
664a80672d7SAlexei Starovoitov copy = kmemdup(ma, sizeof(*ma), GFP_KERNEL);
6659f2c6e96SAlexei Starovoitov if (!copy) {
6669f2c6e96SAlexei Starovoitov /* Slow path with inline barrier-s */
6679f2c6e96SAlexei Starovoitov free_mem_alloc(ma);
6689f2c6e96SAlexei Starovoitov return;
6699f2c6e96SAlexei Starovoitov }
6709f2c6e96SAlexei Starovoitov
6719f2c6e96SAlexei Starovoitov /* Defer barriers into worker to let the rest of map memory to be freed */
672a80672d7SAlexei Starovoitov memset(ma, 0, sizeof(*ma));
6739f2c6e96SAlexei Starovoitov INIT_WORK(©->work, free_mem_alloc_deferred);
6749f2c6e96SAlexei Starovoitov queue_work(system_unbound_wq, ©->work);
6759f2c6e96SAlexei Starovoitov }
6769f2c6e96SAlexei Starovoitov
bpf_mem_alloc_destroy(struct bpf_mem_alloc * ma)6777c8199e2SAlexei Starovoitov void bpf_mem_alloc_destroy(struct bpf_mem_alloc *ma)
6787c8199e2SAlexei Starovoitov {
6797c8199e2SAlexei Starovoitov struct bpf_mem_caches *cc;
6807c8199e2SAlexei Starovoitov struct bpf_mem_cache *c;
6819f2c6e96SAlexei Starovoitov int cpu, i, rcu_in_progress;
6827c8199e2SAlexei Starovoitov
6837c8199e2SAlexei Starovoitov if (ma->cache) {
6849f2c6e96SAlexei Starovoitov rcu_in_progress = 0;
6857c8199e2SAlexei Starovoitov for_each_possible_cpu(cpu) {
6867c8199e2SAlexei Starovoitov c = per_cpu_ptr(ma->cache, cpu);
687d114dde2SAlexei Starovoitov WRITE_ONCE(c->draining, true);
6883d058187SHou Tao irq_work_sync(&c->refill_work);
6897c8199e2SAlexei Starovoitov drain_mem_cache(c);
69012c8d0f4SAlexei Starovoitov rcu_in_progress += atomic_read(&c->call_rcu_ttrace_in_progress);
6915af6807bSAlexei Starovoitov rcu_in_progress += atomic_read(&c->call_rcu_in_progress);
6927c8199e2SAlexei Starovoitov }
693bfc03c15SAlexei Starovoitov /* objcg is the same across cpus */
6947c8199e2SAlexei Starovoitov if (c->objcg)
6957c8199e2SAlexei Starovoitov obj_cgroup_put(c->objcg);
6969f2c6e96SAlexei Starovoitov destroy_mem_alloc(ma, rcu_in_progress);
6977c8199e2SAlexei Starovoitov }
6987c8199e2SAlexei Starovoitov if (ma->caches) {
6999f2c6e96SAlexei Starovoitov rcu_in_progress = 0;
7007c8199e2SAlexei Starovoitov for_each_possible_cpu(cpu) {
7017c8199e2SAlexei Starovoitov cc = per_cpu_ptr(ma->caches, cpu);
7027c8199e2SAlexei Starovoitov for (i = 0; i < NUM_CACHES; i++) {
7037c8199e2SAlexei Starovoitov c = &cc->cache[i];
704d114dde2SAlexei Starovoitov WRITE_ONCE(c->draining, true);
7053d058187SHou Tao irq_work_sync(&c->refill_work);
7067c8199e2SAlexei Starovoitov drain_mem_cache(c);
70712c8d0f4SAlexei Starovoitov rcu_in_progress += atomic_read(&c->call_rcu_ttrace_in_progress);
7085af6807bSAlexei Starovoitov rcu_in_progress += atomic_read(&c->call_rcu_in_progress);
7097c8199e2SAlexei Starovoitov }
7107c8199e2SAlexei Starovoitov }
7117c8199e2SAlexei Starovoitov if (c->objcg)
7127c8199e2SAlexei Starovoitov obj_cgroup_put(c->objcg);
7139f2c6e96SAlexei Starovoitov destroy_mem_alloc(ma, rcu_in_progress);
7147c8199e2SAlexei Starovoitov }
7157c8199e2SAlexei Starovoitov }
7167c8199e2SAlexei Starovoitov
7177c8199e2SAlexei Starovoitov /* notrace is necessary here and in other functions to make sure
7187c8199e2SAlexei Starovoitov * bpf programs cannot attach to them and cause llist corruptions.
7197c8199e2SAlexei Starovoitov */
unit_alloc(struct bpf_mem_cache * c)7207c8199e2SAlexei Starovoitov static void notrace *unit_alloc(struct bpf_mem_cache *c)
7217c8199e2SAlexei Starovoitov {
7227c8199e2SAlexei Starovoitov struct llist_node *llnode = NULL;
7237c8199e2SAlexei Starovoitov unsigned long flags;
7247c8199e2SAlexei Starovoitov int cnt = 0;
7257c8199e2SAlexei Starovoitov
7267c8199e2SAlexei Starovoitov /* Disable irqs to prevent the following race for majority of prog types:
7277c8199e2SAlexei Starovoitov * prog_A
7287c8199e2SAlexei Starovoitov * bpf_mem_alloc
7297c8199e2SAlexei Starovoitov * preemption or irq -> prog_B
7307c8199e2SAlexei Starovoitov * bpf_mem_alloc
7317c8199e2SAlexei Starovoitov *
7327c8199e2SAlexei Starovoitov * but prog_B could be a perf_event NMI prog.
7337c8199e2SAlexei Starovoitov * Use per-cpu 'active' counter to order free_list access between
7347c8199e2SAlexei Starovoitov * unit_alloc/unit_free/bpf_mem_refill.
7357c8199e2SAlexei Starovoitov */
7367c8199e2SAlexei Starovoitov local_irq_save(flags);
7377c8199e2SAlexei Starovoitov if (local_inc_return(&c->active) == 1) {
7387c8199e2SAlexei Starovoitov llnode = __llist_del_first(&c->free_llist);
739822fb26bSAlexei Starovoitov if (llnode) {
7407c8199e2SAlexei Starovoitov cnt = --c->free_cnt;
741822fb26bSAlexei Starovoitov *(struct bpf_mem_cache **)llnode = c;
742822fb26bSAlexei Starovoitov }
7437c8199e2SAlexei Starovoitov }
7447c8199e2SAlexei Starovoitov local_dec(&c->active);
7457c8199e2SAlexei Starovoitov local_irq_restore(flags);
7467c8199e2SAlexei Starovoitov
7477c8199e2SAlexei Starovoitov WARN_ON(cnt < 0);
7487c8199e2SAlexei Starovoitov
7497c266178SAlexei Starovoitov if (cnt < c->low_watermark)
7507c8199e2SAlexei Starovoitov irq_work_raise(c);
7517c8199e2SAlexei Starovoitov return llnode;
7527c8199e2SAlexei Starovoitov }
7537c8199e2SAlexei Starovoitov
7547c8199e2SAlexei Starovoitov /* Though 'ptr' object could have been allocated on a different cpu
7557c8199e2SAlexei Starovoitov * add it to the free_llist of the current cpu.
7567c8199e2SAlexei Starovoitov * Let kfree() logic deal with it when it's later called from irq_work.
7577c8199e2SAlexei Starovoitov */
unit_free(struct bpf_mem_cache * c,void * ptr)7587c8199e2SAlexei Starovoitov static void notrace unit_free(struct bpf_mem_cache *c, void *ptr)
7597c8199e2SAlexei Starovoitov {
7607c8199e2SAlexei Starovoitov struct llist_node *llnode = ptr - LLIST_NODE_SZ;
7617c8199e2SAlexei Starovoitov unsigned long flags;
7627c8199e2SAlexei Starovoitov int cnt = 0;
7637c8199e2SAlexei Starovoitov
7647c8199e2SAlexei Starovoitov BUILD_BUG_ON(LLIST_NODE_SZ > 8);
7657c8199e2SAlexei Starovoitov
766822fb26bSAlexei Starovoitov /*
767822fb26bSAlexei Starovoitov * Remember bpf_mem_cache that allocated this object.
768822fb26bSAlexei Starovoitov * The hint is not accurate.
769822fb26bSAlexei Starovoitov */
770822fb26bSAlexei Starovoitov c->tgt = *(struct bpf_mem_cache **)llnode;
771822fb26bSAlexei Starovoitov
7727c8199e2SAlexei Starovoitov local_irq_save(flags);
7737c8199e2SAlexei Starovoitov if (local_inc_return(&c->active) == 1) {
7747c8199e2SAlexei Starovoitov __llist_add(llnode, &c->free_llist);
7757c8199e2SAlexei Starovoitov cnt = ++c->free_cnt;
7767c8199e2SAlexei Starovoitov } else {
7777c8199e2SAlexei Starovoitov /* unit_free() cannot fail. Therefore add an object to atomic
7787c8199e2SAlexei Starovoitov * llist. free_bulk() will drain it. Though free_llist_extra is
7797c8199e2SAlexei Starovoitov * a per-cpu list we have to use atomic llist_add here, since
7807c8199e2SAlexei Starovoitov * it also can be interrupted by bpf nmi prog that does another
7817c8199e2SAlexei Starovoitov * unit_free() into the same free_llist_extra.
7827c8199e2SAlexei Starovoitov */
7837c8199e2SAlexei Starovoitov llist_add(llnode, &c->free_llist_extra);
7847c8199e2SAlexei Starovoitov }
7857c8199e2SAlexei Starovoitov local_dec(&c->active);
7867c8199e2SAlexei Starovoitov local_irq_restore(flags);
7877c8199e2SAlexei Starovoitov
7887c266178SAlexei Starovoitov if (cnt > c->high_watermark)
7897c8199e2SAlexei Starovoitov /* free few objects from current cpu into global kmalloc pool */
7907c8199e2SAlexei Starovoitov irq_work_raise(c);
7917c8199e2SAlexei Starovoitov }
7927c8199e2SAlexei Starovoitov
unit_free_rcu(struct bpf_mem_cache * c,void * ptr)7935af6807bSAlexei Starovoitov static void notrace unit_free_rcu(struct bpf_mem_cache *c, void *ptr)
7945af6807bSAlexei Starovoitov {
7955af6807bSAlexei Starovoitov struct llist_node *llnode = ptr - LLIST_NODE_SZ;
7965af6807bSAlexei Starovoitov unsigned long flags;
7975af6807bSAlexei Starovoitov
7985af6807bSAlexei Starovoitov c->tgt = *(struct bpf_mem_cache **)llnode;
7995af6807bSAlexei Starovoitov
8005af6807bSAlexei Starovoitov local_irq_save(flags);
8015af6807bSAlexei Starovoitov if (local_inc_return(&c->active) == 1) {
8025af6807bSAlexei Starovoitov if (__llist_add(llnode, &c->free_by_rcu))
8035af6807bSAlexei Starovoitov c->free_by_rcu_tail = llnode;
8045af6807bSAlexei Starovoitov } else {
8055af6807bSAlexei Starovoitov llist_add(llnode, &c->free_llist_extra_rcu);
8065af6807bSAlexei Starovoitov }
8075af6807bSAlexei Starovoitov local_dec(&c->active);
8085af6807bSAlexei Starovoitov local_irq_restore(flags);
8095af6807bSAlexei Starovoitov
8105af6807bSAlexei Starovoitov if (!atomic_read(&c->call_rcu_in_progress))
8115af6807bSAlexei Starovoitov irq_work_raise(c);
8125af6807bSAlexei Starovoitov }
8135af6807bSAlexei Starovoitov
8147c8199e2SAlexei Starovoitov /* Called from BPF program or from sys_bpf syscall.
8157c8199e2SAlexei Starovoitov * In both cases migration is disabled.
8167c8199e2SAlexei Starovoitov */
bpf_mem_alloc(struct bpf_mem_alloc * ma,size_t size)8177c8199e2SAlexei Starovoitov void notrace *bpf_mem_alloc(struct bpf_mem_alloc *ma, size_t size)
8187c8199e2SAlexei Starovoitov {
8197c8199e2SAlexei Starovoitov int idx;
8207c8199e2SAlexei Starovoitov void *ret;
8217c8199e2SAlexei Starovoitov
8227c8199e2SAlexei Starovoitov if (!size)
823*3bc29c78SHou Tao return NULL;
8247c8199e2SAlexei Starovoitov
8257c8199e2SAlexei Starovoitov idx = bpf_mem_cache_idx(size + LLIST_NODE_SZ);
8267c8199e2SAlexei Starovoitov if (idx < 0)
8277c8199e2SAlexei Starovoitov return NULL;
8287c8199e2SAlexei Starovoitov
8297c8199e2SAlexei Starovoitov ret = unit_alloc(this_cpu_ptr(ma->caches)->cache + idx);
8307c8199e2SAlexei Starovoitov return !ret ? NULL : ret + LLIST_NODE_SZ;
8317c8199e2SAlexei Starovoitov }
8327c8199e2SAlexei Starovoitov
bpf_mem_free(struct bpf_mem_alloc * ma,void * ptr)8337c8199e2SAlexei Starovoitov void notrace bpf_mem_free(struct bpf_mem_alloc *ma, void *ptr)
8347c8199e2SAlexei Starovoitov {
835*3bc29c78SHou Tao struct bpf_mem_cache *c;
8367c8199e2SAlexei Starovoitov int idx;
8377c8199e2SAlexei Starovoitov
8387c8199e2SAlexei Starovoitov if (!ptr)
8397c8199e2SAlexei Starovoitov return;
8407c8199e2SAlexei Starovoitov
841*3bc29c78SHou Tao c = *(void **)(ptr - LLIST_NODE_SZ);
842*3bc29c78SHou Tao idx = bpf_mem_cache_idx(c->unit_size);
843*3bc29c78SHou Tao if (WARN_ON_ONCE(idx < 0))
8447c8199e2SAlexei Starovoitov return;
8457c8199e2SAlexei Starovoitov
8467c8199e2SAlexei Starovoitov unit_free(this_cpu_ptr(ma->caches)->cache + idx, ptr);
8477c8199e2SAlexei Starovoitov }
8487c8199e2SAlexei Starovoitov
bpf_mem_free_rcu(struct bpf_mem_alloc * ma,void * ptr)8495af6807bSAlexei Starovoitov void notrace bpf_mem_free_rcu(struct bpf_mem_alloc *ma, void *ptr)
8505af6807bSAlexei Starovoitov {
851*3bc29c78SHou Tao struct bpf_mem_cache *c;
8525af6807bSAlexei Starovoitov int idx;
8535af6807bSAlexei Starovoitov
8545af6807bSAlexei Starovoitov if (!ptr)
8555af6807bSAlexei Starovoitov return;
8565af6807bSAlexei Starovoitov
857*3bc29c78SHou Tao c = *(void **)(ptr - LLIST_NODE_SZ);
858*3bc29c78SHou Tao idx = bpf_mem_cache_idx(c->unit_size);
859*3bc29c78SHou Tao if (WARN_ON_ONCE(idx < 0))
8605af6807bSAlexei Starovoitov return;
8615af6807bSAlexei Starovoitov
8625af6807bSAlexei Starovoitov unit_free_rcu(this_cpu_ptr(ma->caches)->cache + idx, ptr);
8635af6807bSAlexei Starovoitov }
8645af6807bSAlexei Starovoitov
bpf_mem_cache_alloc(struct bpf_mem_alloc * ma)8657c8199e2SAlexei Starovoitov void notrace *bpf_mem_cache_alloc(struct bpf_mem_alloc *ma)
8667c8199e2SAlexei Starovoitov {
8677c8199e2SAlexei Starovoitov void *ret;
8687c8199e2SAlexei Starovoitov
8697c8199e2SAlexei Starovoitov ret = unit_alloc(this_cpu_ptr(ma->cache));
8707c8199e2SAlexei Starovoitov return !ret ? NULL : ret + LLIST_NODE_SZ;
8717c8199e2SAlexei Starovoitov }
8727c8199e2SAlexei Starovoitov
bpf_mem_cache_free(struct bpf_mem_alloc * ma,void * ptr)8737c8199e2SAlexei Starovoitov void notrace bpf_mem_cache_free(struct bpf_mem_alloc *ma, void *ptr)
8747c8199e2SAlexei Starovoitov {
8757c8199e2SAlexei Starovoitov if (!ptr)
8767c8199e2SAlexei Starovoitov return;
8777c8199e2SAlexei Starovoitov
8787c8199e2SAlexei Starovoitov unit_free(this_cpu_ptr(ma->cache), ptr);
8797c8199e2SAlexei Starovoitov }
880e65a5c6eSMartin KaFai Lau
bpf_mem_cache_free_rcu(struct bpf_mem_alloc * ma,void * ptr)8815af6807bSAlexei Starovoitov void notrace bpf_mem_cache_free_rcu(struct bpf_mem_alloc *ma, void *ptr)
8825af6807bSAlexei Starovoitov {
8835af6807bSAlexei Starovoitov if (!ptr)
8845af6807bSAlexei Starovoitov return;
8855af6807bSAlexei Starovoitov
8865af6807bSAlexei Starovoitov unit_free_rcu(this_cpu_ptr(ma->cache), ptr);
8875af6807bSAlexei Starovoitov }
8885af6807bSAlexei Starovoitov
889e65a5c6eSMartin KaFai Lau /* Directly does a kfree() without putting 'ptr' back to the free_llist
890e65a5c6eSMartin KaFai Lau * for reuse and without waiting for a rcu_tasks_trace gp.
891e65a5c6eSMartin KaFai Lau * The caller must first go through the rcu_tasks_trace gp for 'ptr'
892e65a5c6eSMartin KaFai Lau * before calling bpf_mem_cache_raw_free().
893e65a5c6eSMartin KaFai Lau * It could be used when the rcu_tasks_trace callback does not have
894e65a5c6eSMartin KaFai Lau * a hold on the original bpf_mem_alloc object that allocated the
895e65a5c6eSMartin KaFai Lau * 'ptr'. This should only be used in the uncommon code path.
896e65a5c6eSMartin KaFai Lau * Otherwise, the bpf_mem_alloc's free_llist cannot be refilled
897e65a5c6eSMartin KaFai Lau * and may affect performance.
898e65a5c6eSMartin KaFai Lau */
bpf_mem_cache_raw_free(void * ptr)899e65a5c6eSMartin KaFai Lau void bpf_mem_cache_raw_free(void *ptr)
900e65a5c6eSMartin KaFai Lau {
901e65a5c6eSMartin KaFai Lau if (!ptr)
902e65a5c6eSMartin KaFai Lau return;
903e65a5c6eSMartin KaFai Lau
904e65a5c6eSMartin KaFai Lau kfree(ptr - LLIST_NODE_SZ);
905e65a5c6eSMartin KaFai Lau }
906e65a5c6eSMartin KaFai Lau
907e65a5c6eSMartin KaFai Lau /* When flags == GFP_KERNEL, it signals that the caller will not cause
908e65a5c6eSMartin KaFai Lau * deadlock when using kmalloc. bpf_mem_cache_alloc_flags() will use
909e65a5c6eSMartin KaFai Lau * kmalloc if the free_llist is empty.
910e65a5c6eSMartin KaFai Lau */
bpf_mem_cache_alloc_flags(struct bpf_mem_alloc * ma,gfp_t flags)911e65a5c6eSMartin KaFai Lau void notrace *bpf_mem_cache_alloc_flags(struct bpf_mem_alloc *ma, gfp_t flags)
912e65a5c6eSMartin KaFai Lau {
913e65a5c6eSMartin KaFai Lau struct bpf_mem_cache *c;
914e65a5c6eSMartin KaFai Lau void *ret;
915e65a5c6eSMartin KaFai Lau
916e65a5c6eSMartin KaFai Lau c = this_cpu_ptr(ma->cache);
917e65a5c6eSMartin KaFai Lau
918e65a5c6eSMartin KaFai Lau ret = unit_alloc(c);
919e65a5c6eSMartin KaFai Lau if (!ret && flags == GFP_KERNEL) {
920e65a5c6eSMartin KaFai Lau struct mem_cgroup *memcg, *old_memcg;
921e65a5c6eSMartin KaFai Lau
922e65a5c6eSMartin KaFai Lau memcg = get_memcg(c);
923e65a5c6eSMartin KaFai Lau old_memcg = set_active_memcg(memcg);
924e65a5c6eSMartin KaFai Lau ret = __alloc(c, NUMA_NO_NODE, GFP_KERNEL | __GFP_NOWARN | __GFP_ACCOUNT);
925d9105720SHou Tao if (ret)
926d9105720SHou Tao *(struct bpf_mem_cache **)ret = c;
927e65a5c6eSMartin KaFai Lau set_active_memcg(old_memcg);
928e65a5c6eSMartin KaFai Lau mem_cgroup_put(memcg);
929e65a5c6eSMartin KaFai Lau }
930e65a5c6eSMartin KaFai Lau
931e65a5c6eSMartin KaFai Lau return !ret ? NULL : ret + LLIST_NODE_SZ;
932e65a5c6eSMartin KaFai Lau }
933