[PATCH RFC 4/5] memcg: add a helper to get online memcg from memcgid

From: Bingfang Guo <[email protected]>

When swapping in, the folio is charged back to the memcg that swapped
it out, or to one of its ancestors if that memcg is gone.
mem_cgroup_swapin_charge_folio() currently does the id lookup and the
css_tryget_online() check by hand under the RCU read lock.

The objcg behind the id is reparented to an online memcg when its own
memcg is destroyed, so looking the id up and taking a reference
through the objcg is enough to guarantee an online memcg.  Add
mem_cgroup_from_private_id_online() for that purpose and use it in
mem_cgroup_swapin_charge_folio(), dropping the RCU read lock usage.

Signed-off-by: Bingfang Guo <[email protected]>
---
 include/linux/memcontrol.h |  1 +
 mm/memcontrol.c            | 22 ++++++++++++++++++----
 2 files changed, 19 insertions(+), 4 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index c33ec7efad50b..fef8a1c4191b1 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -827,6 +827,7 @@ static inline unsigned short mem_cgroup_private_id(struct mem_cgroup *memcg)
 	return memcg->id;
 }
 struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id);
+struct mem_cgroup *mem_cgroup_from_private_id_online(unsigned short id);
 
 static inline u64 mem_cgroup_id(struct mem_cgroup *memcg)
 {
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 12545ca48194d..fdf2e0d1f17e5 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4012,6 +4012,22 @@ struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id)
 	return obj_cgroup_memcg(objcg);
 }
 
+/**
+ * mem_cgroup_from_private_id - look up an online memcg from a memcg id
+ *                              and get a reference.
+ * @id: the memcg id to look up
+ */
+struct mem_cgroup *mem_cgroup_from_private_id_online(unsigned short id)
+{
+	struct obj_cgroup *objcg;
+
+	objcg = xa_load(&mem_cgroup_private_ids, id);
+	if (!objcg)
+		return NULL;
+
+	return get_mem_cgroup_from_objcg(objcg);
+}
+
 static struct mem_cgroup *mem_cgroup_take_from_private_id(unsigned short id, unsigned int n)
 {
 	struct obj_cgroup *objcg;
@@ -5248,11 +5264,9 @@ int mem_cgroup_swapin_charge_folio(struct folio *folio, unsigned short id,
 	if (mem_cgroup_disabled())
 		return 0;
 
-	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
-	if (!memcg || !css_tryget_online(&memcg->css))
+	memcg = mem_cgroup_from_private_id_online(id);
+	if (!memcg)
 		memcg = get_mem_cgroup_from_mm(mm);
-	rcu_read_unlock();
 
 	ret = charge_memcg(folio, memcg, gfp);
 

-- 
2.43.7