Re: [PATCH 6/6] slab: Introduce kmalloc_nolock() and kfree_nolock().

[Alexei Starovoitov <alexei.starovoitov@gmail.com>]

On Mon, May 5, 2025 at 11:46 AM Shakeel Butt <shakeel.butt@linux.dev> wrote:
>
> On Wed, Apr 30, 2025 at 08:27:18PM -0700, Alexei Starovoitov wrote:
> > From: Alexei Starovoitov <ast@kernel.org>
> >
> > --- a/mm/memcontrol.c
> > +++ b/mm/memcontrol.c
> > @@ -595,7 +595,13 @@ static inline void memcg_rstat_updated(struct mem_cgroup *memcg, int val)
> >       if (!val)
> >               return;
> >
> > -     cgroup_rstat_updated(memcg->css.cgroup, cpu);
> > +     /*
> > +      * If called from NMI via kmalloc_nolock -> memcg_slab_post_alloc_hook
> > +      * -> obj_cgroup_charge -> mod_memcg_state,
> > +      * then delay the update.
> > +      */
> > +     if (!in_nmi())
> > +             cgroup_rstat_updated(memcg->css.cgroup, cpu);
>
> I don't think we can just ignore cgroup_rstat_updated() for nmi as there
> is a chance (though very small) that we will loose these stats updates.

I'm failing to understand why it's an issue.
Not doing cgroup_rstat_updated() can only cause updated_next link
to stay NULL when it should be set,
but it should be harmless, and no different from racy check
that the code already doing:
if (data_race(cgroup_rstat_cpu(cgrp, cpu)->updated_next))
  return;

Imaging it was !NULL, the code would return,
but then preemption, something clears it to NULL,
and here we're skipping a set of updated_next.

> In addition, memcg_rstat_updated() itself is not reentrant safe along
> with couple of functions leading to it like __mod_memcg_lruvec_state().

Sure. __mod_memcg_lruvec_state() is not reentrant,
but it's not an issue for kmalloc_nolock(), since objcg/memcg
charge/uncharge from slub is not calling it (as far as I can tell).

>
> >       statc = this_cpu_ptr(memcg->vmstats_percpu);
> >       for (; statc; statc = statc->parent) {
> >               /*
> > @@ -2895,7 +2901,7 @@ static bool consume_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
> >       unsigned long flags;
> >       bool ret = false;
> >
> > -     local_lock_irqsave(&memcg_stock.stock_lock, flags);
> > +     local_lock_irqsave_check(&memcg_stock.stock_lock, flags);
> >
> >       stock = this_cpu_ptr(&memcg_stock);
> >       if (objcg == READ_ONCE(stock->cached_objcg) && stock->nr_bytes >= nr_bytes) {
> > @@ -2995,7 +3001,7 @@ static void refill_obj_stock(struct obj_cgroup *objcg, unsigned int nr_bytes,
> >       unsigned long flags;
> >       unsigned int nr_pages = 0;
> >
> > -     local_lock_irqsave(&memcg_stock.stock_lock, flags);
> > +     local_lock_irqsave_check(&memcg_stock.stock_lock, flags);
> >
> >       stock = this_cpu_ptr(&memcg_stock);
> >       if (READ_ONCE(stock->cached_objcg) != objcg) { /* reset if necessary */
> > @@ -3088,6 +3094,27 @@ static inline size_t obj_full_size(struct kmem_cache *s)
> >       return s->size + sizeof(struct obj_cgroup *);
> >  }
> >
> > +/*
> > + * Try subtract from nr_charged_bytes without making it negative
> > + */
> > +static bool obj_cgroup_charge_atomic(struct obj_cgroup *objcg, gfp_t flags, size_t sz)
> > +{
> > +     size_t old = atomic_read(&objcg->nr_charged_bytes);
> > +     u32 nr_pages = sz >> PAGE_SHIFT;
> > +     u32 nr_bytes = sz & (PAGE_SIZE - 1);
> > +
> > +     if ((ssize_t)(old - sz) >= 0 &&
> > +         atomic_cmpxchg(&objcg->nr_charged_bytes, old, old - sz) == old)
> > +             return true;
> > +
> > +     nr_pages++;
> > +     if (obj_cgroup_charge_pages(objcg, flags, nr_pages))
> > +             return false;
> > +
> > +     atomic_add(PAGE_SIZE - nr_bytes, &objcg->nr_charged_bytes);
> > +     return true;
> > +}
> > +
> >  bool __memcg_slab_post_alloc_hook(struct kmem_cache *s, struct list_lru *lru,
> >                                 gfp_t flags, size_t size, void **p)
> >  {
> > @@ -3128,6 +3155,21 @@ bool __memcg_slab_post_alloc_hook(struct kmem_cache *s, struct list_lru *lru,
> >                       return false;
> >       }
> >
> > +     if (!gfpflags_allow_spinning(flags)) {
> > +             if (local_lock_is_locked(&memcg_stock.stock_lock)) {
> > +                     /*
> > +                      * Cannot use
> > +                      * lockdep_assert_held(this_cpu_ptr(&memcg_stock.stock_lock));
> > +                      * since lockdep might not have been informed yet
> > +                      * of lock acquisition.
> > +                      */
> > +                     return obj_cgroup_charge_atomic(objcg, flags,
> > +                                                     size * obj_full_size(s));
>
> We can not just ignore the stat updates here.
>
> > +             } else {
> > +                     lockdep_assert_not_held(this_cpu_ptr(&memcg_stock.stock_lock));
> > +             }
> > +     }
> > +
> >       for (i = 0; i < size; i++) {
> >               slab = virt_to_slab(p[i]);
> >
> > @@ -3162,8 +3204,12 @@ bool __memcg_slab_post_alloc_hook(struct kmem_cache *s, struct list_lru *lru,
> >  void __memcg_slab_free_hook(struct kmem_cache *s, struct slab *slab,
> >                           void **p, int objects, struct slabobj_ext *obj_exts)
> >  {
> > +     bool lock_held = local_lock_is_locked(&memcg_stock.stock_lock);
> >       size_t obj_size = obj_full_size(s);
> >
> > +     if (likely(!lock_held))
> > +             lockdep_assert_not_held(this_cpu_ptr(&memcg_stock.stock_lock));
> > +
> >       for (int i = 0; i < objects; i++) {
> >               struct obj_cgroup *objcg;
> >               unsigned int off;
> > @@ -3174,8 +3220,12 @@ void __memcg_slab_free_hook(struct kmem_cache *s, struct slab *slab,
> >                       continue;
> >
> >               obj_exts[off].objcg = NULL;
> > -             refill_obj_stock(objcg, obj_size, true, -obj_size,
> > -                              slab_pgdat(slab), cache_vmstat_idx(s));
> > +             if (unlikely(lock_held)) {
> > +                     atomic_add(obj_size, &objcg->nr_charged_bytes);
>
> objcg->nr_charged_bytes is stats ignorant and the relevant stats need to
> be updated before putting stuff into it.

I'm not following.
It's functionally equivalent to refill_obj_stock() without
__account_obj_stock().
And the stats are not ignored.
The next __memcg_slab_free_hook() from good context will update
them. It's only a tiny delay in update.
I don't see why it's an issue.

> > +             } else {
> > +                     refill_obj_stock(objcg, obj_size, true, -obj_size,
> > +                                      slab_pgdat(slab), cache_vmstat_idx(s));
> > +             }
> >               obj_cgroup_put(objcg);
> >       }
> >  }
>
> I am actually working on making this whole call chain (i.e.
> kmalloc/kmem_cache_alloc to memcg [un]charging) reentrant/nmi safe.

Thank you for working on it!
You mean this set:
https://lore.kernel.org/all/20250429061211.1295443-1-shakeel.butt@linux.dev/
?
it's making css_rstat_updated() re-entrant,
which is renamed/reworked version of memcg_rstat_updated().
That's good, but not enough from slub pov.
It removes the need for the first hunk in this patch from mm/memcontrol.c
+ if (!in_nmi())
+               cgroup_rstat_updated(...);

but hunks in __memcg_slab_post_alloc_hook() and __memcg_slab_free_hook()
are still needed.
And I think the obj_cgroup_charge_atomic() approach in this patch is correct.
The delay in rstat update seems fine.
Please help me understand what I'm missing.