[PATCH 1/5 cgroup/for-6.16-fixes] cgroup: add exclusive css rstat init/exit api for base stats

[JP Kobryn <inwardvessel@gmail.com>]

It is known at call sites of css_rstat_init/exit() whether the given css is
associated with the base stats or a formal subsystem. Instead of passing
&cgrp->self or css to a shared API, offer an additional init/exit API pair
for exclusive use with the base stats.

Make use of this new API to make existing code more readable and also to
prepare for refactoring the complex css rstat init/exit logic.

Signed-off-by: JP Kobryn <inwardvessel@gmail.com>
---
 kernel/cgroup/cgroup-internal.h |  2 ++
 kernel/cgroup/cgroup.c          | 11 ++++++-----
 kernel/cgroup/rstat.c           | 27 +++++++++++++++++++++++----
 3 files changed, 31 insertions(+), 9 deletions(-)

diff --git a/kernel/cgroup/cgroup-internal.h b/kernel/cgroup/cgroup-internal.h
index b14e61c64a34..24fca840bf1c 100644
--- a/kernel/cgroup/cgroup-internal.h
+++ b/kernel/cgroup/cgroup-internal.h
@@ -270,6 +270,8 @@ int cgroup_task_count(const struct cgroup *cgrp);
 /*
  * rstat.c
  */
+int cgroup_rstat_base_init(struct cgroup *cgrp);
+void cgroup_rstat_base_exit(struct cgroup *cgrp);
 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);
diff --git a/kernel/cgroup/cgroup.c b/kernel/cgroup/cgroup.c
index a723b7dc6e4e..d684d33236b8 100644
--- a/kernel/cgroup/cgroup.c
+++ b/kernel/cgroup/cgroup.c
@@ -2149,7 +2149,7 @@ int cgroup_setup_root(struct cgroup_root *root, u16 ss_mask)
 	if (ret)
 		goto destroy_root;
 
-	ret = css_rstat_init(&root_cgrp->self);
+	ret = cgroup_rstat_base_init(root_cgrp);
 	if (ret)
 		goto destroy_root;
 
@@ -2190,7 +2190,7 @@ int cgroup_setup_root(struct cgroup_root *root, u16 ss_mask)
 	goto out;
 
 exit_stats:
-	css_rstat_exit(&root_cgrp->self);
+	cgroup_rstat_base_exit(root_cgrp);
 destroy_root:
 	kernfs_destroy_root(root->kf_root);
 	root->kf_root = NULL;
@@ -5446,13 +5446,13 @@ static void css_free_rwork_fn(struct work_struct *work)
 	struct cgroup *cgrp = css->cgroup;
 
 	percpu_ref_exit(&css->refcnt);
-	css_rstat_exit(css);
 
 	if (!css_is_self(css)) {
 		/* css free path */
 		struct cgroup_subsys_state *parent = css->parent;
 		int id = css->id;
 
+		css_rstat_exit(css);
 		ss->css_free(css);
 		cgroup_idr_remove(&ss->css_idr, id);
 		cgroup_put(cgrp);
@@ -5477,6 +5477,7 @@ static void css_free_rwork_fn(struct work_struct *work)
 			cgroup_put(cgroup_parent(cgrp));
 			kernfs_put(cgrp->kn);
 			psi_cgroup_free(cgrp);
+			cgroup_rstat_base_exit(cgrp);
 			kfree(cgrp);
 		} else {
 			/*
@@ -5742,7 +5743,7 @@ static struct cgroup *cgroup_create(struct cgroup *parent, const char *name,
 	 * Now that init_cgroup_housekeeping() has been called and cgrp->self
 	 * is setup, it is safe to perform rstat initialization on it.
 	 */
-	ret = css_rstat_init(&cgrp->self);
+	ret = cgroup_rstat_base_init(cgrp);
 	if (ret)
 		goto out_kernfs_remove;
 
@@ -5818,7 +5819,7 @@ static struct cgroup *cgroup_create(struct cgroup *parent, const char *name,
 out_psi_free:
 	psi_cgroup_free(cgrp);
 out_stat_exit:
-	css_rstat_exit(&cgrp->self);
+	cgroup_rstat_base_exit(cgrp);
 out_kernfs_remove:
 	kernfs_remove(cgrp->kn);
 out_cancel_ref:
diff --git a/kernel/cgroup/rstat.c b/kernel/cgroup/rstat.c
index cbeaa499a96a..8c3cb4a989ad 100644
--- a/kernel/cgroup/rstat.c
+++ b/kernel/cgroup/rstat.c
@@ -437,11 +437,10 @@ __bpf_kfunc void css_rstat_flush(struct cgroup_subsys_state *css)
 	}
 }
 
-int css_rstat_init(struct cgroup_subsys_state *css)
+static int __css_rstat_init(struct cgroup_subsys_state *css, bool is_self)
 {
 	struct cgroup *cgrp = css->cgroup;
 	int cpu;
-	bool is_self = css_is_self(css);
 
 	if (is_self) {
 		/* the root cgrp has rstat_base_cpu preallocated */
@@ -481,7 +480,7 @@ int css_rstat_init(struct cgroup_subsys_state *css)
 	return 0;
 }
 
-void css_rstat_exit(struct cgroup_subsys_state *css)
+static void __css_rstat_exit(struct cgroup_subsys_state *css, bool is_self)
 {
 	int cpu;
 
@@ -499,7 +498,7 @@ void css_rstat_exit(struct cgroup_subsys_state *css)
 			return;
 	}
 
-	if (css_is_self(css)) {
+	if (is_self) {
 		struct cgroup *cgrp = css->cgroup;
 
 		free_percpu(cgrp->rstat_base_cpu);
@@ -510,6 +509,26 @@ void css_rstat_exit(struct cgroup_subsys_state *css)
 	css->rstat_cpu = NULL;
 }
 
+int css_rstat_init(struct cgroup_subsys_state *css)
+{
+	return __css_rstat_init(css, false);
+}
+
+void css_rstat_exit(struct cgroup_subsys_state *css)
+{
+	return __css_rstat_exit(css, false);
+}
+
+int cgroup_rstat_base_init(struct cgroup *cgrp)
+{
+	return __css_rstat_init(&cgrp->self, true);
+}
+
+void cgroup_rstat_base_exit(struct cgroup *cgrp)
+{
+	__css_rstat_exit(&cgrp->self, true);
+}
+
 /**
  * ss_rstat_init - subsystem-specific rstat initialization
  * @ss: target subsystem
-- 
2.47.1