[PATCH 0/4] memcg: move memcgid refcount to objcg to unpin dying memcgs

[PATCH 0/4] memcg: move memcgid refcount to objcg to unpin dying memcgs

[Bingfang Guo via B4 Relay]

Changes since RFC:
- Make the ID API symmetric: mem_cgroup_private_id_get() returns the ID
  and mem_cgroup_private_id_put() takes the ID. (Muchun Song)
- Drop the "take the ID only after the charge succeeded" and the
  dying-memcg walk-up in the swap charge path. (Muchun Song)
- Fold the list_lru / workingset / MGLRU NULL handling into the patch
  that moves the refcount, so the broken intermediate state is never
  introduced. (Muchun Song)
- Add obj_cgroup_from_private_id() and use it in the swap uncharge path,
  so the final put reuses the objcg that was already looked up instead
  of having the put helper return a memcg. (Muchun Song)
- Rename mem_cgroup->id_objcg to private_id_objcg, and make
  mem_cgroup->private_id an unsigned short. (Muchun Song)
- Require the caller to hold the RCU read lock in the ID helpers instead
  of taking it inside them. (Muchun Song)
- Add mem_cgroup_private_id_is_root() and decide whether a swap entry
  carries a counter charge from the ID's root status instead of the
  resolved memcg's, which keeps charge and uncharge symmetric after
  reparenting.
- Reorder and merge commits; the series is now four patches.
- Link to RFC v2: https://lore.kernel.org/r/20260901-bingfangguo-memcgid-rework-v2-0-8edd7f7a7251@tencent.com

Although the dying memcg problem caused by LRU pages is fixed, I can
still see many dying memcgs on some workloads that use shmem and those
pages are swapped out. For example, programs populating logs to tmpfs or
containers sharing data using shmem.  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.

Unbinding the ID from the memcg has three consequences the series has to
deal with:

  1. The ID stops pinning the memcg, so the paths that relied on the ID
     reference to keep the memcg alive have to hold the RCU read lock
     instead.  (Patch 1.)

  2. Charge and uncharge no longer necessarily happen on the same memcg:
     swapout charges the folio's memcg, while the slot free resolves the
     nearest live ancestor.  The counters are hierarchical, and the
     MEMCG_SWAP stat is either reparented at offline (v1) or not visible
     (v2), so nothing leaks.  But "does this entry carry a counter
     charge at all" can no longer be answered from the resolved memcg:
     root is skipped only because root's swap is not accounted, and a
     non-root ID whose memcg was reparented into root still carries a
     charge that must be released.  Patch 2 adds
     mem_cgroup_private_id_is_root() and makes all three swap paths
     decide on the ID's root status.

  3. An ID can now outlive the memcg it was allocated to, so the memcg
     resolved from an ID is not necessarily the memcg the ID was handed
     out for.  Callers that need exactly that memcg (list_lru, the
     workingset and MGLRU shadow tests) now get NULL and skip the entry,
     while the swap paths, which only need something to account to, get
     the nearest live ancestor.  (Patch 4.)

The series is now four patches:

  memcg: keep swap charging under RCU protection
  memcg: base swap charge accounting on memcgid root status
  memcg: manipulate memcg private ID references by ID
  memcg: move memcg private ID refcount to objcg

Only the last patch changes behavior.  The first three are preparation
that keeps today's semantics while the ID is still bound to the css.

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

** Changelogs **

Changes in RFC 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.
- Link to RFC v1: https://lore.kernel.org/r/20260813-memcgid-objcg-v1-0-83d21c685b77@tencent.com

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
Bingfang Guo (4):
      memcg: keep swap charging under RCU protection
      memcg: base swap charge accounting on memcgid root status
      memcg: manipulate memcg private ID references by ID
      memcg: move memcg private ID refcount to objcg

 include/linux/memcontrol.h |   9 ++--
 mm/list_lru.c              |   2 +-
 mm/memcontrol-v1.c         |  21 ++++-----
 mm/memcontrol-v1.h         |   8 +++-
 mm/memcontrol.c            | 111 ++++++++++++++++++++++++++++++++++-----------
 mm/workingset.c            |   5 +-
 6 files changed, 110 insertions(+), 46 deletions(-)
---
base-commit: 27e4e1835109ef599d72abe6c09711e0b1916033
change-id: 20260827-bingfangguo-memcgid-rework-938e25ecaba2

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

[PATCH 1/4] memcg: keep swap charging under RCU protection

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

This is a preparatory work for unbinding memcgid from memcg. No
functional change.

The swap charging path currently drops its RCU read lock after acquiring
a private ID reference. This is safe because the ID reference pins the
memcg's CSS.

Moving private ID references to objcgs will remove that lifetime
guarantee. Keep the RCU read lock held while accessing the memcg for
counter charging, statistics and failure handling. (This matches what
__memcg1_swapout() already does.).

Save the private ID before dropping the RCU read lock, and use the saved
value when recording the swap entry. The swap cluster locking remains
outside the RCU read-side critical section.

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

diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 791e536efaebe..72522ec827c9a 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -5954,6 +5954,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
 	struct page_counter *counter;
 	struct mem_cgroup *memcg;
 	struct obj_cgroup *objcg;
+	unsigned short private_id;
 
 	if (do_memsw_account())
 		return 0;
@@ -5973,20 +5974,21 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
 
 	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
 	/* memcg is pined by memcg ID. */
-	rcu_read_unlock();
+	private_id = mem_cgroup_private_id(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);
+	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, private_id);
 	swap_cluster_unlock(ci);
 
 	return 0;

-- 
2.43.7

[PATCH 2/4] memcg: base swap charge accounting on memcgid root status

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

A private ID currently pins its original memcg, so testing whether the
ID belongs to root is equivalent to testing whether the resolved memcg
is root.

That equivalence will no longer hold when IDs refer to objcgs. A
non-root ID may resolve to root after reparenting, but its swap entries
still carry counter charges inherited by root. Skipping their uncharge
based on the resolved memcg would leave those charges behind.

Use the private ID's root status to decide whether a swap entry carries
a counter charge. A root-ID entry carries none, while a non-root-ID
entry must release its charge even if its current accounting memcg has
become root.

This patch adds a new helper to check if the memcgid equals to that of
the root memcg.  For swap uncharging and v2 swap charging, simply decide
whether to charge/uncharge memsw or swap counter based on the swap
memcgid is root or not. For v1 swapout, don't recharge the memsw
counter, just cancel the charge if the swap memcg ID refers to the root
memcg.

This makes memory and swap charging have similar semantics: one relys on
the objcg's root status, and the other relys on the memcgid's root
status.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol-v1.c | 16 ++++++++--------
 mm/memcontrol-v1.h |  5 +++++
 mm/memcontrol.c    |  4 ++--
 3 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/mm/memcontrol-v1.c b/mm/memcontrol-v1.c
index bf2c7d53b01b1..ed015fdd95123 100644
--- a/mm/memcontrol-v1.c
+++ b/mm/memcontrol-v1.c
@@ -271,6 +271,7 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
 	struct mem_cgroup *memcg, *swap_memcg;
 	struct obj_cgroup *objcg;
 	unsigned int nr_entries;
+	unsigned short private_id;
 
 	VM_WARN_ON_ONCE_FOLIO(!folio_test_swapcache(folio), folio);
 	VM_WARN_ON_ONCE_FOLIO(!folio_test_locked(folio), folio);
@@ -293,25 +294,24 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
 	/*
 	 * In case the memcg owning these pages has been offlined and doesn't
 	 * have an ID allocated to it anymore, charge the closest online
-	 * ancestor for the swap instead and transfer the memory+swap charge.
+	 * ancestor for the swap instead and cancel the memory+swap charge
+	 * if the ID refers to the root memcg.
 	 */
 	nr_entries = folio_nr_pages(folio);
 	swap_memcg = mem_cgroup_private_id_get_online(memcg, nr_entries);
+	private_id = mem_cgroup_private_id(swap_memcg);
 	mod_memcg_state(swap_memcg, MEMCG_SWAP, nr_entries);
 
-	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries,
-			  mem_cgroup_private_id(swap_memcg));
+	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries, private_id);
 
 	folio_unqueue_deferred_split(folio);
 	folio->memcg_data = 0;
 
-	if (!obj_cgroup_is_root(objcg))
+	if (!obj_cgroup_is_root(objcg)) {
 		page_counter_uncharge(&memcg->memory, nr_entries);
 
-	if (memcg != swap_memcg) {
-		if (!mem_cgroup_is_root(swap_memcg))
-			page_counter_charge(&swap_memcg->memsw, nr_entries);
-		page_counter_uncharge(&memcg->memsw, nr_entries);
+		if (mem_cgroup_private_id_is_root(private_id))
+			page_counter_uncharge(&memcg->memsw, nr_entries);
 	}
 
 	/*
diff --git a/mm/memcontrol-v1.h b/mm/memcontrol-v1.h
index 2cd37e1792d79..23be2512702dc 100644
--- a/mm/memcontrol-v1.h
+++ b/mm/memcontrol-v1.h
@@ -22,6 +22,11 @@ void drain_all_stock(struct mem_cgroup *root_memcg);
 
 int memory_stat_show(struct seq_file *m, void *v);
 
+static inline bool mem_cgroup_private_id_is_root(unsigned short id)
+{
+	return id == mem_cgroup_private_id(root_mem_cgroup);
+}
+
 struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg,
 						    unsigned int n);
 
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 72522ec827c9a..bfe53e4392f09 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -5976,7 +5976,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
 	/* memcg is pined by memcg ID. */
 	private_id = mem_cgroup_private_id(memcg);
 
-	if (!mem_cgroup_is_root(memcg) &&
+	if (!mem_cgroup_private_id_is_root(private_id) &&
 	    !page_counter_try_charge(&memcg->swap, nr_pages, &counter)) {
 		memcg_memory_event(memcg, MEMCG_SWAP_MAX);
 		memcg_memory_event(memcg, MEMCG_SWAP_FAIL);
@@ -6006,7 +6006,7 @@ void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 	rcu_read_lock();
 	memcg = mem_cgroup_from_private_id(id);
 	if (memcg) {
-		if (!mem_cgroup_is_root(memcg)) {
+		if (!mem_cgroup_private_id_is_root(id)) {
 			if (do_memsw_account())
 				page_counter_uncharge(&memcg->memsw, nr_pages);
 			else

-- 
2.43.7

[PATCH 3/4] memcg: manipulate memcg private ID references by ID

[Bingfang Guo via B4 Relay]

From: Bingfang Guo <bingfangguo@tencent.com>

This is a preparatory work for moving memcgid from memcg to objcg.

Swap entries retain a private ID rather than a memcg pointer. Once
private ID references are moved to objcgs, the ID can also outlive the
memcg to which it was originally assigned.  So it's better to make the
get and put functions accept the ID itself instead of the memcg.

Rename mem_cgroup_private_id_get_online() to
mem_cgroup_private_id_get(), and make it return the ID only.  If the
memcg is already dying, the dying memcg will still be used for charging
and stats accounting in v2 swap charging path. But they are hierarchical
and will be reparented after offlining so it doesn't matter.

Make mem_cgroup_private_id_put() take the ID and resolve the reference
holder internally. Convert swap uncharge and charge rollback to release
the reference using that ID.  This introduces an extra xarray lookup for
now, which will be removed in the final patch.

Separate the online-state reference release into
mem_cgroup_private_id_kill(). The offline path already has the memcg
pointer and can call the underlying put helper directly.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 mm/memcontrol-v1.c |  7 +++----
 mm/memcontrol-v1.h |  3 +--
 mm/memcontrol.c    | 32 +++++++++++++++++++++++---------
 3 files changed, 27 insertions(+), 15 deletions(-)

diff --git a/mm/memcontrol-v1.c b/mm/memcontrol-v1.c
index ed015fdd95123..b7f2868885071 100644
--- a/mm/memcontrol-v1.c
+++ b/mm/memcontrol-v1.c
@@ -268,7 +268,7 @@ void memcg1_commit_charge(struct folio *folio, struct mem_cgroup *memcg)
  */
 void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
 {
-	struct mem_cgroup *memcg, *swap_memcg;
+	struct mem_cgroup *memcg;
 	struct obj_cgroup *objcg;
 	unsigned int nr_entries;
 	unsigned short private_id;
@@ -298,9 +298,8 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
 	 * if the ID refers to the root memcg.
 	 */
 	nr_entries = folio_nr_pages(folio);
-	swap_memcg = mem_cgroup_private_id_get_online(memcg, nr_entries);
-	private_id = mem_cgroup_private_id(swap_memcg);
-	mod_memcg_state(swap_memcg, MEMCG_SWAP, nr_entries);
+	private_id = mem_cgroup_private_id_get(memcg, nr_entries);
+	mod_memcg_state(memcg, MEMCG_SWAP, nr_entries);
 
 	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries, private_id);
 
diff --git a/mm/memcontrol-v1.h b/mm/memcontrol-v1.h
index 23be2512702dc..281425273ea97 100644
--- a/mm/memcontrol-v1.h
+++ b/mm/memcontrol-v1.h
@@ -27,8 +27,7 @@ static inline bool mem_cgroup_private_id_is_root(unsigned short id)
 	return id == mem_cgroup_private_id(root_mem_cgroup);
 }
 
-struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg,
-						    unsigned int n);
+unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n);
 
 void reparent_memcg_lruvec_state_local(struct mem_cgroup *memcg,
 				       struct mem_cgroup *parent, int idx);
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index bfe53e4392f09..ed44b3e7ac938 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4082,7 +4082,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->private_id_ref)) {
 		mem_cgroup_private_id_remove(memcg);
@@ -4092,7 +4092,22 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
 	}
 }
 
-struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
+static inline 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);
+}
+
+unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n)
 {
 	while (!refcount_add_not_zero(n, &memcg->private_id_ref)) {
 		/*
@@ -4105,7 +4120,8 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
 		}
 		memcg = parent_mem_cgroup(memcg);
 	}
-	return memcg;
+
+	return mem_cgroup_private_id(memcg);
 }
 
 /**
@@ -4430,7 +4446,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)
@@ -5972,15 +5988,13 @@ 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. */
-	private_id = mem_cgroup_private_id(memcg);
+	private_id = mem_cgroup_private_id_get(memcg, nr_pages);
 
 	if (!mem_cgroup_private_id_is_root(private_id) &&
 	    !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);
+		mem_cgroup_private_id_put(private_id, nr_pages);
 		rcu_read_unlock();
 		return -ENOMEM;
 	}
@@ -6013,7 +6027,7 @@ 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();
 }

-- 
2.43.7

[PATCH 4/4] 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. The objcg is stored in the global memcgid xarray instead
and used for retrieving the online memcg from it.  So swapped out pages
no longer pin the dying memcg.

After the change, a memcgid can refer to a non present memcg. To handle
this situation, when trying to get the original memcg from the id,
compare the memcgid passed in with that of the memcg, and return NULL to
indicate its death if they differ. NULL checks are added for
list_lru_walk_node(), workingset_test_recent() and lru_gen_test_recent()
to skip dead memcgs.

In the earlier patch, an extra xarray lookup was introduced in swap
uncharging path.  Now that we have the objcg pointer in the function,
the extra overhead can be removed by using it for putting directly.

Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
---
 include/linux/memcontrol.h |  9 +++---
 mm/list_lru.c              |  2 +-
 mm/memcontrol.c            | 81 ++++++++++++++++++++++++++++++++++------------
 mm/workingset.c            |  5 ++-
 4 files changed, 71 insertions(+), 26 deletions(-)

diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 46bf724cae7af..3fb18191cfbd1 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -180,6 +180,7 @@ struct obj_cgroup {
 	struct percpu_ref refcnt;
 	struct mem_cgroup *memcg;
 	atomic_t nr_charged_bytes;
+	refcount_t private_id_ref;
 	union {
 		struct list_head list; /* protected by objcg_lock */
 		struct rcu_head rcu;
@@ -225,9 +226,6 @@ struct mem_cgroup {
 	/* vmpressure notifications. Written on every reclaim iteration. */
 	struct vmpressure vmpressure;
 
-	/* Written on every swap charge and uncharge. */
-	refcount_t private_id_ref;
-
 #ifdef CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC
 	/* MEMCG_KMEM for nmi context */
 	atomic_t		kmem_stat;
@@ -324,8 +322,11 @@ struct mem_cgroup {
 	unsigned long zswap_max;
 #endif
 
+	/* The objcg holding private memcg ID. */
+	struct obj_cgroup *private_id_objcg;
+
 	/* Private memcg ID. Used to ID objects that outlive the cgroup */
-	int private_id;
+	unsigned short private_id;
 
 	int kmemcg_id;
 
diff --git a/mm/list_lru.c b/mm/list_lru.c
index 8a6dd0a489e12..7edd79113cc56 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 ed44b3e7ac938..22deee8312856 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4074,6 +4074,18 @@ static void memcg_wb_domain_size_changed(struct mem_cgroup *memcg)
 #define MEM_CGROUP_ID_MAX	((1UL << MEM_CGROUP_ID_SHIFT) - 1)
 static DEFINE_XARRAY_ALLOC1(mem_cgroup_private_ids);
 
+/**
+ * obj_cgroup_from_private_id - look up the objcg holding the memcg id.
+ * @id: the memcg id to look up
+ *
+ * Caller must hold rcu_read_lock().
+ */
+static inline struct obj_cgroup *obj_cgroup_from_private_id(unsigned short id)
+{
+	lockdep_assert_once(rcu_read_lock_held());
+	return xa_load(&mem_cgroup_private_ids, id);
+}
+
 static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
 {
 	if (memcg->private_id > 0) {
@@ -4082,34 +4094,43 @@ static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
 	}
 }
 
-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->private_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->private_id_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);
 	}
 }
 
 static inline void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
 {
-	struct mem_cgroup *memcg;
+	struct obj_cgroup *objcg;
 
 	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
-	__mem_cgroup_private_id_put(memcg, n);
+	objcg = obj_cgroup_from_private_id(id);
+	__mem_cgroup_private_id_put(objcg, id, 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->private_id_objcg, memcg->private_id, 1);
 }
 
 unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n)
 {
-	while (!refcount_add_not_zero(n, &memcg->private_id_ref)) {
+	struct obj_cgroup *objcg;
+	lockdep_assert_once(rcu_read_lock_held());
+
+	objcg = memcg->private_id_objcg;
+
+	while (!refcount_add_not_zero(n, &objcg->private_id_ref)) {
 		/*
 		 * The root cgroup cannot be destroyed, so it's refcount must
 		 * always be >= 1.
@@ -4119,6 +4140,7 @@ unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int
 			break;
 		}
 		memcg = parent_mem_cgroup(memcg);
+		objcg = memcg->private_id_objcg;
 	}
 
 	return mem_cgroup_private_id(memcg);
@@ -4129,11 +4151,24 @@ unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int
  * @id: the memcg id to look up
  *
  * Caller must hold rcu_read_lock().
+ *
+ * @return: the memcg, or NULL if the memcg referred to is already dead.
  */
 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());
-	return xa_load(&mem_cgroup_private_ids, id);
+
+	objcg = obj_cgroup_from_private_id(id);
+	if (!objcg)
+		return NULL;
+
+	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)
@@ -4228,18 +4263,21 @@ static struct mem_cgroup *mem_cgroup_alloc(struct mem_cgroup *parent)
 	struct mem_cgroup *memcg;
 	int node, cpu;
 	int __maybe_unused i;
+	unsigned int private_id;
 	long error;
 
 	memcg = kmem_cache_zalloc(memcg_cachep, GFP_KERNEL);
 	if (!memcg)
 		return ERR_PTR(-ENOMEM);
 
-	error = xa_alloc(&mem_cgroup_private_ids, &memcg->private_id, NULL,
+	error = xa_alloc(&mem_cgroup_private_ids, &private_id, NULL,
 			 XA_LIMIT(1, MEM_CGROUP_ID_MAX), GFP_KERNEL);
 	if (error)
 		goto fail;
 	error = -ENOMEM;
 
+	memcg->private_id = private_id;
+
 	memcg->vmstats = kzalloc_obj(struct memcg_vmstats, GFP_KERNEL_ACCOUNT);
 	if (!memcg->vmstats)
 		goto fail;
@@ -4380,9 +4418,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->private_id_ref, 1);
-	css_get(css);
+	/* CSS pins memcg ID, memcg ID pins obj cgroup */
+	memcg->private_id_objcg = objcg;
+	refcount_set(&memcg->private_id_objcg->private_id_ref, 1);
+	obj_cgroup_get(memcg->private_id_objcg);
 
 	/*
 	 * Ensure mem_cgroup_from_private_id() works once we're fully online.
@@ -4394,7 +4433,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->private_id, memcg, GFP_KERNEL);
+	xa_store(&mem_cgroup_private_ids, memcg->private_id, memcg->private_id_objcg, GFP_KERNEL);
 
 	return 0;
 free_objcg:
@@ -5832,8 +5871,6 @@ static void __init memcg_struct_check(void)
 				      memory_events_local);
 	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
 				      vmpressure);
-	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
-				      private_id_ref);
 #ifdef CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC
 	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
 				      kmem_stat);
@@ -5878,6 +5915,8 @@ static void __init memcg_struct_check(void)
 	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
 				      zswap_writeback);
 #endif
+	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
+				      private_id_objcg);
 	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
 				      private_id);
 	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
@@ -6015,10 +6054,12 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
  */
 void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
 {
+	struct obj_cgroup *objcg;
 	struct mem_cgroup *memcg;
 
 	rcu_read_lock();
-	memcg = mem_cgroup_from_private_id(id);
+	objcg = obj_cgroup_from_private_id(id);
+	memcg = obj_cgroup_memcg(objcg);
 	if (memcg) {
 		if (!mem_cgroup_private_id_is_root(id)) {
 			if (do_memsw_account())
@@ -6027,7 +6068,7 @@ 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);
+		__mem_cgroup_private_id_put(objcg, id, nr_pages);
 	}
 	rcu_read_unlock();
 }
diff --git a/mm/workingset.c b/mm/workingset.c
index 8412f4840ae35..7e4fbc5a786d6 100644
--- a/mm/workingset.c
+++ b/mm/workingset.c
@@ -281,6 +281,9 @@ static bool lru_gen_test_recent(void *shadow, struct lruvec **lruvec,
 	unpack_shadow(shadow, &memcg_id, &pgdat, token, workingset);
 
 	memcg = mem_cgroup_from_private_id(memcg_id);
+	if (!memcg)
+		return false;
+
 	*lruvec = mem_cgroup_lruvec(memcg, pgdat);
 
 	max_seq = READ_ONCE((*lruvec)->lrugen.max_seq);
@@ -470,7 +473,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 1/4] memcg: keep swap charging under RCU protection

[Muchun Song]

> On Sep 18, 2026, at 17:18, Bingfang Guo via B4 Relay <devnull+bingfangguo.tencent.com@kernel.org> wrote:
> 
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> This is a preparatory work for unbinding memcgid from memcg. No
> functional change.
> 
> The swap charging path currently drops its RCU read lock after acquiring
> a private ID reference. This is safe because the ID reference pins the
> memcg's CSS.
> 
> Moving private ID references to objcgs will remove that lifetime
> guarantee. Keep the RCU read lock held while accessing the memcg for
> counter charging, statistics and failure handling. (This matches what
> __memcg1_swapout() already does.).
> 
> Save the private ID before dropping the RCU read lock, and use the saved
> value when recording the swap entry. The swap cluster locking remains
> outside the RCU read-side critical section.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>

Acked-by: Muchun Song <muchun.song@linux.dev>

Thanks.

Re: [PATCH 1/4] memcg: keep swap charging under RCU protection

[Shakeel Butt]

On Fri, Sep 18, 2026 at 05:18:40PM +0800, Bingfang Guo via B4 Relay wrote:
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> This is a preparatory work for unbinding memcgid from memcg. No
> functional change.
> 
> The swap charging path currently drops its RCU read lock after acquiring
> a private ID reference. This is safe because the ID reference pins the
> memcg's CSS.
> 
> Moving private ID references to objcgs will remove that lifetime
> guarantee. Keep the RCU read lock held while accessing the memcg for
> counter charging, statistics and failure handling. (This matches what
> __memcg1_swapout() already does.).
> 
> Save the private ID before dropping the RCU read lock, and use the saved
> value when recording the swap entry. The swap cluster locking remains
> outside the RCU read-side critical section.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> ---
>  mm/memcontrol.c | 8 +++++---
>  1 file changed, 5 insertions(+), 3 deletions(-)
> 
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index 791e536efaebe..72522ec827c9a 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -5954,6 +5954,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
>  	struct page_counter *counter;
>  	struct mem_cgroup *memcg;
>  	struct obj_cgroup *objcg;
> +	unsigned short private_id;
>  
>  	if (do_memsw_account())
>  		return 0;
> @@ -5973,20 +5974,21 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
>  

If there is one more spin of this patch, I think scoped_guard(rcu) would be more
readable here, so please use that here.

>  	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
>  	/* memcg is pined by memcg ID. */
> -	rcu_read_unlock();
> +	private_id = mem_cgroup_private_id(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);
> +	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, private_id);
>  	swap_cluster_unlock(ci);
>  
>  	return 0;
> 
> -- 
> 2.43.7
> 
> 

Re: [PATCH 1/4] memcg: keep swap charging under RCU protection

[Bingfang Guo]

On Fri, Sep 18, 2026 at 09:46:47AM -0700, Shakeel Butt wrote:
> On Fri, Sep 18, 2026 at 05:18:40PM +0800, Bingfang Guo via B4 Relay wrote:
> > From: Bingfang Guo <bingfangguo@tencent.com>
> > 
> > This is a preparatory work for unbinding memcgid from memcg. No
> > functional change.
> > 
> > The swap charging path currently drops its RCU read lock after acquiring
> > a private ID reference. This is safe because the ID reference pins the
> > memcg's CSS.
> > 
> > Moving private ID references to objcgs will remove that lifetime
> > guarantee. Keep the RCU read lock held while accessing the memcg for
> > counter charging, statistics and failure handling. (This matches what
> > __memcg1_swapout() already does.).
> > 
> > Save the private ID before dropping the RCU read lock, and use the saved
> > value when recording the swap entry. The swap cluster locking remains
> > outside the RCU read-side critical section.
> > 
> > Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> > ---
> >  mm/memcontrol.c | 8 +++++---
> >  1 file changed, 5 insertions(+), 3 deletions(-)
> > 
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index 791e536efaebe..72522ec827c9a 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> > @@ -5954,6 +5954,7 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
> >  	struct page_counter *counter;
> >  	struct mem_cgroup *memcg;
> >  	struct obj_cgroup *objcg;
> > +	unsigned short private_id;
> >  
> >  	if (do_memsw_account())
> >  		return 0;
> > @@ -5973,20 +5974,21 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
> >  
> 
> If there is one more spin of this patch, I think scoped_guard(rcu) would be more
> readable here, so please use that here.
> 

Hi, Shakeel. Thanks for your review and suggestion!

It looks like a good idea. I will include that in the next version!

> >  	memcg = mem_cgroup_private_id_get_online(memcg, nr_pages);
> >  	/* memcg is pined by memcg ID. */
> > -	rcu_read_unlock();
> > +	private_id = mem_cgroup_private_id(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);
> > +	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, private_id);
> >  	swap_cluster_unlock(ci);
> >  
> >  	return 0;
> > 
> > -- 
> > 2.43.7
> > 
> > 

Re: [PATCH 3/4] memcg: manipulate memcg private ID references by ID

[Shakeel Butt]

On Fri, Sep 18, 2026 at 05:18:42PM +0800, Bingfang Guo via B4 Relay wrote:
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> This is a preparatory work for moving memcgid from memcg to objcg.
> 
> Swap entries retain a private ID rather than a memcg pointer. Once
> private ID references are moved to objcgs, the ID can also outlive the
> memcg to which it was originally assigned.  So it's better to make the
> get and put functions accept the ID itself instead of the memcg.
> 
> Rename mem_cgroup_private_id_get_online() to
> mem_cgroup_private_id_get(), and make it return the ID only.  If the
> memcg is already dying, the dying memcg will still be used for charging
> and stats accounting in v2 swap charging path. But they are hierarchical
> and will be reparented after offlining so it doesn't matter.
> 
> Make mem_cgroup_private_id_put() take the ID and resolve the reference
> holder internally. Convert swap uncharge and charge rollback to release
> the reference using that ID.  This introduces an extra xarray lookup for
> now, which will be removed in the final patch.
> 
> Separate the online-state reference release into
> mem_cgroup_private_id_kill(). The offline path already has the memcg
> pointer and can call the underlying put helper directly.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> ---
>  mm/memcontrol-v1.c |  7 +++----
>  mm/memcontrol-v1.h |  3 +--
>  mm/memcontrol.c    | 32 +++++++++++++++++++++++---------
>  3 files changed, 27 insertions(+), 15 deletions(-)
> 
> diff --git a/mm/memcontrol-v1.c b/mm/memcontrol-v1.c
> index ed015fdd95123..b7f2868885071 100644
> --- a/mm/memcontrol-v1.c
> +++ b/mm/memcontrol-v1.c
> @@ -268,7 +268,7 @@ void memcg1_commit_charge(struct folio *folio, struct mem_cgroup *memcg)
>   */
>  void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
>  {
> -	struct mem_cgroup *memcg, *swap_memcg;
> +	struct mem_cgroup *memcg;
>  	struct obj_cgroup *objcg;
>  	unsigned int nr_entries;
>  	unsigned short private_id;
> @@ -298,9 +298,8 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
>  	 * if the ID refers to the root memcg.
>  	 */
>  	nr_entries = folio_nr_pages(folio);
> -	swap_memcg = mem_cgroup_private_id_get_online(memcg, nr_entries);
> -	private_id = mem_cgroup_private_id(swap_memcg);
> -	mod_memcg_state(swap_memcg, MEMCG_SWAP, nr_entries);
> +	private_id = mem_cgroup_private_id_get(memcg, nr_entries);
> +	mod_memcg_state(memcg, MEMCG_SWAP, nr_entries);
>  
>  	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries, private_id);
>  
> diff --git a/mm/memcontrol-v1.h b/mm/memcontrol-v1.h
> index 23be2512702dc..281425273ea97 100644
> --- a/mm/memcontrol-v1.h
> +++ b/mm/memcontrol-v1.h
> @@ -27,8 +27,7 @@ static inline bool mem_cgroup_private_id_is_root(unsigned short id)
>  	return id == mem_cgroup_private_id(root_mem_cgroup);
>  }
>  
> -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg,
> -						    unsigned int n);
> +unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n);
>  
>  void reparent_memcg_lruvec_state_local(struct mem_cgroup *memcg,
>  				       struct mem_cgroup *parent, int idx);
> diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> index bfe53e4392f09..ed44b3e7ac938 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -4082,7 +4082,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->private_id_ref)) {
>  		mem_cgroup_private_id_remove(memcg);
> @@ -4092,7 +4092,22 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
>  	}
>  }
>  
> -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
> +static inline void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
> +{
> +	struct mem_cgroup *memcg;
> +
> +	rcu_read_lock();

Use lockdep_assert_in_rcu_read_lock() here instead of taking rcu as both callers
already taking rcu read lock.

Re: [PATCH 4/4] memcg: move memcg private ID refcount to objcg

[Shakeel Butt]

On Fri, Sep 18, 2026 at 05:18:43PM +0800, Bingfang Guo via B4 Relay wrote:
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> diff --git a/mm/workingset.c b/mm/workingset.c
> index 8412f4840ae35..7e4fbc5a786d6 100644
> --- a/mm/workingset.c
> +++ b/mm/workingset.c
> @@ -281,6 +281,9 @@ static bool lru_gen_test_recent(void *shadow, struct lruvec **lruvec,
>  	unpack_shadow(shadow, &memcg_id, &pgdat, token, workingset);
>  
>  	memcg = mem_cgroup_from_private_id(memcg_id);
> +	if (!memcg)
> +		return false;

You can not return here without initializing *lruvec as it might be used in the
caller. Also mem_cgroup_lruvec() can handle null memcg and will substitute with
root_mem_cgroup.

> +
>  	*lruvec = mem_cgroup_lruvec(memcg, pgdat);
>  
>  	max_seq = READ_ONCE((*lruvec)->lrugen.max_seq);
> @@ -470,7 +473,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 3/4] memcg: manipulate memcg private ID references by ID

[Bingfang Guo]

On Fri, Sep 18, 2026 at 11:14:03AM +0800, Shakeel Butt wrote:
> On Fri, Sep 18, 2026 at 05:18:42PM +0800, Bingfang Guo via B4 Relay wrote:
> > From: Bingfang Guo <bingfangguo@tencent.com>
> > 
> > This is a preparatory work for moving memcgid from memcg to objcg.
> > 
> > Swap entries retain a private ID rather than a memcg pointer. Once
> > private ID references are moved to objcgs, the ID can also outlive the
> > memcg to which it was originally assigned.  So it's better to make the
> > get and put functions accept the ID itself instead of the memcg.
> > 
> > Rename mem_cgroup_private_id_get_online() to
> > mem_cgroup_private_id_get(), and make it return the ID only.  If the
> > memcg is already dying, the dying memcg will still be used for charging
> > and stats accounting in v2 swap charging path. But they are hierarchical
> > and will be reparented after offlining so it doesn't matter.
> > 
> > Make mem_cgroup_private_id_put() take the ID and resolve the reference
> > holder internally. Convert swap uncharge and charge rollback to release
> > the reference using that ID.  This introduces an extra xarray lookup for
> > now, which will be removed in the final patch.
> > 
> > Separate the online-state reference release into
> > mem_cgroup_private_id_kill(). The offline path already has the memcg
> > pointer and can call the underlying put helper directly.
> > 
> > Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> > ---
> >  mm/memcontrol-v1.c |  7 +++----
> >  mm/memcontrol-v1.h |  3 +--
> >  mm/memcontrol.c    | 32 +++++++++++++++++++++++---------
> >  3 files changed, 27 insertions(+), 15 deletions(-)
> > 
> > diff --git a/mm/memcontrol-v1.c b/mm/memcontrol-v1.c
> > index ed015fdd95123..b7f2868885071 100644
> > --- a/mm/memcontrol-v1.c
> > +++ b/mm/memcontrol-v1.c
> > @@ -268,7 +268,7 @@ void memcg1_commit_charge(struct folio *folio, struct mem_cgroup *memcg)
> >   */
> >  void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
> >  {
> > -	struct mem_cgroup *memcg, *swap_memcg;
> > +	struct mem_cgroup *memcg;
> >  	struct obj_cgroup *objcg;
> >  	unsigned int nr_entries;
> >  	unsigned short private_id;
> > @@ -298,9 +298,8 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
> >  	 * if the ID refers to the root memcg.
> >  	 */
> >  	nr_entries = folio_nr_pages(folio);
> > -	swap_memcg = mem_cgroup_private_id_get_online(memcg, nr_entries);
> > -	private_id = mem_cgroup_private_id(swap_memcg);
> > -	mod_memcg_state(swap_memcg, MEMCG_SWAP, nr_entries);
> > +	private_id = mem_cgroup_private_id_get(memcg, nr_entries);
> > +	mod_memcg_state(memcg, MEMCG_SWAP, nr_entries);
> >  
> >  	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries, private_id);
> >  
> > diff --git a/mm/memcontrol-v1.h b/mm/memcontrol-v1.h
> > index 23be2512702dc..281425273ea97 100644
> > --- a/mm/memcontrol-v1.h
> > +++ b/mm/memcontrol-v1.h
> > @@ -27,8 +27,7 @@ static inline bool mem_cgroup_private_id_is_root(unsigned short id)
> >  	return id == mem_cgroup_private_id(root_mem_cgroup);
> >  }
> >  
> > -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg,
> > -						    unsigned int n);
> > +unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n);
> >  
> >  void reparent_memcg_lruvec_state_local(struct mem_cgroup *memcg,
> >  				       struct mem_cgroup *parent, int idx);
> > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > index bfe53e4392f09..ed44b3e7ac938 100644
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> > @@ -4082,7 +4082,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->private_id_ref)) {
> >  		mem_cgroup_private_id_remove(memcg);
> > @@ -4092,7 +4092,22 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
> >  	}
> >  }
> >  
> > -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
> > +static inline void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
> > +{
> > +	struct mem_cgroup *memcg;
> > +
> > +	rcu_read_lock();
> 
> Use lockdep_assert_in_rcu_read_lock() here instead of taking rcu as both callers
> already taking rcu read lock.
> 

Thanks for pointing out this.

Agreed.  Both two callers are already holding the rcu lock so
taking the lock here is unnecessary. So I will drop the
rcu_read_lock() and use that in the next version!

My concern is that: mem_cgroup_private_id_put() looks like a
universal put function, requiring rcu held (which is true today)
is not that obvious to the users.  So I think adding a short kdoc
comment to make it clear later might be a good idea.

Re: [PATCH 4/4] memcg: move memcg private ID refcount to objcg

[Bingfang Guo]

On Fri, Sep 18, 2026 at 11:19:08AM +0800, Shakeel Butt wrote:
> On Fri, Sep 18, 2026 at 05:18:43PM +0800, Bingfang Guo via B4 Relay wrote:
> > From: Bingfang Guo <bingfangguo@tencent.com>
> > 
> > diff --git a/mm/workingset.c b/mm/workingset.c
> > index 8412f4840ae35..7e4fbc5a786d6 100644
> > --- a/mm/workingset.c
> > +++ b/mm/workingset.c
> > @@ -281,6 +281,9 @@ static bool lru_gen_test_recent(void *shadow, struct lruvec **lruvec,
> >  	unpack_shadow(shadow, &memcg_id, &pgdat, token, workingset);
> >  
> >  	memcg = mem_cgroup_from_private_id(memcg_id);
> > +	if (!memcg)
> > +		return false;
> 
> You can not return here without initializing *lruvec as it might be used in the
> caller. Also mem_cgroup_lruvec() can handle null memcg and will substitute with
> root_mem_cgroup.
> 

Oops. Thanks! Nice catch!

I made a mistake here and it managed to work so I didn't notice
it...

In the previous patch I left here unchanged and used the root.
But then it struck me that maybe it is more consistent to make
the classical LRU and MGLRU behave similarly by skipping the
recency check for both of them.

I'd like to hear about your suggestions:  Whether to simply
remove the if statement and use the root only for MGLRU? Or to
use the root and keep the test speculative for both the two?

Thanks
Bingfang

> > +
> >  	*lruvec = mem_cgroup_lruvec(memcg, pgdat);
> >  
> >  	max_seq = READ_ONCE((*lruvec)->lrugen.max_seq);
> > @@ -470,7 +473,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 3/4] memcg: manipulate memcg private ID references by ID

[Shakeel Butt]

On Sat, Sep 19, 2026 at 02:46:52AM +0800, Bingfang Guo wrote:
> On Fri, Sep 18, 2026 at 11:14:03AM +0800, Shakeel Butt wrote:
> > On Fri, Sep 18, 2026 at 05:18:42PM +0800, Bingfang Guo via B4 Relay wrote:
> > > From: Bingfang Guo <bingfangguo@tencent.com>
> > > 
> > > This is a preparatory work for moving memcgid from memcg to objcg.
> > > 
> > > Swap entries retain a private ID rather than a memcg pointer. Once
> > > private ID references are moved to objcgs, the ID can also outlive the
> > > memcg to which it was originally assigned.  So it's better to make the
> > > get and put functions accept the ID itself instead of the memcg.
> > > 
> > > Rename mem_cgroup_private_id_get_online() to
> > > mem_cgroup_private_id_get(), and make it return the ID only.  If the
> > > memcg is already dying, the dying memcg will still be used for charging
> > > and stats accounting in v2 swap charging path. But they are hierarchical
> > > and will be reparented after offlining so it doesn't matter.
> > > 
> > > Make mem_cgroup_private_id_put() take the ID and resolve the reference
> > > holder internally. Convert swap uncharge and charge rollback to release
> > > the reference using that ID.  This introduces an extra xarray lookup for
> > > now, which will be removed in the final patch.
> > > 
> > > Separate the online-state reference release into
> > > mem_cgroup_private_id_kill(). The offline path already has the memcg
> > > pointer and can call the underlying put helper directly.
> > > 
> > > Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> > > ---
> > >  mm/memcontrol-v1.c |  7 +++----
> > >  mm/memcontrol-v1.h |  3 +--
> > >  mm/memcontrol.c    | 32 +++++++++++++++++++++++---------
> > >  3 files changed, 27 insertions(+), 15 deletions(-)
> > > 
> > > diff --git a/mm/memcontrol-v1.c b/mm/memcontrol-v1.c
> > > index ed015fdd95123..b7f2868885071 100644
> > > --- a/mm/memcontrol-v1.c
> > > +++ b/mm/memcontrol-v1.c
> > > @@ -268,7 +268,7 @@ void memcg1_commit_charge(struct folio *folio, struct mem_cgroup *memcg)
> > >   */
> > >  void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
> > >  {
> > > -	struct mem_cgroup *memcg, *swap_memcg;
> > > +	struct mem_cgroup *memcg;
> > >  	struct obj_cgroup *objcg;
> > >  	unsigned int nr_entries;
> > >  	unsigned short private_id;
> > > @@ -298,9 +298,8 @@ void __memcg1_swapout(struct folio *folio, struct swap_cluster_info *ci)
> > >  	 * if the ID refers to the root memcg.
> > >  	 */
> > >  	nr_entries = folio_nr_pages(folio);
> > > -	swap_memcg = mem_cgroup_private_id_get_online(memcg, nr_entries);
> > > -	private_id = mem_cgroup_private_id(swap_memcg);
> > > -	mod_memcg_state(swap_memcg, MEMCG_SWAP, nr_entries);
> > > +	private_id = mem_cgroup_private_id_get(memcg, nr_entries);
> > > +	mod_memcg_state(memcg, MEMCG_SWAP, nr_entries);
> > >  
> > >  	__swap_cgroup_set(ci, swp_cluster_offset(folio->swap), nr_entries, private_id);
> > >  
> > > diff --git a/mm/memcontrol-v1.h b/mm/memcontrol-v1.h
> > > index 23be2512702dc..281425273ea97 100644
> > > --- a/mm/memcontrol-v1.h
> > > +++ b/mm/memcontrol-v1.h
> > > @@ -27,8 +27,7 @@ static inline bool mem_cgroup_private_id_is_root(unsigned short id)
> > >  	return id == mem_cgroup_private_id(root_mem_cgroup);
> > >  }
> > >  
> > > -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg,
> > > -						    unsigned int n);
> > > +unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n);
> > >  
> > >  void reparent_memcg_lruvec_state_local(struct mem_cgroup *memcg,
> > >  				       struct mem_cgroup *parent, int idx);
> > > diff --git a/mm/memcontrol.c b/mm/memcontrol.c
> > > index bfe53e4392f09..ed44b3e7ac938 100644
> > > --- a/mm/memcontrol.c
> > > +++ b/mm/memcontrol.c
> > > @@ -4082,7 +4082,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->private_id_ref)) {
> > >  		mem_cgroup_private_id_remove(memcg);
> > > @@ -4092,7 +4092,22 @@ static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned
> > >  	}
> > >  }
> > >  
> > > -struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, unsigned int n)
> > > +static inline void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
> > > +{
> > > +	struct mem_cgroup *memcg;
> > > +
> > > +	rcu_read_lock();
> > 
> > Use lockdep_assert_in_rcu_read_lock() here instead of taking rcu as both callers
> > already taking rcu read lock.
> > 
> 
> Thanks for pointing out this.
> 
> Agreed.  Both two callers are already holding the rcu lock so
> taking the lock here is unnecessary. So I will drop the
> rcu_read_lock() and use that in the next version!
> 
> My concern is that: mem_cgroup_private_id_put() looks like a
> universal put function, requiring rcu held (which is true today)
> is not that obvious to the users.  So I think adding a short kdoc
> comment to make it clear later might be a good idea.

This function is internal to memcg code, so whoever is going to call it better
know to call it with rcu held because of lockdep assert. No need to have a
comment.

Re: [PATCH 2/4] memcg: base swap charge accounting on memcgid root status

[Muchun Song]

> On Sep 18, 2026, at 17:18, Bingfang Guo via B4 Relay <devnull+bingfangguo.tencent.com@kernel.org> wrote:
> 
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> A private ID currently pins its original memcg, so testing whether the
> ID belongs to root is equivalent to testing whether the resolved memcg
> is root.
> 
> That equivalence will no longer hold when IDs refer to objcgs. A
> non-root ID may resolve to root after reparenting, but its swap entries
> still carry counter charges inherited by root. Skipping their uncharge
> based on the resolved memcg would leave those charges behind.
> 
> Use the private ID's root status to decide whether a swap entry carries
> a counter charge. A root-ID entry carries none, while a non-root-ID
> entry must release its charge even if its current accounting memcg has
> become root.
> 
> This patch adds a new helper to check if the memcgid equals to that of
> the root memcg.  For swap uncharging and v2 swap charging, simply decide
> whether to charge/uncharge memsw or swap counter based on the swap
> memcgid is root or not. For v1 swapout, don't recharge the memsw
> counter, just cancel the charge if the swap memcg ID refers to the root
> memcg.
> 
> This makes memory and swap charging have similar semantics: one relys on
> the objcg's root status, and the other relys on the memcgid's root
> status.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>

Acked-by: Muchun Song <muchun.song@linux.dev>

Re: [PATCH 3/4] memcg: manipulate memcg private ID references by ID

[Muchun Song]

> On Sep 18, 2026, at 17:18, Bingfang Guo via B4 Relay <devnull+bingfangguo.tencent.com@kernel.org> wrote:
> 
> From: Bingfang Guo <bingfangguo@tencent.com>
> 
> This is a preparatory work for moving memcgid from memcg to objcg.
> 
> Swap entries retain a private ID rather than a memcg pointer. Once
> private ID references are moved to objcgs, the ID can also outlive the
> memcg to which it was originally assigned.  So it's better to make the
> get and put functions accept the ID itself instead of the memcg.
> 
> Rename mem_cgroup_private_id_get_online() to
> mem_cgroup_private_id_get(), and make it return the ID only.  If the
> memcg is already dying, the dying memcg will still be used for charging
> and stats accounting in v2 swap charging path. But they are hierarchical
> and will be reparented after offlining so it doesn't matter.
> 
> Make mem_cgroup_private_id_put() take the ID and resolve the reference
> holder internally. Convert swap uncharge and charge rollback to release
> the reference using that ID.  This introduces an extra xarray lookup for
> now, which will be removed in the final patch.
> 
> Separate the online-state reference release into
> mem_cgroup_private_id_kill(). The offline path already has the memcg
> pointer and can call the underlying put helper directly.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>

With Shakeel's suggestion.

Acked-by: Muchun Song <muchun.song@linux.dev>

Re: [PATCH 4/4] memcg: move memcg private ID refcount to objcg

[Muchun Song]

On 2026/9/18 17:18, Bingfang Guo via B4 Relay wrote:
> 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. The objcg is stored in the global memcgid xarray instead
> and used for retrieving the online memcg from it.  So swapped out pages
> no longer pin the dying memcg.
>
> After the change, a memcgid can refer to a non present memcg. To handle
> this situation, when trying to get the original memcg from the id,
> compare the memcgid passed in with that of the memcg, and return NULL to
> indicate its death if they differ. NULL checks are added for
> list_lru_walk_node(), workingset_test_recent() and lru_gen_test_recent()
> to skip dead memcgs.
>
> In the earlier patch, an extra xarray lookup was introduced in swap
> uncharging path.  Now that we have the objcg pointer in the function,
> the extra overhead can be removed by using it for putting directly.
>
> Signed-off-by: Bingfang Guo <bingfangguo@tencent.com>
> ---
>   include/linux/memcontrol.h |  9 +++---
>   mm/list_lru.c              |  2 +-
>   mm/memcontrol.c            | 81 ++++++++++++++++++++++++++++++++++------------
>   mm/workingset.c            |  5 ++-
>   4 files changed, 71 insertions(+), 26 deletions(-)
>
> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
> index 46bf724cae7af..3fb18191cfbd1 100644
> --- a/include/linux/memcontrol.h
> +++ b/include/linux/memcontrol.h
> @@ -180,6 +180,7 @@ struct obj_cgroup {
>   	struct percpu_ref refcnt;
>   	struct mem_cgroup *memcg;
>   	atomic_t nr_charged_bytes;
> +	refcount_t private_id_ref;
>   	union {
>   		struct list_head list; /* protected by objcg_lock */
>   		struct rcu_head rcu;
> @@ -225,9 +226,6 @@ struct mem_cgroup {
>   	/* vmpressure notifications. Written on every reclaim iteration. */
>   	struct vmpressure vmpressure;
>   
> -	/* Written on every swap charge and uncharge. */
> -	refcount_t private_id_ref;
> -
>   #ifdef CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC
>   	/* MEMCG_KMEM for nmi context */
>   	atomic_t		kmem_stat;
> @@ -324,8 +322,11 @@ struct mem_cgroup {
>   	unsigned long zswap_max;
>   #endif
>   
> +	/* The objcg holding private memcg ID. */
> +	struct obj_cgroup *private_id_objcg;
> +
>   	/* Private memcg ID. Used to ID objects that outlive the cgroup */
> -	int private_id;
> +	unsigned short private_id;
>   
>   	int kmemcg_id;
>   
> diff --git a/mm/list_lru.c b/mm/list_lru.c
> index 8a6dd0a489e12..7edd79113cc56 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 ed44b3e7ac938..22deee8312856 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -4074,6 +4074,18 @@ static void memcg_wb_domain_size_changed(struct mem_cgroup *memcg)
>   #define MEM_CGROUP_ID_MAX	((1UL << MEM_CGROUP_ID_SHIFT) - 1)
>   static DEFINE_XARRAY_ALLOC1(mem_cgroup_private_ids);
>   
> +/**
> + * obj_cgroup_from_private_id - look up the objcg holding the memcg id.
> + * @id: the memcg id to look up
> + *
> + * Caller must hold rcu_read_lock().

This comment is really redundant since you already include a lockdep
assertion. My advice for writing good comments: when the code is
self-explanatory, explain why rather than what.

> + */
> +static inline struct obj_cgroup *obj_cgroup_from_private_id(unsigned short id)
> +{
> +	lockdep_assert_once(rcu_read_lock_held());
> +	return xa_load(&mem_cgroup_private_ids, id);
> +}
> +
>   static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
>   {
>   	if (memcg->private_id > 0) {
> @@ -4082,34 +4094,43 @@ static void mem_cgroup_private_id_remove(struct mem_cgroup *memcg)
>   	}
>   }
>   
> -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->private_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->private_id_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);
>   	}
>   }
>   
>   static inline void mem_cgroup_private_id_put(unsigned short id, unsigned int n)
>   {
> -	struct mem_cgroup *memcg;
> +	struct obj_cgroup *objcg;
>   
>   	rcu_read_lock();
> -	memcg = mem_cgroup_from_private_id(id);
> -	__mem_cgroup_private_id_put(memcg, n);
> +	objcg = obj_cgroup_from_private_id(id);
> +	__mem_cgroup_private_id_put(objcg, id, 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->private_id_objcg, memcg->private_id, 1);
>   }
>   
>   unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n)
>   {
> -	while (!refcount_add_not_zero(n, &memcg->private_id_ref)) {
> +	struct obj_cgroup *objcg;

Please add a blank line between the variable definition and the assertion.

> +	lockdep_assert_once(rcu_read_lock_held());
> +
> +	objcg = memcg->private_id_objcg;
> +
> +	while (!refcount_add_not_zero(n, &objcg->private_id_ref)) {
>   		/*
>   		 * The root cgroup cannot be destroyed, so it's refcount must
>   		 * always be >= 1.
> @@ -4119,6 +4140,7 @@ unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int
>   			break;
>   		}
>   		memcg = parent_mem_cgroup(memcg);
> +		objcg = memcg->private_id_objcg;
>   	}
>   
>   	return mem_cgroup_private_id(memcg);
> @@ -4129,11 +4151,24 @@ unsigned short mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int
>    * @id: the memcg id to look up
>    *
>    * Caller must hold rcu_read_lock().
> + *
> + * @return: the memcg, or NULL if the memcg referred to is already dead.
>    */
>   struct mem_cgroup *mem_cgroup_from_private_id(unsigned short id)
>   {
> +	struct obj_cgroup *objcg;
> +	struct mem_cgroup *memcg;

It would be better to add a blank line here as well.

>   	WARN_ON_ONCE(!rcu_read_lock_held());
> -	return xa_load(&mem_cgroup_private_ids, id);
> +
> +	objcg = obj_cgroup_from_private_id(id);
> +	if (!objcg)
> +		return NULL;
> +
> +	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)
> @@ -4228,18 +4263,21 @@ static struct mem_cgroup *mem_cgroup_alloc(struct mem_cgroup *parent)
>   	struct mem_cgroup *memcg;
>   	int node, cpu;
>   	int __maybe_unused i;
> +	unsigned int private_id;

This could be a separate cleanup patch.

>   	long error;
>   
>   	memcg = kmem_cache_zalloc(memcg_cachep, GFP_KERNEL);
>   	if (!memcg)
>   		return ERR_PTR(-ENOMEM);
>   
> -	error = xa_alloc(&mem_cgroup_private_ids, &memcg->private_id, NULL,
> +	error = xa_alloc(&mem_cgroup_private_ids, &private_id, NULL,
>   			 XA_LIMIT(1, MEM_CGROUP_ID_MAX), GFP_KERNEL);
>   	if (error)
>   		goto fail;
>   	error = -ENOMEM;
>   
> +	memcg->private_id = private_id;
> +
>   	memcg->vmstats = kzalloc_obj(struct memcg_vmstats, GFP_KERNEL_ACCOUNT);
>   	if (!memcg->vmstats)
>   		goto fail;
> @@ -4380,9 +4418,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->private_id_ref, 1);
> -	css_get(css);
> +	/* CSS pins memcg ID, memcg ID pins obj cgroup */
> +	memcg->private_id_objcg = objcg;
> +	refcount_set(&memcg->private_id_objcg->private_id_ref, 1);
> +	obj_cgroup_get(memcg->private_id_objcg);
>   
>   	/*
>   	 * Ensure mem_cgroup_from_private_id() works once we're fully online.
> @@ -4394,7 +4433,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->private_id, memcg, GFP_KERNEL);
> +	xa_store(&mem_cgroup_private_ids, memcg->private_id, memcg->private_id_objcg, GFP_KERNEL);
>   
>   	return 0;
>   free_objcg:
> @@ -5832,8 +5871,6 @@ static void __init memcg_struct_check(void)
>   				      memory_events_local);
>   	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
>   				      vmpressure);
> -	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
> -				      private_id_ref);
>   #ifdef CONFIG_MEMCG_NMI_SAFETY_REQUIRES_ATOMIC
>   	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_write_hot,
>   				      kmem_stat);
> @@ -5878,6 +5915,8 @@ static void __init memcg_struct_check(void)
>   	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
>   				      zswap_writeback);
>   #endif
> +	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
> +				      private_id_objcg);
>   	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
>   				      private_id);
>   	CACHELINE_ASSERT_GROUP_MEMBER(struct mem_cgroup, memcg_read_mostly,
> @@ -6015,10 +6054,12 @@ int __mem_cgroup_try_charge_swap(struct folio *folio)
>    */
>   void __mem_cgroup_uncharge_swap(unsigned short id, unsigned int nr_pages)
>   {
> +	struct obj_cgroup *objcg;
>   	struct mem_cgroup *memcg;
>   
>   	rcu_read_lock();
> -	memcg = mem_cgroup_from_private_id(id);
> +	objcg = obj_cgroup_from_private_id(id);
> +	memcg = obj_cgroup_memcg(objcg);
>   	if (memcg) {
>   		if (!mem_cgroup_private_id_is_root(id)) {
>   			if (do_memsw_account())
> @@ -6027,7 +6068,7 @@ 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);
> +		__mem_cgroup_private_id_put(objcg, id, nr_pages);
>   	}
>   	rcu_read_unlock();
>   }
> diff --git a/mm/workingset.c b/mm/workingset.c
> index 8412f4840ae35..7e4fbc5a786d6 100644
> --- a/mm/workingset.c
> +++ b/mm/workingset.c
> @@ -281,6 +281,9 @@ static bool lru_gen_test_recent(void *shadow, struct lruvec **lruvec,
>   	unpack_shadow(shadow, &memcg_id, &pgdat, token, workingset);
>   
>   	memcg = mem_cgroup_from_private_id(memcg_id);
> +	if (!memcg)
> +		return false;
> +
>   	*lruvec = mem_cgroup_lruvec(memcg, pgdat);
>   
>   	max_seq = READ_ONCE((*lruvec)->lrugen.max_seq);
> @@ -470,7 +473,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();
>   
>