Re: [PATCH bpf-next v6 1/2] mm/bpf: Add bpf_proactive_reclaim kfunc

[JP Kobryn <jp.kobryn@linux.dev>]

On 9/1/26 11:57 PM, Hui Zhu wrote:
> Hi JP,
> 
> Thanks for the review!
> 
>> Hi Hui,
>>
>> On 8/31/26 7:21 PM, Hui Zhu wrote:
>>> From: Hui Zhu <zhuhui@kylinos.cn>
>>>
>>> Add bpf_proactive_reclaim(), a sleepable kfunc which performs one
>>> proactive reclaim pass on a given memory cgroup, similar to a write
>>> to memory.reclaim but without retrying until the target is reached.
>>>
>>> The kfunc is restricted to BPF_PROG_TYPE_SYSCALL so that reclaim
>>> always runs in a clean process context. Generic sleepable programs
>>> may execute with filesystem locks held or in NOFS/NOIO contexts,
>>> where the reclaim path could deadlock in filesystem shrinkers. A
>>> SYSCALL program can still drive reclaim asynchronously through
>>> bpf_wq or task_work callbacks, which run in process context and
>>> keep the SYSCALL program type, so they can call the kfunc too.
>>>
>>> The kfunc refuses to reclaim if the calling task is already in a
>>> reclaim context, as a nested reclaim would corrupt the outer reclaim
>>> state.
>>>
>>> Signed-off-by: Hui Zhu <zhuhui@kylinos.cn>
>>> ---
>>
>> The difflog is missing in these patches.
> Sorry, my mistake. Will add the changelog (changes since v4) in the
> next version.
>>
>>>   mm/bpf_memcontrol.c | 76 +++++++++++++++++++++++++++++++++++++++++++--
>>>   1 file changed, 74 insertions(+), 2 deletions(-)
>>>
>>> diff --git a/mm/bpf_memcontrol.c b/mm/bpf_memcontrol.c
>>> index 716df49d7647..fd48faa5f8b0 100644
>>> --- a/mm/bpf_memcontrol.c
>>> +++ b/mm/bpf_memcontrol.c
>>> @@ -6,6 +6,7 @@
>>>    */
>>>     #include <linux/memcontrol.h>
>>> +#include <linux/swap.h>
>>>   #include <linux/bpf.h>
>>>     __bpf_kfunc_start_defs();
>>> @@ -159,6 +160,55 @@ __bpf_kfunc void 
>>> bpf_mem_cgroup_flush_stats(struct mem_cgroup *memcg)
>>>       mem_cgroup_flush_stats(memcg);
>>>   }
>>>   +/*
>>> + * Reclaim must not recurse: try_to_free_mem_cgroup_pages() overwrites
>>> + * current->reclaim_state, so a nested call would corrupt the outer
>>> + * reclaim state. Reclaim windows are marked with PF_MEMALLOC;
>>> + * reclaim_state is also checked because it is installed slightly
>>> + * before PF_MEMALLOC.
>>> + */
>>> +static bool bpf_in_reclaim_context(void)
>>> +{
>>> +    return (current->flags & PF_MEMALLOC) || current->reclaim_state;
>>> +}
>>> +
>>> +/**
>>> + * bpf_proactive_reclaim - proactively reclaim memory from a memory
>>> + *                         cgroup
>>> + * @memcg: the target memory cgroup to reclaim from
>>> + * @size:  the amount of memory to reclaim, in bytes
>>> + *
>>> + * Trigger one proactive reclaim pass on @memcg, similar to a write to
>>> + * memory.reclaim, but without retrying until @size is reached.
>>> + *
>>> + * This kfunc is restricted to BPF_PROG_TYPE_SYSCALL to ensure it runs
>>> + * in a clean process context. The SYSCALL program can schedule the
>>> + * actual reclaim work via bpf_wq or timers, which also execute in
>>> + * safe process context (workqueue, task_work).
>>
>> On the workqueue aspect, I could see potential issues. The target size
>> has no upper bound so the total scan/execution time on the shared
>> wq can easily stall other work. Contention on the lru_lock can make
>> matters worse because interrupts are disabled while holding the lock. So
>> the contention would not only stall other work, but can delay IPI
>> handling leading to CSD lock stalls.
> 
> 
> Agreed. I previously misread
> 
>      .nr_to_reclaim = max(nr_pages, SWAP_CLUSTER_MAX)
> 
> in try_to_free_mem_cgroup_pages() as an upper bound on the reclaim
> target; it is actually a lower bound, so nothing inside the reclaim
> path limits how long a single kfunc call can run. The next version
> will cap the per-invocation reclaim target.
> 
> 
>>
>> It looks like the only existing path that explicitly calls
>> try_to_free_mem_cgroup_pages() from a shared wq is the memory.high
>> fallback used when the limit is exceeded outside of task context. But
>> even in that case, it's more constrained. The reclaim request is bounded
>> at MEMCG_CHARGE_BATCH (in high_work_func()) and is limited to one work
>> item per memcg.
>>
>> Would it make sense to follow the existing precedent and use the same
>> bound in your kfunc? You could then batch the wq submissions and you
>> would also be able to stop submitting in between if needed, like in the
>> case of the cgroup dying.
>>
> Yes, the next version will cap the reclaim target of a single
> bpf_proactive_reclaim() call at MEMCG_CHARGE_BATCH, following the
> high_work_func() precedent, so each invocation is a bounded unit of
> work on the shared wq.
> 
> For the "one work item per memcg" side, I'd like to hear your
> thoughts on the following approach: allow only one in-flight
> bpf_proactive_reclaim() per memcg. The kfunc would take a per-memcg
> flag (atomic cmpxchg) on entry and return 0 if another BPF reclaim
> pass is already running on the same memcg, mirroring the
> one-work-item-per-memcg property of high_work. This would prevent wq
> work items from piling up reclaiming the same memcg.
> 
> The reason for doing this with a per-memcg in-flight check rather
> than a fixed per-memcg work item is to keep bpf_proactive_reclaim()
> flexible: it bounds how much reclaim can run against one memcg at any
> moment, while leaving the reclaim policy -- when to reclaim, how many
> passes to batch, and when to stop (e.g. if the target cgroup is
> dying) -- entirely in the BPF program. Do you think this is a
> reasonable way to bound the total reclaim activity per memcg, or
> would you prefer something else?

Let's not add state to the memcg to emulate high_work in bpf. I think
you should document and implement this on the bpf side. Change the
selftest example so that bpf_proactive_reclaim() is invoked only once
per callback (instead of a in a loop) and requeue the same work if
another reclaim batch is needed. You could also check if the cgroup is
dying and discontinue the work.

I would also document that each memcg should use its own bpf_wq item.
Correct me if I'm wrong, but it looks like that is your intention in the
selftest wq_map.

> 
> With the per-invocation cap, each kfunc call becomes a small, bounded
> unit of work, and the BPF program does the batching: it schedules
> successive wq submissions and can stop submitting between passes when
> needed. This keeps the "when and how hard to reclaim" policy in BPF
> while bounding the kernel-side cost of each invocation.
> 
> Best,
> Hui
> 
> 
>>> + *
>>> + * Must not be called with a filesystem lock held: the reclaim path
>>> + * may deadlock on it via filesystem shrinkers.
>>> + *
>>> + * Return: The amount of memory reclaimed, in bytes, or 0 if @size is
>>> + * smaller than a page or the task is already in a reclaim context.
>>> + */
>>> +__bpf_kfunc unsigned long bpf_proactive_reclaim(struct mem_cgroup 
>>> *memcg,
>>> +                        unsigned long size)
>>> +{
>>> +    unsigned long nr_reclaimed;
>>> +
>>> +    if (size < PAGE_SIZE || unlikely(bpf_in_reclaim_context()))
>>> +        return 0;
>>> +
>>> +    nr_reclaimed = try_to_free_mem_cgroup_pages(memcg, size / 
>>> PAGE_SIZE,
>>> +                            GFP_KERNEL,
>>> +                            MEMCG_RECLAIM_MAY_SWAP |
>>> +                            MEMCG_RECLAIM_PROACTIVE,
>>> +                            NULL);
>>> +
>>> +    return nr_reclaimed * PAGE_SIZE;
>>> +}
>>> +
>>>   __bpf_kfunc_end_defs();
>>>     BTF_KFUNCS_START(bpf_memcontrol_kfuncs)
>>> @@ -171,22 +221,44 @@ BTF_ID_FLAGS(func, bpf_mem_cgroup_memory_events)
>>>   BTF_ID_FLAGS(func, bpf_mem_cgroup_usage)
>>>   BTF_ID_FLAGS(func, bpf_mem_cgroup_page_state)
>>>   BTF_ID_FLAGS(func, bpf_mem_cgroup_flush_stats, KF_SLEEPABLE)
>>> -
>>>   BTF_KFUNCS_END(bpf_memcontrol_kfuncs)
>>>   +/*
>>> + * Proactive reclaim needs a clean process context, so it is restricted
>>> + * to BPF_PROG_TYPE_SYSCALL. The bpf_wq and task_work callbacks that a
>>> + * SYSCALL program schedules run as the same program type, so they can
>>> + * still invoke it; generic sleepable programs (e.g. fentry on reclaim
>>> + * paths, inode_rmdir) cannot.
>>> + */
>>> +BTF_KFUNCS_START(bpf_memcontrol_reclaim_kfuncs)
>>> +BTF_ID_FLAGS(func, bpf_proactive_reclaim, KF_SLEEPABLE)
>>> +BTF_KFUNCS_END(bpf_memcontrol_reclaim_kfuncs)
>>> +
>>>   static const struct btf_kfunc_id_set bpf_memcontrol_kfunc_set = {
>>>       .owner          = THIS_MODULE,
>>>       .set            = &bpf_memcontrol_kfuncs,
>>>   };
>>>   +static const struct btf_kfunc_id_set 
>>> bpf_memcontrol_reclaim_kfunc_set = {
>>> +    .owner          = THIS_MODULE,
>>> +    .set            = &bpf_memcontrol_reclaim_kfuncs,
>>> +};
>>> +
>>>   static int __init bpf_memcontrol_init(void)
>>>   {
>>>       int err;
>>>         err = register_btf_kfunc_id_set(BPF_PROG_TYPE_UNSPEC,
>>>                       &bpf_memcontrol_kfunc_set);
>>> -    if (err)
>>> +    if (err) {
>>>           pr_warn("error while registering bpf memcontrol kfuncs: 
>>> %d", err);
>>> +        return err;
>>> +    }
>>> +
>>> +    err = register_btf_kfunc_id_set(BPF_PROG_TYPE_SYSCALL,
>>> +                    &bpf_memcontrol_reclaim_kfunc_set);
>>> +    if (err)
>>> +        pr_warn("error registering bpf reclaim kfuncs: %d", err);
>>>         return err;
>>>   }
>>