Re: [PATCH RFC 1/5] memcg: move memcg private ID refcount to objcg

[Bingfang Guo]

> 2026年8月13日 16:52，Bingfang Guo via B4 Relay <devnull+bingfangguo.tencent.com@xxxxxxxxxx> 写道：
> 
> From: Bingfang Guo <bingfangguo@xxxxxxxxxxx>
> 
> In the previous series by Muchun Song and Qi Zheng, folios are charged
> to the objcg and reparented as the memcg offlines. Same can be done to
> memcg private ID and its main user: swap entries. Make the memcgid
> xarray hold a pointer and a reference to an objcg of the memcg, which is
> used to find the memcg (or its parent) later on. The online state now
> pins the objcg instead of the css, so swapped out pages no longer pin
> the dying memcg.
> 
> The id reference held by the online state is released in css_released()
> after reparenting instead of in css_offline(). This is the key
> invariant the rest of the series builds on: css_offline() runs while
> other css references may still be held, but css_released() only runs
> once the last reference is gone, so a caller holding a memcg reference
> can always count on the id refcount being alive. To prevent races
> between memcgid put in css offline and memcgid get, the release is put
> off till css_released, which could delay the release of the memcgid and
> the objcg it pins, but overall it should be fine.
> 
> After reparenting, the objcg points to a live ancestor, so
> mem_cgroup_from_private_id() now returns that ancestor instead of the
> memcg the ID originally belonged to. Callers that need the exact memcg
> are fixed in patch 5.
> 
> Signed-off-by: Bingfang Guo <bingfangguo@xxxxxxxxxxx>
> ---
> include/linux/memcontrol.h | 10 +++--
> mm/memcontrol.c            | 99 ++++++++++++++++++++++++++++++++--------------
> 2 files changed, 76 insertions(+), 33 deletions(-)
> 
> diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
> index 8170bb8066a22..c33ec7efad50b 100644
> --- a/include/linux/memcontrol.h
> +++ b/include/linux/memcontrol.h
> @@ -191,6 +191,7 @@ struct obj_cgroup {
> struct rcu_head rcu;
> };
> bool is_root;
> + refcount_t id_ref;
> };
> 
> /*
> @@ -202,8 +203,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 */
> @@ -270,6 +271,9 @@ struct mem_cgroup {
> #endif
> int kmemcg_id;
> 
> + /* Private memcg ID. Used to ID objects that outlive the cgroup */
> + int id;
> +
> struct memcg_vmstats_percpu __percpu *vmstats_percpu;
> 
> #ifdef CONFIG_CGROUP_WRITEBACK
> @@ -820,7 +824,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 8319ad8c5c23a..5f30e76ee93d7 100644
> --- a/mm/memcontrol.c
> +++ b/mm/memcontrol.c
> @@ -3697,7 +3697,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)
> @@ -3956,25 +3956,40 @@ 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 inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n)
> +/**
> + * @objcg: the objcg returned by mem_cgroup_private_id_objcg
> + * @id: the corresponding memcg private id
> + */
> +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->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(struct mem_cgroup *memcg, unsigned int n)
> +{
> + __mem_cgroup_private_id_put(memcg->id_objcg, memcg->id, n);
> +}
> +
> 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 = memcg->id_objcg;
> +
> + while (!refcount_add_not_zero(n, &objcg->id_ref)) {
> /*
> * The root cgroup cannot be destroyed, so it's refcount must
> * always be >= 1.
> @@ -3984,6 +3999,7 @@ struct mem_cgroup *mem_cgroup_private_id_get_online(struct mem_cgroup *memcg, un
> break;
> }
> memcg = parent_mem_cgroup(memcg);
> + objcg = memcg->id_objcg;
> }
> return memcg;
> }
> @@ -3996,8 +4012,29 @@ 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);
> +}
> +
> +static struct mem_cgroup *mem_cgroup_take_from_private_id(unsigned short id, unsigned int n)
> +{
> + struct obj_cgroup *objcg;
> + struct mem_cgroup *memcg;
> +
> + objcg = xa_load(&mem_cgroup_private_ids, id);
> + if (!objcg)
> + return NULL;
> +
> + memcg = get_mem_cgroup_from_objcg(objcg);
> +
> + __mem_cgroup_private_id_put(objcg, id, n);
> + return memcg;
> }
> 
> struct mem_cgroup *mem_cgroup_get_from_id(u64 id)
> @@ -4098,7 +4135,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;
> @@ -4243,9 +4280,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 = memcg->nodeinfo[0]->objcg;
> + refcount_set(&memcg->id_objcg->id_ref, 1);
> + obj_cgroup_get(memcg->id_objcg);
> 
> /*
> * Ensure mem_cgroup_from_private_id() works once we're fully online.
> @@ -4257,7 +4295,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, objcg, GFP_KERNEL);

Sashiko pointed out that objcg here is pointing to the wrong node here.
I made a mistake here while rebasing the patch set. Also the problem reported by syzbot and in
patch 2 is also caused by this…

It should be like this:

@@ -4257,7 +4295,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);

But it also says that it could go wrong if node 0 is not present on sparse NUMA setups, so I
think it might be better to just use the last objcg set up above:

@@ -4294,7 +4295,7 @@ static int mem_cgroup_css_online(struct cgroup_subsys_state *css)
        lru_gen_online_memcg(memcg);
 
        /* CSS pins memcg ID, memcg ID pins obj cgroup */
-       memcg->id_objcg = memcg->nodeinfo[0]->objcg;
+       memcg->id_objcg = objcg;
        refcount_set(&memcg->id_objcg->id_ref, 1);
        obj_cgroup_get(memcg->id_objcg);

> 
> return 0;
> free_objcg:
> @@ -4308,8 +4346,6 @@ static void mem_cgroup_css_offline(struct cgroup_subsys_state *css)
> lru_gen_offline_memcg(memcg);
> 
> drain_all_stock(memcg);
> -
> - mem_cgroup_private_id_put(memcg, 1);
> }
> 
> static void mem_cgroup_css_released(struct cgroup_subsys_state *css)
> @@ -4318,6 +4354,9 @@ static void mem_cgroup_css_released(struct cgroup_subsys_state *css)
> 
> invalidate_reclaim_iterators(memcg);
> lru_gen_release_memcg(memcg);
> +
> + mem_cgroup_private_id_put(memcg, 1);

Sashiko says that if memcg onlining fails early, this could result in null pointer dereference.
I think we can fix it like this since mem_cgroup_private_id_put is only used in css releasing:

@@ -3980,9 +3980,13 @@ static void __mem_cgroup_private_id_put(struct obj_cgroup *objcg,
        }
 }
 
-static inline void mem_cgroup_private_id_put(struct mem_cgroup *memcg, unsigned int n)
+static inline void mem_cgroup_private_id_release(struct mem_cgroup *memcg, unsigned int n)
 {
+       if (!memcg->id_objcg)
+               return;
+
        __mem_cgroup_private_id_put(memcg->id_objcg, memcg->id, n);
+       memcg->id_objcg = NULL;
 }
 
 void mem_cgroup_private_id_get(struct mem_cgroup *memcg, unsigned int n)
@@ -4367,9 +4371,7 @@ static void mem_cgroup_css_released(struct cgroup_subsys_state *css)
 
        invalidate_reclaim_iterators(memcg);
        lru_gen_release_memcg(memcg);
-
-       mem_cgroup_private_id_put(memcg, 1);
-       memcg->id_objcg = NULL;
+       mem_cgroup_private_id_release(memcg, 1);
 }
 
 static void mem_cgroup_css_free(struct cgroup_subsys_state *css)

> + memcg->id_objcg = NULL;
> }
> 
> static void mem_cgroup_css_free(struct cgroup_subsys_state *css)
> @@ -5651,19 +5690,19 @@ 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);
> - if (memcg) {
> - if (!mem_cgroup_is_root(memcg)) {
> - if (do_memsw_account())
> - page_counter_uncharge(&memcg->memsw, nr_pages);
> - else
> - page_counter_uncharge(&memcg->swap, nr_pages);
> - }
> - mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
> - mem_cgroup_private_id_put(memcg, nr_pages);
> + memcg = mem_cgroup_take_from_private_id(id, nr_pages);
> + if (!memcg)
> + return;
> +
> + if (!mem_cgroup_is_root(memcg)) {
> + if (do_memsw_account())
> + page_counter_uncharge(&memcg->memsw, nr_pages);
> + else
> + page_counter_uncharge(&memcg->swap, nr_pages);
> }
> - rcu_read_unlock();
> + mod_memcg_state(memcg, MEMCG_SWAP, -nr_pages);
> +
> + mem_cgroup_put(memcg);
> }
> 
> long mem_cgroup_get_nr_swap_pages(struct mem_cgroup *memcg)
> 
> -- 
> 2.43.7
> 
>