bpf: Introduce lock-free bpf_async_update_prog_callback()

Introduce bpf_async_update_prog_callback(): lock-free update of cb->prog
and cb->callback_fn. This function allows updating prog and callback_fn
fields of the struct bpf_async_cb without holding lock.
For now use it under the lock from __bpf_async_set_callback(), in the
next patches that lock will be removed.

Lock-free algorithm:
 * Acquire a guard reference on prog to prevent it from being freed
   during the retry loop.
 * Retry loop:
    1. Each iteration acquires a new prog reference and stores it
       in cb->prog via xchg. The previous prog is released.
    2. The loop condition checks if both cb->prog and cb->callback_fn
       match what we just wrote. If either differs, a concurrent writer
       overwrote our value, and we must retry.
    3. When we retry, our previously-stored prog was already released by
       the concurrent writer or will be released by us after
       overwriting.
 * Release guard reference.

Acked-by: Andrii Nakryiko <andrii@kernel.org>
Signed-off-by: Mykyta Yatsenko <yatsenko@meta.com>
Link: https://lore.kernel.org/r/20260120-timer_nolock-v6-3-670ffdd787b4@meta.com
Signed-off-by: Alexei Starovoitov <ast@kernel.org>

diff --git a/kernel/bpf/helpers.c b/kernel/bpf/helpers.c
index 6eadb66b8c67ea407eb0103ece1eb003298be57d..2a2df867bfe72865bf942617695fb0a7a9bbb7c9 100644 (file)
--- a/kernel/bpf/helpers.c
+++ b/kernel/bpf/helpers.c
@@ -1354,10 +1354,43 @@ static const struct bpf_func_proto bpf_timer_init_proto = {
        .arg3_type      = ARG_ANYTHING,
 };
 
+static int bpf_async_update_prog_callback(struct bpf_async_cb *cb, void *callback_fn,
+                                         struct bpf_prog *prog)
+{
+       struct bpf_prog *prev;
+
+       /* Acquire a guard reference on prog to prevent it from being freed during the loop */
+       if (prog) {
+               prog = bpf_prog_inc_not_zero(prog);
+               if (IS_ERR(prog))
+                       return PTR_ERR(prog);
+       }
+
+       do {
+               if (prog)
+                       prog = bpf_prog_inc_not_zero(prog);
+               prev = xchg(&cb->prog, prog);
+               rcu_assign_pointer(cb->callback_fn, callback_fn);
+
+               /*
+                * Release previous prog, make sure that if other CPU is contending,
+                * to set bpf_prog, references are not leaked as each iteration acquires and
+                * releases one reference.
+                */
+               if (prev)
+                       bpf_prog_put(prev);
+
+       } while (READ_ONCE(cb->prog) != prog || READ_ONCE(cb->callback_fn) != callback_fn);
+
+       if (prog)
+               bpf_prog_put(prog);
+
+       return 0;
+}
+
 static int __bpf_async_set_callback(struct bpf_async_kern *async, void *callback_fn,
                                    struct bpf_prog *prog)
 {
-       struct bpf_prog *prev;
        struct bpf_async_cb *cb;
        int ret = 0;
 
@@ -1378,22 +1411,7 @@ static int __bpf_async_set_callback(struct bpf_async_kern *async, void *callback
                ret = -EPERM;
                goto out;
        }
-       prev = cb->prog;
-       if (prev != prog) {
-               /* Bump prog refcnt once. Every bpf_timer_set_callback()
-                * can pick different callback_fn-s within the same prog.
-                */
-               prog = bpf_prog_inc_not_zero(prog);
-               if (IS_ERR(prog)) {
-                       ret = PTR_ERR(prog);
-                       goto out;
-               }
-               if (prev)
-                       /* Drop prev prog refcnt when swapping with new prog */
-                       bpf_prog_put(prev);
-               cb->prog = prog;
-       }
-       rcu_assign_pointer(cb->callback_fn, callback_fn);
+       ret = bpf_async_update_prog_callback(cb, callback_fn, prog);
 out:
        __bpf_spin_unlock_irqrestore(&async->lock);
        return ret;
@@ -1453,17 +1471,6 @@ static const struct bpf_func_proto bpf_timer_start_proto = {
        .arg3_type      = ARG_ANYTHING,
 };
 
-static void drop_prog_refcnt(struct bpf_async_cb *async)
-{
-       struct bpf_prog *prog = async->prog;
-
-       if (prog) {
-               bpf_prog_put(prog);
-               async->prog = NULL;
-               rcu_assign_pointer(async->callback_fn, NULL);
-       }
-}
-
 BPF_CALL_1(bpf_timer_cancel, struct bpf_async_kern *, timer)
 {
        struct bpf_hrtimer *t, *cur_t;
@@ -1514,7 +1521,7 @@ BPF_CALL_1(bpf_timer_cancel, struct bpf_async_kern *, timer)
                goto out;
        }
 drop:
-       drop_prog_refcnt(&t->cb);
+       bpf_async_update_prog_callback(&t->cb, NULL, NULL);
 out:
        __bpf_spin_unlock_irqrestore(&timer->lock);
        /* Cancel the timer and wait for associated callback to finish
@@ -1547,7 +1554,7 @@ static struct bpf_async_cb *__bpf_async_cancel_and_free(struct bpf_async_kern *a
        cb = async->cb;
        if (!cb)
                goto out;
-       drop_prog_refcnt(cb);
+       bpf_async_update_prog_callback(cb, NULL, NULL);
        /* The subsequent bpf_timer_start/cancel() helpers won't be able to use
         * this timer, since it won't be initialized.
         */