[RFC PATCH bpf-next 2/8] cgroup: Add new helpers for cgroup1 hierarchy

[Yafang Shao <laoar.shao@gmail.com>]

Two new helpers are added for cgroup1 hierarchy:

- task_cgroup1_id_within_hierarchy
  Retrieves the associated cgroup ID of a task within a specific cgroup1
  hierarchy. The cgroup1 hierarchy is identified by its hierarchy ID.
- task_ancestor_cgroup1_id_within_hierarchy
  Retrieves the associated ancestor cgroup ID of a task whithin a
  specific cgroup1 hierarchy. The specific ancestor level is determined by
  its ancestor level.

These helper functions have been added to facilitate the tracing of tasks
within a particular container or cgroup in BPF programs. It's important to
note that these helpers are designed specifically for cgroup1.

Signed-off-by: Yafang Shao <laoar.shao@gmail.com>
---
 include/linux/cgroup.h          |  9 ++++-
 kernel/cgroup/cgroup-internal.h |  2 -
 kernel/cgroup/cgroup-v1.c       | 67 +++++++++++++++++++++++++++++++++
 3 files changed, 75 insertions(+), 3 deletions(-)

diff --git a/include/linux/cgroup.h b/include/linux/cgroup.h
index b307013b9c6c..65bde6eb41ef 100644
--- a/include/linux/cgroup.h
+++ b/include/linux/cgroup.h
@@ -71,6 +71,8 @@ struct css_task_iter {
 extern struct file_system_type cgroup_fs_type;
 extern struct cgroup_root cgrp_dfl_root;
 extern struct css_set init_css_set;
+extern struct list_head cgroup_roots;
+extern spinlock_t css_set_lock;
 
 #define SUBSYS(_x) extern struct cgroup_subsys _x ## _cgrp_subsys;
 #include <linux/cgroup_subsys.h>
@@ -159,6 +161,8 @@ void css_task_iter_start(struct cgroup_subsys_state *css, unsigned int flags,
 			 struct css_task_iter *it);
 struct task_struct *css_task_iter_next(struct css_task_iter *it);
 void css_task_iter_end(struct css_task_iter *it);
+struct cgroup *task_cgroup_from_root(struct task_struct *task,
+				     struct cgroup_root *root);
 
 /**
  * css_for_each_child - iterate through children of a css
@@ -388,7 +392,6 @@ static inline void cgroup_unlock(void)
  * as locks used during the cgroup_subsys::attach() methods.
  */
 #ifdef CONFIG_PROVE_RCU
-extern spinlock_t css_set_lock;
 #define task_css_set_check(task, __c)					\
 	rcu_dereference_check((task)->cgroups,				\
 		rcu_read_lock_sched_held() ||				\
@@ -855,4 +858,8 @@ static inline void cgroup_bpf_put(struct cgroup *cgrp) {}
 
 #endif /* CONFIG_CGROUP_BPF */
 
+u64 task_cgroup1_id_within_hierarchy(struct task_struct *tsk, int hierarchy_id);
+u64 task_ancestor_cgroup1_id_within_hierarchy(struct task_struct *tsk, int hierarchy_id,
+					      int ancestor_level);
+
 #endif /* _LINUX_CGROUP_H */
diff --git a/kernel/cgroup/cgroup-internal.h b/kernel/cgroup/cgroup-internal.h
index c56071f150f2..2c32a80c1334 100644
--- a/kernel/cgroup/cgroup-internal.h
+++ b/kernel/cgroup/cgroup-internal.h
@@ -164,9 +164,7 @@ struct cgroup_mgctx {
 #define DEFINE_CGROUP_MGCTX(name)						\
 	struct cgroup_mgctx name = CGROUP_MGCTX_INIT(name)
 
-extern spinlock_t css_set_lock;
 extern struct cgroup_subsys *cgroup_subsys[];
-extern struct list_head cgroup_roots;
 
 /* iterate across the hierarchies */
 #define for_each_root(root)						\
diff --git a/kernel/cgroup/cgroup-v1.c b/kernel/cgroup/cgroup-v1.c
index c487ffef6652..18064de0a883 100644
--- a/kernel/cgroup/cgroup-v1.c
+++ b/kernel/cgroup/cgroup-v1.c
@@ -1263,6 +1263,73 @@ int cgroup1_get_tree(struct fs_context *fc)
 	return ret;
 }
 
+/**
+ * task_cgroup_id_within_hierarchy - Retrieves the associated cgroup ID from
+ * a task within a specific cgroup1 hierarchy.
+ * @task: The task to be tested
+ * @hierarchy_id: The hierarchy ID of a cgroup1
+ *
+ * We limit it to cgroup1 only.
+ */
+u64 task_cgroup1_id_within_hierarchy(struct task_struct *tsk, int hierarchy_id)
+{
+	struct cgroup_root *root;
+	struct cgroup *cgrp;
+	u64 cgid = 0;
+
+	spin_lock_irq(&css_set_lock);
+	list_for_each_entry(root, &cgroup_roots, root_list) {
+		/* cgroup1 only*/
+		if (root == &cgrp_dfl_root)
+			continue;
+		if (root->hierarchy_id != hierarchy_id)
+			continue;
+		cgrp = task_cgroup_from_root(tsk, root);
+		WARN_ON_ONCE(!cgrp);
+		cgid = cgroup_id(cgrp);
+		break;
+	}
+	spin_unlock_irq(&css_set_lock);
+	return cgid;
+}
+
+/**
+ * task_ancestor_cgroup_id_within_hierarchy - Retrieves the associated ancestor
+ * cgroup ID from a task within a specific cgroup1 hierarchy.
+ * @task: The task to be tested
+ * @hierarchy_id: The hierarchy ID of a cgroup1
+ * @ancestor_level: level of ancestor to find starting from root
+ *
+ * We limit it to cgroup1 only.
+ */
+u64 task_ancestor_cgroup1_id_within_hierarchy(struct task_struct *tsk, int hierarchy_id,
+					      int ancestor_level)
+{
+	struct cgroup *cgrp, *ancestor;
+	struct cgroup_root *root;
+	u64 cgid = 0;
+
+	spin_lock_irq(&css_set_lock);
+	list_for_each_entry(root, &cgroup_roots, root_list) {
+		/* cgroup1 only*/
+		if (root == &cgrp_dfl_root)
+			continue;
+		if (root->hierarchy_id != hierarchy_id)
+			continue;
+
+		cgrp = task_cgroup_from_root(tsk, root);
+		WARN_ON_ONCE(!cgrp);
+		ancestor = cgroup_ancestor(cgrp, ancestor_level);
+		if (!ancestor)
+			break;
+
+		cgid = cgroup_id(ancestor);
+		break;
+	}
+	spin_unlock_irq(&css_set_lock);
+	return cgid;
+}
+
 static int __init cgroup1_wq_init(void)
 {
 	/*
-- 
2.30.1 (Apple Git-130)