1#include "bpf_experimental.h"
2
3struct val_t {
4	long b, c, d;
5};
6
7struct elem {
8	long sum;
9	struct val_t __percpu_kptr *pc;
10};
11
12struct {
13	__uint(type, BPF_MAP_TYPE_CGRP_STORAGE);
14	__uint(map_flags, BPF_F_NO_PREALLOC);
15	__type(key, int);
16	__type(value, struct elem);
17} cgrp SEC(".maps");
18
19const volatile int nr_cpus;
20
21/* Initialize the percpu object */
22SEC("fentry/bpf_fentry_test1")
23int BPF_PROG(test_cgrp_local_storage_1)
24{
25	struct task_struct *task;
26	struct val_t __percpu_kptr *p;
27	struct elem *e;
28
29	task = bpf_get_current_task_btf();
30	e = bpf_cgrp_storage_get(&cgrp, task->cgroups->dfl_cgrp, 0,
31				 BPF_LOCAL_STORAGE_GET_F_CREATE);
32	if (!e)
33		return 0;
34
35	p = bpf_percpu_obj_new(struct val_t);
36	if (!p)
37		return 0;
38
39	p = bpf_kptr_xchg(&e->pc, p);
40	if (p)
41		bpf_percpu_obj_drop(p);
42
43	return 0;
44}
45
46/* Percpu data collection */
47SEC("fentry/bpf_fentry_test2")
48int BPF_PROG(test_cgrp_local_storage_2)
49{
50	struct task_struct *task;
51	struct val_t __percpu_kptr *p;
52	struct val_t *v;
53	struct elem *e;
54
55	task = bpf_get_current_task_btf();
56	e = bpf_cgrp_storage_get(&cgrp, task->cgroups->dfl_cgrp, 0, 0);
57	if (!e)
58		return 0;
59
60	p = e->pc;
61	if (!p)
62		return 0;
63
64	v = bpf_per_cpu_ptr(p, 0);
65	if (!v)
66		return 0;
67	v->c = 1;
68	v->d = 2;
69	return 0;
70}
71
72int cpu0_field_d, sum_field_c;
73int my_pid;
74
75/* Summarize percpu data collection */
76SEC("fentry/bpf_fentry_test3")
77int BPF_PROG(test_cgrp_local_storage_3)
78{
79	struct task_struct *task;
80	struct val_t __percpu_kptr *p;
81	struct val_t *v;
82	struct elem *e;
83	int i;
84
85	if ((bpf_get_current_pid_tgid() >> 32) != my_pid)
86		return 0;
87
88	task = bpf_get_current_task_btf();
89	e = bpf_cgrp_storage_get(&cgrp, task->cgroups->dfl_cgrp, 0, 0);
90	if (!e)
91		return 0;
92
93	p = e->pc;
94	if (!p)
95		return 0;
96
97	bpf_for(i, 0, nr_cpus) {
98		v = bpf_per_cpu_ptr(p, i);
99		if (v) {
100			if (i == 0)
101				cpu0_field_d = v->d;
102			sum_field_c += v->c;
103		}
104	}
105
106	return 0;
107}
108
109char _license[] SEC("license") = "GPL";
110