Re: [PATCH bpf-next v2 1/2] bpf: Migrate bpf_task_work to kmalloc_nolock

[Mykyta Yatsenko <mykyta.yatsenko5@gmail.com>]

On 3/31/26 1:00 AM, Andrii Nakryiko wrote:
> On Mon, Mar 30, 2026 at 3:28 PM Mykyta Yatsenko
> <mykyta.yatsenko5@gmail.com> wrote:
>>
>> From: Mykyta Yatsenko <yatsenko@meta.com>
>>
>> Replace bpf_mem_alloc/bpf_mem_free with
>> kmalloc_nolock/kfree_rcu for bpf_task_work_ctx.
>>
>> Replace guard(rcu_tasks_trace)() with guard(rcu)() in
>> bpf_task_work_irq(). The function only accesses ctx struct members
>> (not map values), so tasks trace protection is not needed - regular
>> RCU is sufficient since ctx is freed via kfree_rcu. The guard in
>> bpf_task_work_callback() remains as tasks trace since it accesses map
>> values from process context.
> 
> I didn't quite get if this change was necessary for correctness or
> it's just an optimization?
> 

Correctness - ctx is freed via kfree_rcu(), so we need to hold rcu read 
lock when we pass refcnt to the task_work_add() callback. It worked 
before on tasks trace rcu because bpf_mem_alloc() used it (with normal 
rcu chaining) before freeing the ctx.

>>
>> Sleepable BPF programs hold rcu_read_lock_trace but not
>> regular rcu_read_lock. Since kfree_rcu
>> waits for a regular RCU grace period, the ctx memory can be freed
>> while a sleepable program is still running. Add scoped_guard(rcu)
>> around the pointer read and refcount tryget in
>> bpf_task_work_acquire_ctx to close this race window.
>>
>> Since kfree_rcu uses call_rcu internally which is not safe from
>> NMI context, defer destruction via irq_work when irqs are disabled.
>>
>> For the lost-cmpxchg path the ctx was never published, so
>> kfree_nolock is safe.
>>
>> Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
>> ---
>>   kernel/bpf/helpers.c | 56 ++++++++++++++++++++++++++++++++++------------------
>>   1 file changed, 37 insertions(+), 19 deletions(-)
>>
> 
> [...]