[PATCH RFC v2 0/6] mm/memcg: move memcgid refcount to objcg to unpin dying memcgs

[PATCH RFC v2 0/6] mm/memcg: move memcgid refcount to objcg to unpin dying memcgs

[Bingfang Guo via B4 Relay]

Although the dying memcg problem caused by LRU pages is fixed, I can
still see the issue on some workloads that use shmem after those pages
are swapped out.  This series binds the memcgid refcount to objcgs so
dying memcgs can be freed normally in this case.

The memcg private ID identifies memcgs for objects that can outlive the
cgroup itself: swap entries and workingset shadows.  Today the ID's
refcount is embedded in the css, and every outstanding ID reference
(mostly swap entries) pins the css, keeping the entire memcg alive,
which brings a problem: A swapped-out page holds a memcgid reference
that pins the css, so the memcg cannot be freed until the page is
swapped back in and charged back to its online parent.

The work done by Muchun Song and Qi Zheng already charges folios to the
objcg, which is reparented to its parent when the memcg offlines.  This
series applies similar idea to the memcg private ID: the ID's refcount
moves from the css into the objcg, and the memcgid xarray holds a
reference to an objcg instead of pinning the css.  When the memcg
offlines, the objcg is reparented and any remaining memcgid references
resolve to the ancestor, so swapped-out pages no longer pin the dying
memcg and get the online parent naturally on swapin.

Patches 1-4 are self-contained preparatory cleanups:
  - add a helper to drop the online-state reference on offlining
  - take the memcgid reference only after swap charging succeeds
  - make mem_cgroup_private_id_put() take the ID instead of the memcg
  - return the memcg from mem_cgroup_private_id_put() for uncharging

Patch 5 actually moves the refcount to objcgs, and patch 6 makes the
remaining callers (list_lru, workingset) to filter out memcgs that have
already been reparented.

Changes in v2:
- Rework the logic to try to keep changes small.
- Change order of the commits to make it cleaner.
- Fix problems reported by sashiko.
- Reparent to mm-unstable.

RFC v1 contains some scripts and codes for reproducing the problem and
testing the fix:

https://lore.kernel.org/linux-mm/20260813-memcgid-objcg-v1-0-83d21c685b77@tencent.com/T/#m0e8ce07568e16239ba113c803f27caf308b299e7

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
Bingfang Guo (6):
      mm/memcg: add a helper to kill the memcgid on offlining
      mm/memcg: get memcgid reference only after swap charging success
      mm/memcg: pass the id itself instead of memcg for putting ID
      mm/memcg: return the memcg when putting memcgid
      mm/memcg: move memcg private ID refcount to objcg
      mm/memcg: filter out reparented memcgs got from memcgid

 include/linux/memcontrol.h |  15 +++---
 mm/list_lru.c              |   2 +-
 mm/memcontrol.c            | 117 +++++++++++++++++++++++++++++++++++----------
 mm/workingset.c            |   2 +-
 4 files changed, 100 insertions(+), 36 deletions(-)
---
base-commit: 88297631d4d42f6004cb39c0ba3da7d2d10a616f
change-id: 20260827-bingfangguo-memcgid-rework-938e25ecaba2

Best regards,
-- 
Bingfang Guo <bingfangguo@tencent.com>

[PATCH RFC v2 1/6] mm/memcg: add a helper to kill the memcgid on offlining

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

The online state holds one reference on the memcg private id; it is
released when the memcg goes offline. Put that release behind a small
helper, mem_cgroup_private_id_kill(), so css_offline() reads clearly
and the id refcount helpers stay symmetric: get_online() / kill() /
put().

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol.c | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 256b68ffca70e..31cec9dde55f0 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4048,6 +4048,11 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
 	}
 }
 
+static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
+{
+	mem_cgroup_private_id_put(memcg, 1);
+}
+
 struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
 {
 	while (!refcount_add_not_zero(n, &memcg->id.ref)) {
@@ -4385,7 +4390,7 @@ static void mem_cgroup_css_offline(struct cgroup_subsys_state *css)
 
 	drain_all_stock(memcg);
 
-	mem_cgroup_private_id_put(memcg, 1);
+	mem_cgroup_private_id_kill(memcg);
 }
 
 static void mem_cgroup_css_released(struct cgroup_subsys_state *css)

-- 
2.43.7

[PATCH RFC v2 2/6] mm/memcg: get memcgid reference only after swap charging success

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

__mem_cgroup_try_charge_swap() pinned the memcg private id before the
swap counter was charged and had to undo the pin on the failure path.
Hold RCU lock for an extended period (which should be fine,
__memcg1_swapout() does this as well) so concurrent memcg release can
be avoided, and take the id reference to its online parent only after
charging has succeeded.

The failure path is now a plain return, and the id is only pinned for
entries that actually end up charged to swap.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol.c | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 31cec9dde55f0..ecb4fb07d7735 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -5755,6 +5755,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
 	struct page_counter *counter;
 	struct mem_cgroup *memcg;
 	struct obj_cgroup *objcg;
+	unsigned short memcgid;
 
 	if (do_memsw_account())
 		return 0;
@@ -5772,22 +5773,24 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
 		return 0;
 	}
 
-	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
-	/* memcg is pined by memcg ID. */
-	rcu_read_unlock();
+	while (memcg_is_dying(memcg))
+		memcg = parent_mem_cgroup(memcg);
 
 	if (!mem_cgroup_is_root(memcg) &&
 	    !page_counter_try_charge(&memcg->swap, nr_pages, &counter)) {
 		memcg_memory_event(memcg, MEMCG_SWAP_MAX);
 		memcg_memory_event(memcg, MEMCG_SWAP_FAIL);
-		mem_cgroup_private_id_put(memcg, nr_pages);
+		rcu_read_unlock();
 		return -ENOMEM;
 	}
 	mod_memcg_state(memcg, MEMCG_SWAP, nr_pages);
 
+	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
+	memcgid = mem_cgroup_private_id(memcg);
+	rcu_read_unlock();
+
 	ci = swap_cluster_get_and_lock(folio);
-	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_pages,
-			  mem_cgroup_private_id(memcg));
+	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_pages, memcgid);
 	swap_cluster_unlock(ci);
 
 	return 0;

-- 
2.43.7

Re: [PATCH RFC v2 2/6] mm/memcg: get memcgid reference only after swap charging success

[Bingfang Guo]

Forwarding reviews from Sashiko:
> > From: Bingfang Guo <bingfangguo@tencent.com>
> > 
> > __mem_cgroup_try_charge_swap() pinned the memcg private id before the
> > swap counter was charged and had to undo the pin on the failure path.
> > Hold RCU lock for an extended period (which should be fine,
> > __memcg1_swapout() does this as well) so concurrent memcg release can
> > be avoided, and take the id reference to its online parent only after
> > charging has succeeded.
> > 
> > The failure path is now a plain return, and the id is only pinned for
> > entries that actually end up charged to swap.
> > 
> > Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> > ---
> >  mm/memcontrol.c | 15 +++++++++------
> >  1 file changed, 9 insertions(+), 6 deletions(-)
> > 
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index 31cec9dde55f0..ecb4fb07d7735 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> > @@ -5755,6 +5755,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
> >  	struct page_counter *counter;
> >  	struct mem_cgroup *memcg;
> >  	struct obj_cgroup *objcg;
> > +	unsigned short memcgid;
> >  
> >  	if (do_memsw_account())
> >  		return 0;
> > @@ -5772,22 +5773,24 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
> >  		return 0;
> >  	}
> >  
> > -	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
> > -	/* memcg is pined by memcg ID. */
> > -	rcu_read_unlock();
> > +	while (memcg_is_dying(memcg))
> > +		memcg = parent_mem_cgroup(memcg);
> >  
> >  	if (!mem_cgroup_is_root(memcg) &&
> >  	    !page_counter_try_charge(&memcg->swap, nr_pages, &counter)) {
> >  		memcg_memory_event(memcg, MEMCG_SWAP_MAX);
> >  		memcg_memory_event(memcg, MEMCG_SWAP_FAIL);
> > -		mem_cgroup_private_id_put(memcg, nr_pages);
> > +		rcu_read_unlock();
> >  		return -ENOMEM;
> >  	}
> >  	mod_memcg_state(memcg, MEMCG_SWAP, nr_pages);
> >  
> > +	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
> > +	memcgid = mem_cgroup_private_id(memcg);
> > +	rcu_read_unlock();
> 
> Does this introduce a Time-of-Check to Time-of-Use (TOCTOU) race?
> If a concurrent cgroup rmdir (triggering css_offline) occurs exactly between
> page_counter_try_charge() and mem_cgroup_private_id_get_online(), could it
> drop the child's ID refcount to zero?
> The subsequent ID lookup would climb to the parent, recording the parent's ID
> in the swap cache while the child's counter was already charged. This could
> leave the child's swap page_counter and MEMCG_SWAP stat permanently charged
> when destroyed.
> Additionally, could the parent cgroup's MEMCG_SWAP stat receive a spurious
> uncharge upon swap-in, causing it to permanently under-report swap usage and
> corrupting the memory.stat output?
> 

Just as I replied in patch 4, offlining shouldn't happen with rcu read lock
held, so actually this won't happen and we should always get the memcg itself...
And even if we are getting the ancestor it should be fine, uncharging can be
done for its online ancestor and stat will be consistent.

Maybe we can release the rcu lock earlier or remove the retry logic when we get
the memcgid reference?

> > +
> >  	ci = swap_cluster_get_and_lock(folio);
> > -	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_pages,
> > -			  mem_cgroup_private_id(memcg));
> > +	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_pages, memcgid);
> >  	swap_cluster_unlock(ci);
> >  
> >  	return 0;
> > 
> > -- 
> > 2.43.7
> > 
> > 

[PATCH RFC v2 3/6] mm/memcg: pass the id itself instead of memcg for putting ID

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

Swap uncharge knows the memcg only by its private id, and the id can
outlive the memcg it used to belong to after we rebind memcgid to
objcgs.  Make mem_cgroup_private_id_put() take the id and resolve the
memcg containing the refcount internally, and keep the underlying
__mem_cgroup_private_id_put() for the offline path that still holds a
memcg pointer.

In the uncharge path, the memcg pointer will have to be read from the
xarray twice, but we'll fix that later by returning the memcg from the
put path, so the uncharge path can obtain a reference in the same step.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol.c | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index ecb4fb07d7735..048c9bb0fad79 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4038,7 +4038,7 @@ static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
 	}
 }
 
-static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n)
+static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n)
 {
 	if (refcount_sub_and_test(n, &memcg->id.ref)) {
 		mem_cgroup_private_id_remove(memcg);
@@ -4048,9 +4048,19 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
 	}
 }
 
+static void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
+{
+	struct mem_cgroup *memcg;
+
+	rcu_read_lock();
+	memcg = mem_cgroup_from_private_id(id);
+	__mem_cgroup_private_id_put(memcg, n);
+	rcu_read_unlock();
+}
+
 static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
 {
-	mem_cgroup_private_id_put(memcg, 1);
+	__mem_cgroup_private_id_put(memcg, 1);
 }
 
 struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
@@ -5815,9 +5825,10 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 				page_counter_uncharge(&memcg->swap, nr_pages);
 		}
 		mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
-		mem_cgroup_private_id_put(memcg, nr_pages);
+		mem_cgroup_private_id_put(id, nr_pages);
 	}
 	rcu_read_unlock();
+
 }
 
 long mem_cgroup_get_nr_swap_pages(struct mem_cgroup *memcg)

-- 
2.43.7

[PATCH RFC v2 4/6] mm/memcg: return the memcg when putting memcgid

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

__mem_cgroup_uncharge_swap() needs both the memcg and the id refcount
drop.  Right now it looks the memcg up by id, uncharges it, then looks
it up again inside mem_cgroup_private_id_put() to drop the reference.

Make mem_cgroup_private_id_put() resolve the id once, drop the
reference, and return the nearest online memcg with a reference held for
the caller.  __mem_cgroup_uncharge_swap() then uses that memcg directly
and drops the reference after uncharging, avoiding the second xarray
lookup.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol.c | 20 +++++++++++++++++---
 1 file changed, 17 insertions(+), 3 deletions(-)

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 048c9bb0fad79..f0503a1e5492d 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4048,14 +4048,28 @@ static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n
 	}
 }
 
-static void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
+/**
+ * mem_cgroup_private_id_put - put memcgid and get the nearest online memcg
+ * @id: the memcg private id got from mem_cgroup_id_get_online
+ * @n: count of references to put
+ */
+static struct mem_cgroup *mem_cgroup_private_id_put(unsigned short id, unsigned int n)
 {
 	struct mem_cgroup *memcg;
 
 	rcu_read_lock();
 	memcg = mem_cgroup_from_private_id(id);
+	if (!memcg)
+		goto out;
+
 	__mem_cgroup_private_id_put(memcg, n);
+
+	while (memcg_is_dying(memcg) || !mem_cgroup_tryget(memcg))
+		memcg = parent_mem_cgroup(memcg);
+
+out:
 	rcu_read_unlock();
+	return memcg;
 }
 
 static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
@@ -5816,7 +5830,7 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 	struct mem_cgroup *memcg;
 
 	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
+	memcg = mem_cgroup_private_id_put(id, nr_pages);
 	if (memcg) {
 		if (!mem_cgroup_is_root(memcg)) {
 			if (do_memsw_account())
@@ -5825,10 +5839,10 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 				page_counter_uncharge(&memcg->swap, nr_pages);
 		}
 		mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
-		mem_cgroup_private_id_put(id, nr_pages);
 	}
 	rcu_read_unlock();
 
+	mem_cgroup_put(memcg);
 }
 
 long mem_cgroup_get_nr_swap_pages(struct mem_cgroup *memcg)

-- 
2.43.7

Re: [PATCH RFC v2 4/6] mm/memcg: return the memcg when putting memcgid

[Bingfang Guo]

Forwarding review from Sashiko:
> > From: Bingfang Guo <bingfangguo@tencent.com>
> > 
> > __mem_cgroup_uncharge_swap() needs both the memcg and the id refcount
> > drop.  Right now it looks the memcg up by id, uncharges it, then looks
> > it up again inside mem_cgroup_private_id_put() to drop the reference.
> > 
> > Make mem_cgroup_private_id_put() resolve the id once, drop the
> > reference, and return the nearest online memcg with a reference held for
> > the caller.  __mem_cgroup_uncharge_swap() then uses that memcg directly
> > and drops the reference after uncharging, avoiding the second xarray
> > lookup.
> > 
> > Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> > ---
> >  mm/memcontrol.c | 20 +++++++++++++++++---
> >  1 file changed, 17 insertions(+), 3 deletions(-)
> > 
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index 048c9bb0fad79..f0503a1e5492d 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> > @@ -4048,14 +4048,28 @@ static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n
> >  	}
> >  }
> >  
> > -static void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
> > +/**
> > + * mem_cgroup_private_id_put - put memcgid and get the nearest online memcg
> > + * @id: the memcg private id got from mem_cgroup_id_get_online
> > + * @n: count of references to put
> > + */
> > +static struct mem_cgroup *mem_cgroup_private_id_put(unsigned short id, unsigned int n)
> >  {
> >  	struct mem_cgroup *memcg;
> >  
> >  	rcu_read_lock();
> >  	memcg = mem_cgroup_from_private_id(id);
> > +	if (!memcg)
> > +		goto out;
> > +
> >  	__mem_cgroup_private_id_put(memcg, n);
> > +
> > +	while (memcg_is_dying(memcg) || !mem_cgroup_tryget(memcg))
> > +		memcg = parent_mem_cgroup(memcg);
> 
> Does returning the nearest online ancestor instead of the originally recorded
> dying memcg cause issues for callers that need to uncharge the exact memcg
> that was charged?

This seems to be a false positive?

The page counters are charged recursively to the root. When we get some ancestor
of the dying (or dead) memcg, the dead ones will never be accessed again so just
uncharging the online parents should be enough? And these counters are also left
untouched when the memcg is offline and pages are reparented to the ancestors
together with the objcg.  So I think this should be fine.

> > +
> > +out:
> >  	rcu_read_unlock();
> > +	return memcg;
> >  }
> >  
> >  static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
> > @@ -5816,7 +5830,7 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
> >  	struct mem_cgroup *memcg;
> >  
> >  	rcu_read_lock();
> > -	memcg = mem_cgroup_from_private_id(id);
> > +	memcg = mem_cgroup_private_id_put(id, nr_pages);
> >  	if (memcg) {
> >  		if (!mem_cgroup_is_root(memcg)) {
> >  			if (do_memsw_account())
> > @@ -5825,10 +5839,10 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
> >  				page_counter_uncharge(&memcg->swap, nr_pages);
> >  		}
> >  		mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
> 
> When a swap entry is charged while a memcg is online, the original memcg's
> local MEMCG_SWAP per-cpu stat is incremented and its swap page_counter is
> charged.
> If this memcg is subsequently offlined, it becomes dying. When the swap is
> eventually freed, mem_cgroup_private_id_put() now returns the nearest
> online ancestor instead of the original dying memcg.
> Will applying the uncharge to this ancestor permanently leak the counter in
> the original dying memcg?
> Furthermore, does this erroneously decrement the ancestor's local MEMCG_SWAP
> stat, causing a permanent stat underflow since the ancestor did not receive
> the corresponding local charge?
> 

We are holding rcu read lock and checked !memcg_is_dying(memcg) in
mem_cgroup_private_id_put(). When the css is being offlined, CSS_DYING is first
set and then mem_cgroup_css_offline() is called behind call_rcu, so no
concurrent offlining of this memcg can take place here.

For cgroup v1, MEMCG_SWAP stats will be reparented so eventually we will get
consistent result. And for cgroup v2, the stat is not visible (and therefore not
reparented) so no problem as well.

> > -		mem_cgroup_private_id_put(id, nr_pages);
> >  	}
> >  	rcu_read_unlock();
> >  
> > +	mem_cgroup_put(memcg);
> >  }
> >  
> >  long mem_cgroup_get_nr_swap_pages(struct mem_cgroup *memcg)
> > 
> > -- 
> > 2.43.7
> > 
> > 

[PATCH RFC v2 5/6] mm/memcg: move memcg private ID refcount to objcg

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

The memcg private ID is used by objects that can't afford storing a
whole pointer and can outlive memcgs to track the memcg (notably swap
entries). The current design holds a refcount to the css, preventing the
memcg from being freed.  This patch unbinds the lifetime of memcgid from
the memcg so it can be freed.

The idea is to move the refcount of memcgid to one of the memcg's objcg
and hold a pointer and a reference to the objcg in the global memcgid
xarray. No more css reference to the memcg so swapped out pages no
longer pin the dying memcg.

When retrieving the online memcg from the id, the objcg is taken out of
the xarray, and resolves to the online parent memcg naturally, which is
exactly what is expected in normal swapin folio charging path.  For swap
uncharging, the objcg is used for putting the id refcount and getting
the online ancestor in one go.

The exceptions are list_lru and workingset recent test, which require
exact the memcg the id points to. Those callers are fixed in the next
patch.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 include/linux/memcontrol.h | 15 +++++-----
 mm/memcontrol.c            | 71 +++++++++++++++++++++++++++++++---------------
 2 files changed, 55 insertions(+), 31 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index f227348a3f24a..eafc817ff244c 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -66,11 +66,6 @@ struct mem_cgroup_reclaim_cookie {
 
 #define MEM_CGROUP_ID_SHIFT	16
 
-struct mem_cgroup_private_id {
-	int id;
-	refcount_t ref;
-};
-
 struct memcg_vmstats_percpu;
 struct memcg1_events_percpu;
 struct memcg_vmstats;
@@ -173,6 +168,7 @@ struct obj_cgroup {
 	struct percpu_ref refcnt;
 	struct mem_cgroup *memcg;
 	atomic_t nr_charged_bytes;
+	refcount_t memcgid_ref;
 	union {
 		struct list_head list; /* protected by objcg_lock */
 		struct rcu_head rcu;
@@ -189,8 +185,8 @@ struct obj_cgroup {
 struct mem_cgroup {
 	struct cgroup_subsys_state css;
 
-	/* Private memcg ID. Used to ID objects that outlive the cgroup */
-	struct mem_cgroup_private_id id;
+	/* The objcg holding private memcg ID. */
+	struct obj_cgroup *id_objcg;
 
 	/* Accounted resources */
 	struct page_counter memory;		/* Both v1 & v2 */
@@ -255,6 +251,9 @@ struct mem_cgroup {
 #endif
 	int kmemcg_id;
 
+	/* Private memcg ID. Used to ID objects that outlive the cgroup */
+	int id;
+
 #ifdef CONFIG_CGROUP_WRITEBACK
 	struct list_head cgwb_list;
 #endif
@@ -810,7 +809,7 @@ static inline unsigned short mem_cgroup_private_id(struct mem_cgroup *memcg)
 	if (mem_cgroup_disabled())
 		return 0;
 
-	return memcg->id.id;
+	return memcg->id;
 }
 struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id);
 
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index f0503a1e5492d..38d2b00657a7a 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -3773,7 +3773,7 @@ static void memcg_online_kmem(struct mem_cgroup *memcg)
 
 	static_branch_enable(&memcg_kmem_online_key);
 
-	memcg->kmemcg_id = memcg->id.id;
+	memcg->kmemcg_id = memcg->id;
 }
 
 static void memcg_offline_kmem(struct mem_cgroup *memcg)
@@ -4032,19 +4032,23 @@ static DEFINE_XARRAY_ALLOC1(mem_cgroup_private_ids);
 
 static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
 {
-	if (memcg->id.id > 0) {
-		xa_erase(&mem_cgroup_private_ids, memcg->id.id);
-		memcg->id.id = 0;
+	if (memcg->id > 0) {
+		xa_erase(&mem_cgroup_private_ids, memcg->id);
+		memcg->id = 0;
 	}
 }
 
-static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n)
+static void __mem_cgroup_private_id_put(struct obj_cgroup *objcg,
+		unsigned short id, unsigned int n)
 {
-	if (refcount_sub_and_test(n, &memcg->id.ref)) {
-		mem_cgroup_private_id_remove(memcg);
+	struct obj_cgroup *objcg_free;
 
-		/* Memcg ID pins CSS */
-		css_put(&memcg->css);
+	if (refcount_sub_and_test(n, &objcg->memcgid_ref)) {
+		objcg_free = xa_erase(&mem_cgroup_private_ids, id);
+		VM_WARN_ON(objcg_free != objcg);
+
+		/* Memcg ID pins the objcg */
+		obj_cgroup_put(objcg);
 	}
 }
 
@@ -4055,18 +4059,20 @@ static void __mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n
  */
 static struct mem_cgroup *mem_cgroup_private_id_put(unsigned short id, unsigned int n)
 {
-	struct mem_cgroup *memcg;
+	struct mem_cgroup *memcg = NULL;
+	struct obj_cgroup *objcg;
 
 	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
-	if (!memcg)
+	objcg = xa_load(&mem_cgroup_private_ids, id);
+	if (unlikely(!objcg))
 		goto out;
 
-	__mem_cgroup_private_id_put(memcg, n);
-
+	memcg = obj_cgroup_memcg(objcg);
 	while (memcg_is_dying(memcg) || !mem_cgroup_tryget(memcg))
 		memcg = parent_mem_cgroup(memcg);
 
+	__mem_cgroup_private_id_put(objcg, id, n);
+
 out:
 	rcu_read_unlock();
 	return memcg;
@@ -4074,12 +4080,17 @@ static struct mem_cgroup *mem_cgroup_private_id_put(unsigned short id, unsigned
 
 static void mem_cgroup_private_id_kill(struct mem_cgroup *memcg)
 {
-	__mem_cgroup_private_id_put(memcg, 1);
+	__mem_cgroup_private_id_put(memcg->id_objcg, memcg->id, 1);
 }
 
 struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
 {
-	while (!refcount_add_not_zero(n, &memcg->id.ref)) {
+	struct obj_cgroup *objcg;
+
+	rcu_read_lock();
+	objcg = memcg->id_objcg;
+
+	while (!refcount_add_not_zero(n, &objcg->memcgid_ref)) {
 		/*
 		 * The root cgroup cannot be destroyed, so it's refcount must
 		 * always be >= 1.
@@ -4089,7 +4100,10 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
 			break;
 		}
 		memcg = parent_mem_cgroup(memcg);
+		objcg = memcg->id_objcg;
 	}
+
+	rcu_read_unlock();
 	return memcg;
 }
 
@@ -4101,8 +4115,14 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
  */
 struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id)
 {
+	struct obj_cgroup *objcg;
 	WARN_ON_ONCE(!rcu_read_lock_held());
-	return xa_load(&mem_cgroup_private_ids, id);
+
+	objcg = xa_load(&mem_cgroup_private_ids, id);
+	if (!objcg)
+		return NULL;
+
+	return obj_cgroup_memcg(objcg);
 }
 
 struct mem_cgroup *mem_cgroup_get_from_id(u64 id)
@@ -4203,7 +4223,7 @@ static struct mem_cgroup *mem_cgroup_alloc(struct mem_cgroup *parent)
 	if (!memcg)
 		return ERR_PTR(-ENOMEM);
 
-	error = xa_alloc(&mem_cgroup_private_ids, &memcg->id.id, NULL,
+	error = xa_alloc(&mem_cgroup_private_ids, &memcg->id, NULL,
 			 XA_LIMIT(1, MEM_CGROUP_ID_MAX), GFP_KERNEL);
 	if (error)
 		goto fail;
@@ -4348,9 +4368,10 @@ static int mem_cgroup_css_online(struct cgroup_subsys_state *css)
 				   FLUSH_TIME);
 	lru_gen_online_memcg(memcg);
 
-	/* Online state pins memcg ID, memcg ID pins CSS */
-	refcount_set(&memcg->id.ref, 1);
-	css_get(css);
+	/* CSS pins memcg ID, memcg ID pins obj cgroup */
+	memcg->id_objcg = objcg;
+	refcount_set(&memcg->id_objcg->memcgid_ref, 1);
+	obj_cgroup_get(memcg->id_objcg);
 
 	/*
 	 * Ensure mem_cgroup_from_private_id() works once we're fully online.
@@ -4362,7 +4383,7 @@ static int mem_cgroup_css_online(struct cgroup_subsys_state *css)
 	 * publish it here at the end of onlining. This matches the
 	 * regular ID destruction during offlining.
 	 */
-	xa_store(&mem_cgroup_private_ids, memcg->id.id, memcg, GFP_KERNEL);
+	xa_store(&mem_cgroup_private_ids, memcg->id, memcg->id_objcg, GFP_KERNEL);
 
 	return 0;
 free_objcg:
@@ -5832,7 +5853,11 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 	rcu_read_lock();
 	memcg = mem_cgroup_private_id_put(id, nr_pages);
 	if (memcg) {
-		if (!mem_cgroup_is_root(memcg)) {
+		/*
+		 * If the memcg was offline and reparented to root, swap needs
+		 * uncharging as well. We check this by comparing the memcgid.
+		 */
+		if (!mem_cgroup_is_root(memcg) || id != mem_cgroup_private_id(memcg)) {
 			if (do_memsw_account())
 				page_counter_uncharge(&memcg->memsw, nr_pages);
 			else

-- 
2.43.7

[PATCH RFC v2 6/6] mm/memcg: filter out reparented memcgs got from memcgid

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

mem_cgroup_from_private_id() looks up the objcg that owns the id and
returns the objcg's current memcg.  After reparenting, that memcg can
differ from the one the id originally belonged to.

Callers such as the list lru and workingset refault code expect to get
back exactly the memcg referred to by the memcgid, so check that the
returned memcg still owns the id and return NULL otherwise. In this
case, make list_lru just skip and workingset tests use the root (as the
comments said in workingset.c:460).

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/list_lru.c   | 2 +-
 mm/memcontrol.c | 9 ++++++++-
 mm/workingset.c | 2 +-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/mm/list_lru.c b/mm/list_lru.c
index 36662d02ff963..bc956267f6835 100644
--- a/mm/list_lru.c
+++ b/mm/list_lru.c
@@ -428,7 +428,7 @@ unsigned long list_lru_walk_node(struct list_lru *lru, int nid,
 		xa_for_each(&lru->xa, index, mlru) {
 			rcu_read_lock();
 			memcg = mem_cgroup_from_private_id(index);
-			if (!mem_cgroup_tryget(memcg)) {
+			if (!memcg || !mem_cgroup_tryget(memcg)) {
 				rcu_read_unlock();
 				continue;
 			}
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 38d2b00657a7a..84dcab3acb8a2 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4112,17 +4112,24 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
  * @id: the memcg id to look up
  *
  * Caller must hold rcu_read_lock().
+ *
+ * @return: the memcg, or NULL if the memcg is already reparented.
  */
 struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id)
 {
 	struct obj_cgroup *objcg;
+	struct mem_cgroup *memcg;
 	WARN_ON_ONCE(!rcu_read_lock_held());
 
 	objcg = xa_load(&mem_cgroup_private_ids, id);
 	if (!objcg)
 		return NULL;
 
-	return obj_cgroup_memcg(objcg);
+	memcg = obj_cgroup_memcg(objcg);
+	if (mem_cgroup_private_id(memcg) != id)
+		return NULL;
+
+	return memcg;
 }
 
 struct mem_cgroup *mem_cgroup_get_from_id(u64 id)
diff --git a/mm/workingset.c b/mm/workingset.c
index 7ac2b88c80ae5..4e9b7b163c01e 100644
--- a/mm/workingset.c
+++ b/mm/workingset.c
@@ -470,7 +470,7 @@ bool workingset_test_recent(void *shadow, bool file, bool *workingset,
 	 * configurations instead.
 	 */
 	eviction_memcg = mem_cgroup_from_private_id(memcgid);
-	if (!mem_cgroup_tryget(eviction_memcg))
+	if (eviction_memcg && !mem_cgroup_tryget(eviction_memcg))
 		eviction_memcg = NULL;
 	rcu_read_unlock();
 

-- 
2.43.7

Re: [PATCH RFC v2 6/6] mm/memcg: filter out reparented memcgs got from memcgid

[Bingfang Guo]

Forwarding reviews from Sashiko:
>> From: Bingfang Guo <bingfangguo@tencent.com>
>> 
>> mem_cgroup_from_private_id() looks up the objcg that owns the id and
>> returns the objcg's current memcg.  After reparenting, that memcg can
>> differ from the one the id originally belonged to.
>> 
>> Callers such as the list lru and workingset refault code expect to get
>> back exactly the memcg referred to by the memcgid, so check that the
>> returned memcg still owns the id and return NULL otherwise. In this
>> case, make list_lru just skip and workingset tests use the root (as the
>> comments said in workingset.c:460).
>> 
>> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
>> ---
>> mm/list_lru.c   | 2 +-
>> mm/memcontrol.c | 9 ++++++++-
>> mm/workingset.c | 2 +-
>> 3 files changed, 10 insertions(+), 3 deletions(-)
>> 
>> diff --git a/mm/list_lru.c b/mm/list_lru.c
>> index 36662d02ff963..bc956267f6835 100644
>> --- a/mm/list_lru.c
>> +++ b/mm/list_lru.c
>> @@ -428,7 +428,7 @@ unsigned long list_lru_walk_node(struct list_lru *lru, int nid,
>> 		xa_for_each(&lru->xa, index, mlru) {
>> 			rcu_read_lock();
>> 			memcg = mem_cgroup_from_private_id(index);
>> -			if (!mem_cgroup_tryget(memcg)) {
>> +			if (!memcg || !mem_cgroup_tryget(memcg)) {
>> 				rcu_read_unlock();
>> 				continue;
>> 			}
>> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
>> index 38d2b00657a7a..84dcab3acb8a2 100644
>> --- a/mm/memcontrol.c
>> +++ b/mm/memcontrol.c
>> @@ -4112,17 +4112,24 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
>>  * @id: the memcg id to look up
>>  *
>>  * Caller must hold rcu_read_lock().
>> + *
>> + * @return: the memcg, or NULL if the memcg is already reparented.
>>  */
>> struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id)
>> {
>> 	struct obj_cgroup *objcg;
>> +	struct mem_cgroup *memcg;
>> 	WARN_ON_ONCE(!rcu_read_lock_held());
>> 
>> 	objcg = xa_load(&mem_cgroup_private_ids, id);
>> 	if (!objcg)
>> 		return NULL;
>> 
>> -	return obj_cgroup_memcg(objcg);
>> +	memcg = obj_cgroup_memcg(objcg);
>> +	if (mem_cgroup_private_id(memcg) != id)
>> +		return NULL;
>> +
>> +	return memcg;
>> }
> 
> Does returning NULL here for reparented memcgs break swapin accounting
> in callers like mem_cgroup_swapin_charge_folio()?
> If an offlined memcg is reparented, its original ID no longer matches the
> inheritor's ID. When mem_cgroup_swapin_charge_folio() calls this function:
> 	memcg = mem_cgroup_from_private_id(id);
> 	if (!memcg || !css_tryget_online(&memcg->css))
> 		memcg = get_mem_cgroup_from_mm(mm);
> Would receiving NULL force it to wrongfully charge the faulting task instead
> of the parent memcg that inherited the swap charge, creating a permanent
> accounting mismatch?
> 

The original behavior uses css_tryget_online(), which filters out offlined
memcgs as well so I suppose no behavioral change here.

And swap (or memsw) counter of the memcg from which the page was swapped out
should be uncharged in __mem_cgroup_uncharge_swap() when the swap entry is
removed.

I will put this in the commit message to explain it later.

>> 
>> struct mem_cgroup *mem_cgroup_get_from_id(u64 id)
>> diff --git a/mm/workingset.c b/mm/workingset.c
>> index 7ac2b88c80ae5..4e9b7b163c01e 100644
>> --- a/mm/workingset.c
>> +++ b/mm/workingset.c
>> @@ -470,7 +470,7 @@ bool workingset_test_recent(void *shadow, bool file, bool *workingset,
>> 	 * configurations instead.
>> 	 */
>> 	eviction_memcg = mem_cgroup_from_private_id(memcgid);
>> -	if (!mem_cgroup_tryget(eviction_memcg))
>> +	if (eviction_memcg && !mem_cgroup_tryget(eviction_memcg))
>> 		eviction_memcg = NULL;
>> 	rcu_read_unlock();
> 
> The commit message states that workingset tests use the root (as the
> comments said in workingset.c:460) when NULL is returned. However, is
> this what actually happens?
> Looking immediately below this diff in workingset_test_recent(), it seems
> receiving NULL triggers an early return:
> 	if (!mem_cgroup_disabled() && !eviction_memcg)
> 		return false;
> Does this abort the refault calculation entirely, treating the reparented
> shadow entry as cold and causing page cache thrashing?

I didn't make it clear in the commit message here.  For MGLRU the root is used.
But for classical LRU, it's just skipped because refault reading of other memcgs
should be meaningless in this case.

I will update the commit message to match the behavior...

> Additionally, can this cause an issue in MGLRU lru_gen_test_recent()?
> If mem_cgroup_from_private_id() returns NULL there:
> 	memcg = mem_cgroup_from_private_id(memcg_id);
> 	*lruvec = mem_cgroup_lruvec(memcg, pgdat);
> Would this silently fall back to the root memcg LRU vector instead of the
> correct parent memcg, incorrectly evaluating the child eviction token
> against the root completely unrelated max_seq?
> 

Yes, but the activation is speculative as documented so it should be okay? Are
there better ways to handle this?

But I do think it may be good to keep the two methods behaves consistently, to
either skip the check or get a random result for both two cases.

>> 
>> 
>> -- 
>> 2.43.7
>> 
>>