[PATCH v3 3/4] block: add BPF kfuncs to read blkcg io.stat

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

Collecting cgroup statistics is expensive because the existing method
opens and parses a cgroup file for every cgroup. memcg already provides
an efficient BPF interface; extend that model to the block controller.

Add bpf_cgroup_css() and bpf_css_release() to acquire a controller's
css from a cgroup. The reference keeps the css alive across the
sleepable css_rstat_flush().

Add bpf_css_to_blkcg() as a checked RCU-protected css-to-blkcg
conversion and an open-coded iterator for the per-device blkgs.

Suggested-by: Shakeel Butt <shakeel.butt@linux.dev>
Suggested-by: Tejun Heo <tj@kernel.org>
Assisted-by: Claude:claude-opus-5
Signed-off-by: Ziyang Men <ziyang.meme@gmail.com>
---
 MAINTAINERS                |   1 +
 block/Makefile             |   3 +
 block/bpf_blkcg.c          | 138 +++++++++++++++++++++++++++++++++++++
 kernel/cgroup/bpf_cgroup.c |  62 ++++++++++++++---
 4 files changed, 194 insertions(+), 10 deletions(-)
 create mode 100644 block/bpf_blkcg.c

diff --git a/MAINTAINERS b/MAINTAINERS
index 2f9472c1a090..87c56e955577 100644
--- a/MAINTAINERS
+++ b/MAINTAINERS
@@ -6617,6 +6617,7 @@ F:	block/blk-cgroup.c
 F:	block/blk-iocost.c
 F:	block/blk-iolatency.c
 F:	block/blk-throttle.c
+F:	block/bpf_blkcg.c
 F:	include/linux/blk-cgroup.h
 
 CONTROL GROUP - CPUSET
diff --git a/block/Makefile b/block/Makefile
index e7bd320e3d69..572e49988c8e 100644
--- a/block/Makefile
+++ b/block/Makefile
@@ -17,6 +17,9 @@ obj-$(CONFIG_BLK_ERROR_INJECTION) += error-injection.o
 obj-$(CONFIG_BLK_DEV_BSG_COMMON) += bsg.o
 obj-$(CONFIG_BLK_DEV_BSGLIB)	+= bsg-lib.o
 obj-$(CONFIG_BLK_CGROUP)	+= blk-cgroup.o
+ifdef CONFIG_BPF_SYSCALL
+obj-$(CONFIG_BLK_CGROUP)	+= bpf_blkcg.o
+endif
 obj-$(CONFIG_BLK_CGROUP_RWSTAT)	+= blk-cgroup-rwstat.o
 obj-$(CONFIG_BLK_CGROUP_FC_APPID) += blk-cgroup-fc-appid.o
 obj-$(CONFIG_BLK_DEV_THROTTLING)	+= blk-throttle.o
diff --git a/block/bpf_blkcg.c b/block/bpf_blkcg.c
new file mode 100644
index 000000000000..d8ab8006bc57
--- /dev/null
+++ b/block/bpf_blkcg.c
@@ -0,0 +1,138 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+/*
+ * Block I/O Controller-related BPF kfuncs and auxiliary code
+ */
+
+#include "blk-cgroup.h"
+
+#include <linux/bpf.h>
+#include <linux/btf_ids.h>
+#include <linux/rculist.h>
+
+__bpf_kfunc_start_defs();
+
+/**
+ * bpf_css_to_blkcg - Cast an io controller css to its block cgroup
+ * @css: io controller css
+ *
+ * Must be called under RCU.
+ *
+ * Return: The block cgroup, or NULL if @css belongs to another controller.
+ */
+__bpf_kfunc struct blkcg *
+bpf_css_to_blkcg(struct cgroup_subsys_state *css)
+{
+	if (unlikely(css->ss != &io_cgrp_subsys))
+		return NULL;
+
+	return css_to_blkcg(css);
+}
+
+struct bpf_iter_blkg {
+	__u64 __opaque[2];
+} __aligned(8);
+
+struct bpf_iter_blkg_kern {
+	struct blkcg *blkcg;
+	struct blkcg_gq *pos;
+} __aligned(8);
+
+/**
+ * bpf_iter_blkg_new - Start iterating a block cgroup's per-device blkgs
+ * @it: iterator to initialize
+ * @blkcg: block cgroup to iterate
+ *
+ * Each blkg holds one device's io.stat counters. Offline blkgs are skipped.
+ * A blkg without a disk can be returned. Root blkgs do not contain the
+ * system-wide statistics shown by root io.stat. Must run under RCU.
+ *
+ * Return: 0 on success.
+ */
+__bpf_kfunc int bpf_iter_blkg_new(struct bpf_iter_blkg *it,
+				  struct blkcg *blkcg)
+{
+	struct bpf_iter_blkg_kern *kit = (void *)it;
+
+	BUILD_BUG_ON(sizeof(struct bpf_iter_blkg_kern) > sizeof(struct bpf_iter_blkg));
+	BUILD_BUG_ON(__alignof__(struct bpf_iter_blkg_kern) !=
+		     __alignof__(struct bpf_iter_blkg));
+
+	kit->pos = NULL;
+	kit->blkcg = blkcg;
+	return 0;
+}
+
+/**
+ * bpf_iter_blkg_next - Return the next online blkg of the iterated block cgroup
+ * @it: iterator
+ *
+ * Return: the next online blkg, or NULL when the walk is done.
+ */
+__bpf_kfunc struct blkcg_gq *bpf_iter_blkg_next(struct bpf_iter_blkg *it)
+{
+	struct bpf_iter_blkg_kern *kit = (void *)it;
+	struct blkcg_gq *blkg = kit->pos;
+	struct hlist_node *node;
+
+	if (!kit->blkcg)
+		return NULL;
+
+	if (!blkg)
+		node = rcu_dereference(hlist_first_rcu(&kit->blkcg->blkg_list));
+	else
+		node = rcu_dereference(hlist_next_rcu(&blkg->blkcg_node));
+
+	/* Skip offline blkgs, matching io.stat. */
+	while (node) {
+		blkg = hlist_entry(node, struct blkcg_gq, blkcg_node);
+		/* A race only changes whether this blkg is returned. */
+		if (data_race(blkg->online)) {
+			kit->pos = blkg;
+			return blkg;
+		}
+		node = rcu_dereference(hlist_next_rcu(&blkg->blkcg_node));
+	}
+
+	/* The iterator must keep returning NULL after completion. */
+	kit->pos = NULL;
+	kit->blkcg = NULL;
+	return NULL;
+}
+
+/**
+ * bpf_iter_blkg_destroy - Tear down a blkg iterator
+ * @it: iterator
+ */
+__bpf_kfunc void bpf_iter_blkg_destroy(struct bpf_iter_blkg *it)
+{
+}
+
+__bpf_kfunc_end_defs();
+
+BTF_KFUNCS_START(bpf_blkcg_kfuncs)
+BTF_ID_FLAGS(func, bpf_css_to_blkcg,
+	     KF_RCU | KF_RCU_PROTECTED | KF_RET_NULL)
+
+BTF_ID_FLAGS(func, bpf_iter_blkg_new,
+	     KF_ITER_NEW | KF_RCU | KF_RCU_PROTECTED)
+BTF_ID_FLAGS(func, bpf_iter_blkg_next, KF_ITER_NEXT | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_iter_blkg_destroy, KF_ITER_DESTROY)
+BTF_KFUNCS_END(bpf_blkcg_kfuncs)
+
+static const struct btf_kfunc_id_set bpf_blkcg_kfunc_set = {
+	.owner		= THIS_MODULE,
+	.set		= &bpf_blkcg_kfuncs,
+};
+
+static int __init bpf_blkcg_init(void)
+{
+	int err;
+
+	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
+					&bpf_blkcg_kfunc_set);
+	if (err)
+		pr_warn("error while registering bpf blkcg kfuncs: %d\n", err);
+
+	return err;
+}
+late_initcall(bpf_blkcg_init);
diff --git a/kernel/cgroup/bpf_cgroup.c b/kernel/cgroup/bpf_cgroup.c
index cd28c838dc7b..e253633e8278 100644
--- a/kernel/cgroup/bpf_cgroup.c
+++ b/kernel/cgroup/bpf_cgroup.c
@@ -8,12 +8,50 @@
 #include <linux/bpf.h>
 #include <linux/btf_ids.h>
 #include <linux/cgroup.h>
+#include <linux/rcupdate.h>
 
+#ifdef CONFIG_CGROUP_SCHED
 #include "../sched/sched.h"
+#endif
 
-#ifdef CONFIG_CGROUP_SCHED
 __bpf_kfunc_start_defs();
 
+/**
+ * bpf_cgroup_css - Get a reference to one controller's css
+ * @cgrp: cgroup to look in
+ * @ssid: controller ID
+ *
+ * The returned css must be released with bpf_css_release().
+ *
+ * Return: The referenced css, or NULL.
+ */
+__bpf_kfunc struct cgroup_subsys_state *
+bpf_cgroup_css(struct cgroup *cgrp, int ssid)
+{
+	struct cgroup_subsys_state *css;
+
+	if (unlikely(ssid < 0 || ssid >= CGROUP_SUBSYS_COUNT))
+		return NULL;
+
+	rcu_read_lock();
+	css = rcu_dereference(cgrp->subsys[ssid]);
+	if (css && !css_tryget(css))
+		css = NULL;
+	rcu_read_unlock();
+
+	return css;
+}
+
+/**
+ * bpf_css_release - Release a css reference
+ * @css: css to release
+ */
+__bpf_kfunc void bpf_css_release(struct cgroup_subsys_state *css)
+{
+	css_put(css);
+}
+
+#ifdef CONFIG_CGROUP_SCHED
 /**
  * bpf_css_to_task_group - Cast a CPU controller css to its task group
  * @css: CPU controller css
@@ -30,29 +68,33 @@ bpf_css_to_task_group(struct cgroup_subsys_state *css)
 
 	return container_of(css, struct task_group, css);
 }
+#endif /* CONFIG_CGROUP_SCHED */
 
 __bpf_kfunc_end_defs();
 
-BTF_KFUNCS_START(bpf_cpu_cgroup_kfunc_ids)
+BTF_KFUNCS_START(bpf_cgroup_kfunc_ids)
+BTF_ID_FLAGS(func, bpf_cgroup_css, KF_ACQUIRE | KF_RCU | KF_RET_NULL)
+BTF_ID_FLAGS(func, bpf_css_release, KF_RELEASE)
+#ifdef CONFIG_CGROUP_SCHED
 BTF_ID_FLAGS(func, bpf_css_to_task_group,
 	     KF_RCU | KF_RCU_PROTECTED | KF_RET_NULL)
-BTF_KFUNCS_END(bpf_cpu_cgroup_kfunc_ids)
+#endif
+BTF_KFUNCS_END(bpf_cgroup_kfunc_ids)
 
-static const struct btf_kfunc_id_set bpf_cpu_cgroup_kfunc_set = {
+static const struct btf_kfunc_id_set bpf_cgroup_kfunc_set = {
 	.owner		= THIS_MODULE,
-	.set		= &bpf_cpu_cgroup_kfunc_ids,
+	.set		= &bpf_cgroup_kfunc_ids,
 };
 
-static int __init bpf_cpu_cgroup_kfunc_init(void)
+static int __init bpf_cgroup_kfunc_init(void)
 {
 	int err;
 
 	err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
-					&bpf_cpu_cgroup_kfunc_set);
+					&bpf_cgroup_kfunc_set);
 	if (err)
-		pr_warn("error while registering cpu cgroup kfuncs: %d\n", err);
+		pr_warn("error while registering cgroup kfuncs: %d\n", err);
 
 	return err;
 }
-late_initcall(bpf_cpu_cgroup_kfunc_init);
-#endif /* CONFIG_CGROUP_SCHED */
+late_initcall(bpf_cgroup_kfunc_init);
-- 
2.53.0-Meta