Re: [PATCH bpf-next v11 1/6] bpf: Add sleepable support for raw tracepoint programs

[Peter Zijlstra <peterz@infradead.org>]

On Tue, Apr 21, 2026 at 10:14:17AM -0700, Mykyta Yatsenko wrote:
> From: Mykyta Yatsenko <yatsenko@meta.com>
> 
> Rework __bpf_trace_run() to support sleepable BPF programs by using
> explicit RCU flavor selection, following the uprobe_prog_run() pattern.
> 
> For sleepable programs, use rcu_read_lock_tasks_trace() for lifetime
> protection with migrate_disable().

Why the migrate_disable() ? This is a new kind of BPF program; would it
not be a good opportunity to not have this constraint?

> For non-sleepable programs, use the
> regular rcu_read_lock_dont_migrate().
> 
> Remove the preempt_disable_notrace/preempt_enable_notrace pair from
> the faultable tracepoint BPF probe wrapper in bpf_probe.h, since
> migration protection and RCU locking are now handled per-program
> inside __bpf_trace_run().
> 
> Adapt bpf_prog_test_run_raw_tp() for sleepable programs: reject
> BPF_F_TEST_RUN_ON_CPU since sleepable programs cannot run in hardirq
> or preempt-disabled context, and call __bpf_prog_test_run_raw_tp()
> directly instead of via smp_call_function_single(). Rework
> __bpf_prog_test_run_raw_tp() to select RCU flavor per-program and
> add per-program recursion context guard for private stack safety.
> 
> Acked-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
> Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
> ---
>  include/trace/bpf_probe.h |  2 --
>  kernel/trace/bpf_trace.c  | 20 ++++++++++++---
>  net/bpf/test_run.c        | 65 ++++++++++++++++++++++++++++++++++++-----------
>  3 files changed, 67 insertions(+), 20 deletions(-)
> 
> diff --git a/include/trace/bpf_probe.h b/include/trace/bpf_probe.h
> index 9391d54d3f12..d1de8f9aa07f 100644
> --- a/include/trace/bpf_probe.h
> +++ b/include/trace/bpf_probe.h
> @@ -58,9 +58,7 @@ static notrace void							\
>  __bpf_trace_##call(void *__data, proto)					\
>  {									\
>  	might_fault();							\
> -	preempt_disable_notrace();					\
>  	CONCATENATE(bpf_trace_run, COUNT_ARGS(args))(__data, CAST_TO_U64(args));	\
> -	preempt_enable_notrace();					\
>  }
>  
>  #undef DECLARE_EVENT_SYSCALL_CLASS
> diff --git a/kernel/trace/bpf_trace.c b/kernel/trace/bpf_trace.c
> index e916f0ccbed9..7276c72c1d31 100644
> --- a/kernel/trace/bpf_trace.c
> +++ b/kernel/trace/bpf_trace.c
> @@ -2072,11 +2072,19 @@ void bpf_put_raw_tracepoint(struct bpf_raw_event_map *btp)
>  static __always_inline
>  void __bpf_trace_run(struct bpf_raw_tp_link *link, u64 *args)
>  {
> +	struct srcu_ctr __percpu *scp = NULL;
>  	struct bpf_prog *prog = link->link.prog;
> +	bool sleepable = prog->sleepable;
>  	struct bpf_run_ctx *old_run_ctx;
>  	struct bpf_trace_run_ctx run_ctx;
>  
> -	rcu_read_lock_dont_migrate();
> +	if (sleepable) {
> +		scp = rcu_read_lock_tasks_trace();
> +		migrate_disable();
> +	} else {
> +		rcu_read_lock_dont_migrate();
> +	}
> +
>  	if (unlikely(!bpf_prog_get_recursion_context(prog))) {
>  		bpf_prog_inc_misses_counter(prog);
>  		goto out;
> @@ -2085,12 +2093,18 @@ void __bpf_trace_run(struct bpf_raw_tp_link *link, u64 *args)
>  	run_ctx.bpf_cookie = link->cookie;
>  	old_run_ctx = bpf_set_run_ctx(&run_ctx.run_ctx);
>  
> -	(void) bpf_prog_run(prog, args);
> +	(void)bpf_prog_run(prog, args);
>  
>  	bpf_reset_run_ctx(old_run_ctx);
>  out:
>  	bpf_prog_put_recursion_context(prog);
> -	rcu_read_unlock_migrate();
> +
> +	if (sleepable) {
> +		migrate_enable();
> +		rcu_read_unlock_tasks_trace(scp);
> +	} else {
> +		rcu_read_unlock_migrate();
> +	}
>  }

Since you have a clear distinction between __BPF_DECLARE_TRACE() and
__BPF_DELARE_TRACE_SYSCALL(), would it not make more sense to have
different versions of __bpf_trace_run() instead of having a runtime
condition on ->sleepable ?