[PATCH 1/2] cgroup, sched: add BPF kfuncs to read a cpu cgroup's stats

[Ziyang Men <ziyang.meme@gmail.com>]

This series adds bpf kfuncs for the cgroup CPU controller, following the
memory controller kfuncs in mm/bpf_memcontrol.c.

Collecting cgroup statistics is expensive: the existing
method is to open and parse a cgroup file. memcg already has an
efficient alternative through BPF; this series extends that idea to cpu.

Design:
- Leave reading the CFS bandwidth counters to the BPF program. They are
  plain fields of tg->cfs_bandwidth, so they need no kernel code.
- Add one kfunc to compute the throttled time. This is necessary because
  it is a sum over every possible cpu, which a user cannot do itself.
- The bpf_cpu_cgroup_cputime() returns all five base CPU-time values in one
  call with one cputime_adjust().

The only part it touches the scheduler part is to discard the static for
throttled_time_self() in order to use externally.

The two kfuncs that take the rstat lock are KF_SLEEPABLE following idea
in the mm/bpf_memcontrol.c

Suggested-by: Shakeel Butt <shakeel.butt@linux.dev>
Assisted-by: Claude:claude-opus-5
Signed-off-by: Ziyang Men <ziyang.meme@gmail.com>
---
 include/linux/cgroup.h          | 15 +++++++
 kernel/cgroup/Makefile          |  2 +
 kernel/cgroup/bpf_cpu.c         | 80 +++++++++++++++++++++++++++++++++
 kernel/cgroup/cgroup-internal.h |  3 ++
 kernel/cgroup/rstat.c           | 42 +++++++++++++++++
 kernel/sched/core.c             |  2 +-
 6 files changed, 143 insertions(+), 1 deletion(-)
 create mode 100644 kernel/cgroup/bpf_cpu.c

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index f2aa46a4f871..d2a6b5efad51 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -923,4 +923,19 @@ struct cgroup *task_get_cgroup1(struct task_struct *tsk, int hierarchy_id);
 
 struct cgroup_of_peak *of_peak(struct kernfs_open_file *of);
 
+/* A cgroup's base CPU-time counters in microseconds, as cpu.stat prints them */
+struct cpu_cgroup_cputime {
+	u64 usage_usec;
+	u64 user_usec;
+	u64 system_usec;
+	u64 nice_usec;
+	u64 forceidle_usec;	/* 0 without CONFIG_SCHED_CORE */
+};
+
+/* A task_group's own throttled time in nanoseconds; see cpu.stat.local */
+struct task_group;
+#ifdef CONFIG_CFS_BANDWIDTH
+u64 throttled_time_self(struct task_group *tg);
+#endif
+
 #endif /* _LINUX_CGROUP_H */
diff --git a/kernel/cgroup/Makefile b/kernel/cgroup/Makefile
index ede31601a363..0ba59b7eef48 100644
--- a/kernel/cgroup/Makefile
+++ b/kernel/cgroup/Makefile
@@ -1,6 +1,8 @@
 # SPDX-License-Identifier: GPL-2.0
 obj-y := cgroup.o rstat.o namespace.o cgroup-v1.o freezer.o
 
+obj-$(CONFIG_BPF_SYSCALL) += bpf_cpu.o
+
 obj-$(CONFIG_CGROUP_FREEZER) += legacy_freezer.o
 obj-$(CONFIG_CGROUP_PIDS) += pids.o
 obj-$(CONFIG_CGROUP_RDMA) += rdma.o
diff --git a/kernel/cgroup/bpf_cpu.c b/kernel/cgroup/bpf_cpu.c
new file mode 100644
index 000000000000..6eb89c8e84fd
--- /dev/null
+++ b/kernel/cgroup/bpf_cpu.c
@@ -0,0 +1,80 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * CPU Controller-related BPF kfuncs
+ *
+ * bpf_cpu_cgroup_cputime() is defined in rstat.c, which owns the locking it
+ * needs, and only registered here.
+ *
+ * Author: Ziyang Men <ziyang.meme@gmail.com>
+ */
+
+#include <linux/bpf.h>
+#include <linux/btf_ids.h>
+#include <linux/cgroup.h>
+
+#include "cgroup-internal.h"
+
+__bpf_kfunc_start_defs();
+
+/**
+ * bpf_cpu_cgroup_flush_stats - Flush a cgroup's base CPU-time statistics
+ * @cgrp: cgroup to flush
+ *
+ * Propagate the cgroup's base CPU-time statistics up the cgroup tree.
+ */
+__bpf_kfunc void bpf_cpu_cgroup_flush_stats(struct cgroup *cgrp)
+{
+	css_rstat_flush(&cgrp->self);
+}
+
+/**
+ * bpf_cpu_cgroup_throttled_self - Read a cgroup's own throttled time
+ * @cgrp: cgroup to read from
+ *
+ * Return: The throttled time in microseconds, or 0 if config is off.
+ */
+__bpf_kfunc u64 bpf_cpu_cgroup_throttled_self(struct cgroup *cgrp)
+{
+/* cpu_cgrp_id needs the cpu controller, which CFS bandwidth depends on */
+#ifdef CONFIG_CFS_BANDWIDTH
+	struct cgroup_subsys_state *css;
+
+	guard(rcu)();
+
+	css = rcu_dereference(cgrp->subsys[cpu_cgrp_id]);
+	if (!css)
+		return 0;
+
+	return div_u64(throttled_time_self((struct task_group *)css),
+		       NSEC_PER_USEC);
+#else
+	return 0;
+#endif
+}
+
+__bpf_kfunc_end_defs();
+
+/* KF_SLEEPABLE keeps the rstat spinlock out of NMI */
+BTF_KFUNCS_START(bpf_cpu_cgroup_kfunc_ids)
+BTF_ID_FLAGS(func, bpf_cpu_cgroup_flush_stats, KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_cpu_cgroup_cputime, KF_SLEEPABLE)
+BTF_ID_FLAGS(func, bpf_cpu_cgroup_throttled_self)
+BTF_KFUNCS_END(bpf_cpu_cgroup_kfunc_ids)
+
+static const struct btf_kfunc_id_set bpf_cpu_cgroup_kfunc_set = {
+	.owner		= THIS_MODULE,
+	.set		= &bpf_cpu_cgroup_kfunc_ids,
+};
+
+static int __init bpf_cpu_cgroup_kfunc_init(void)
+{
+	int err;
+
+	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
+					&bpf_cpu_cgroup_kfunc_set);
+	if (err)
+		pr_warn("error while registering cpu cgroup kfuncs: %d\n", err);
+
+	return err;
+}
+late_initcall(bpf_cpu_cgroup_kfunc_init);
diff --git a/kernel/cgroup/cgroup-internal.h b/kernel/cgroup/cgroup-internal.h
index 58797123b752..65f5b6318289 100644
--- a/kernel/cgroup/cgroup-internal.h
+++ b/kernel/cgroup/cgroup-internal.h
@@ -271,6 +271,9 @@ int css_rstat_init(struct cgroup_subsys_state *css);
 void css_rstat_exit(struct cgroup_subsys_state *css);
 int ss_rstat_init(struct cgroup_subsys *ss);
 void cgroup_base_stat_cputime_show(struct seq_file *seq);
+#ifdef CONFIG_BPF_SYSCALL
+void bpf_cpu_cgroup_cputime(struct cgroup *cgrp, struct cpu_cgroup_cputime *out);
+#endif
 
 /*
  * namespace.c
diff --git a/kernel/cgroup/rstat.c b/kernel/cgroup/rstat.c
index de816a43db9f..f9e30719068e 100644
--- a/kernel/cgroup/rstat.c
+++ b/kernel/cgroup/rstat.c
@@ -752,6 +752,48 @@ void cgroup_base_stat_cputime_show(struct seq_file *seq)
 	cgroup_force_idle_show(seq, &bstat);
 }
 
+#ifdef CONFIG_BPF_SYSCALL
+
+__bpf_kfunc_start_defs();
+
+/**
+ * bpf_cpu_cgroup_cputime - Read a cgroup's base CPU-time data
+ * @cgrp: cgroup to read from
+ * @out: the data in microseconds. Zero it first: the verifier reads the
+ *       whole struct.
+ *
+ * Adjust once and fill all values.
+ */
+__bpf_kfunc void bpf_cpu_cgroup_cputime(struct cgroup *cgrp,
+					struct cpu_cgroup_cputime *out)
+{
+	struct cgroup_base_stat bstat;
+
+	if (cgroup_parent(cgrp)) {
+		__css_rstat_lock(&cgrp->self, -1);
+		bstat = cgrp->bstat;
+		cputime_adjust(&cgrp->bstat.cputime, &cgrp->prev_cputime,
+			       &bstat.cputime.utime, &bstat.cputime.stime);
+		__css_rstat_unlock(&cgrp->self, -1);
+	} else {
+		root_cgroup_cputime(&bstat);
+	}
+
+	out->usage_usec = div_u64(bstat.cputime.sum_exec_runtime, NSEC_PER_USEC);
+	out->user_usec = div_u64(bstat.cputime.utime, NSEC_PER_USEC);
+	out->system_usec = div_u64(bstat.cputime.stime, NSEC_PER_USEC);
+	out->nice_usec = div_u64(bstat.ntime, NSEC_PER_USEC);
+#ifdef CONFIG_SCHED_CORE
+	out->forceidle_usec = div_u64(bstat.forceidle_sum, NSEC_PER_USEC);
+#else
+	out->forceidle_usec = 0;
+#endif
+}
+
+__bpf_kfunc_end_defs();
+
+#endif /* CONFIG_BPF_SYSCALL */
+
 /* Add bpf kfuncs for css_rstat_updated() and css_rstat_flush() */
 BTF_KFUNCS_START(bpf_rstat_kfunc_ids)
 BTF_ID_FLAGS(func, css_rstat_updated)
diff --git a/kernel/sched/core.c b/kernel/sched/core.c
index 96226707c2f6..75735e0e81ef 100644
--- a/kernel/sched/core.c
+++ b/kernel/sched/core.c
@@ -10027,7 +10027,7 @@ static int cpu_cfs_stat_show(struct seq_file *sf, void *v)
 	return 0;
 }
 
-static u64 throttled_time_self(struct task_group *tg)
+u64 throttled_time_self(struct task_group *tg)
 {
 	int i;
 	u64 total = 0;
-- 
2.53.0-Meta