Re: [PATCH bpf-next v2 1/2] bpf: Support private stack for bpf progs

[Yonghong Song <yonghong.song@linux.dev>]

On 7/22/24 10:30 PM, Yonghong Song wrote:
>
> On 7/22/24 6:05 PM, Alexei Starovoitov wrote:
>> On Mon, Jul 22, 2024 at 1:58 PM Andrii Nakryiko
>> <andrii.nakryiko@gmail.com> wrote:
>>> On Fri, Jul 19, 2024 at 8:28 PM Andrii Nakryiko
>>> <andrii.nakryiko@gmail.com> wrote:
>>>> On Thu, Jul 18, 2024 at 1:52 PM Yonghong Song 
>>>> <yonghong.song@linux.dev> wrote:
>>>>> The main motivation for private stack comes from nested
>>>>> scheduler in sched-ext from Tejun. The basic idea is that
>>>>>   - each cgroup will its own associated bpf program,
>>>>>   - bpf program with parent cgroup will call bpf programs
>>>>>     in immediate child cgroups.
>>>>>
>>>>> Let us say we have the following cgroup hierarchy:
>>>>>    root_cg (prog0):
>>>>>      cg1 (prog1):
>>>>>        cg11 (prog11):
>>>>>          cg111 (prog111)
>>>>>          cg112 (prog112)
>>>>>        cg12 (prog12):
>>>>>          cg121 (prog121)
>>>>>          cg122 (prog122)
>>>>>      cg2 (prog2):
>>>>>        cg21 (prog21)
>>>>>        cg22 (prog22)
>>>>>        cg23 (prog23)
>>>>>
>>>>> In the above example, prog0 will call a kfunc which will
>>>>> call prog1 and prog2 to get sched info for cg1 and cg2 and
>>>>> then the information is summarized and sent back to prog0.
>>>>> Similarly, prog11 and prog12 will be invoked in the kfunc
>>>>> and the result will be summarized and sent back to prog1, etc.
>>>>>
>>>>> Currently, for each thread, the x86 kernel allocate 8KB stack.
>>>>> The each bpf program (including its subprograms) has maximum
>>>>> 512B stack size to avoid potential stack overflow.
>>>>> And nested bpf programs increase the risk of stack overflow.
>>>>> To avoid potential stack overflow caused by bpf programs,
>>>>> this patch implemented a private stack so bpf program stack
>>>>> space is allocated dynamically when the program is jited.
>>>>> Such private stack is applied to tracing programs like
>>>>> kprobe/uprobe, perf_event, tracepoint, raw tracepoint and
>>>>> tracing.
>>>>>
>>>>> But more than one instance of the same bpf program may
>>>>> run in the system. To make things simple, percpu private
>>>>> stack is allocated for each program, so if the same program
>>>>> is running on different cpus concurrently, we won't have
>>>>> any issue. Note that the kernel already have logic to prevent
>>>>> the recursion for the same bpf program on the same cpu
>>>>> (kprobe, fentry, etc.).
>>>>>
>>>>> The patch implemented a percpu private stack based approach
>>>>> for x86 arch.
>>>>>    - The stack size will be 0 and any stack access is from
>>>>>      jit-time allocated percpu storage.
>>>>>    - In the beginning of jit, r9 is used to save percpu
>>>>>      private stack pointer.
>>>>>    - Each rbp in the bpf asm insn is replaced by r9.
>>>>>    - For each call, push r9 before the call and pop r9
>>>>>      after the call to preserve r9 value.
>>>>>
>>>>> Compared to previous RFC patch [1], this patch added
>>>>> some conditions to enable private stack, e.g., verifier
>>>>> calculated stack size, prog type, etc. The new patch
>>>>> also added a performance test to compare private stack
>>>>> vs. no private stack.
>>>>>
>>>>> The following are some code example to illustrate the idea
>>>>> for selftest cgroup_skb_sk_lookup:
>>>>>
>>>>>     the existing code                        the private-stack 
>>>>> approach code
>>>>>     endbr64                                  endbr64
>>>>>     nop    DWORD PTR [rax+rax*1+0x0]         nop    DWORD PTR 
>>>>> [rax+rax*1+0x0]
>>>>>     xchg   ax,ax                             xchg   ax,ax
>>>>>     push   rbp                               push   rbp
>>>>>     mov    rbp,rsp                           mov rbp,rsp
>>>>>     endbr64                                  endbr64
>>>>>     sub    rsp,0x68
>>>>>     push   rbx                               push   rbx
>>>>>     ...                                      ...
>>>>>     ...                                      mov r9d,0x8c1c860
>>>>>     ...                                      add r9,QWORD PTR 
>>>>> gs:0x21a00
>>>>>     ...                                      ...
>>>>>     mov    rdx,rbp                           mov    rdx, r9
>>>>>     add    rdx,0xffffffffffffffb4 rdx,0xffffffffffffffb4
>>>>>     ...                                      ...
>>>>>     mov    ecx,0x28                          mov ecx,0x28
>>>>>                                              push   r9
>>>>>     call   0xffffffffe305e474                call 0xffffffffe305e524
>>>>>                                              pop    r9
>>>>>     mov    rdi,rax                           mov rdi,rax
>>>>>     ...                                      ...
>>>>>     movzx  rdi,BYTE PTR [rbp-0x46]           movzx rdi,BYTE PTR 
>>>>> [r9-0x46]
>>>>>     ...                                      ...
>>>>>
>>>> Eduard nerd-sniped me today with this a bit... :)
>>>>
>>>> I have a few questions and suggestions.
>>>>
>>>> So it seems like each *subprogram* (not the entire BPF program) gets
>>>> its own per-CPU private stack allocation. Is that intentional? That
>>>> seems a bit unnecessary. It also prevents any sort of actual
>>>> recursion. Not sure if it's possible to write recursive BPF subprogram
>>>> today, verifier seems to reject obvious limited recursion cases, but
>>>> still, eventually we might need/want to support that, and this will be
>>>> just another hurdle to overcome (so it's best to avoid adding it in
>>>> the first place).
>>>>
>>>> I'm sure Eduard is going to try something like below and it will
>>>> probably break badly (I haven't tried, sorry):
>>>>
>>>> int entry(void *ctx);
>>>>
>>>> struct {
>>>>          __uint(type, BPF_MAP_TYPE_PROG_ARRAY);
>>>>          __uint(max_entries, 1);
>>>>          __uint(key_size, sizeof(__u32));
>>>>          __array(values, int (void *));
>>>> } prog_array_init SEC(".maps") = {
>>>>          .values = {
>>>>                  [0] = (void *)&entry,
>>>>          },
>>>> };
>>>>
>>>> static __noinline int subprog1(void)
>>>> {
>>>>      <some state on the stack>
>>>>
>>>>      /* here entry will replace subprog1, and so we'll have
>>>>       * entry -> entry -> entry -> ..... <tail call limit> -> subprog1
>>>>       */
>>>>      bpf_tail_call(ctx, &prog_array_init, 0);
>>>>
>>>>      return 0;
>>>> }
>>>>
>>>>
>>>> SEC("raw_tp/sys_enter")
>>>> int entry(void *ctx)
>>>> {
>>>>       <some state on the stack>
>>>>
>>>>       subprog1();
>>>> }
>>>>
>>>> And we effectively have limited recursion where the entry's stack
>>>> state is clobbered, no?
>>>>
>>>> So it seems like we need to support recursion.
>>>>
>>> How come everyone just completely ignored the main point of my entire
>>> email and a real problem that has to be solved?...
>>>
>>> Anyways, I did write a below program:
>>>
>>> $ cat minimal.bpf.c
>>> // SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause
>>> /* Copyright (c) 2020 Facebook */
>>> #include <linux/bpf.h>
>>> #include <bpf/bpf_helpers.h>
>>>
>>> char LICENSE[] SEC("license") = "Dual BSD/GPL";
>>>
>>> int my_pid = 0;
>>>
>>> int handle_tp(void *ctx);
>>>
>>> struct {
>>>          __uint(type, BPF_MAP_TYPE_PROG_ARRAY);
>>>          __uint(max_entries, 1);
>>>          __uint(key_size, sizeof(__u32));
>>>          __array(values, int (void *));
>>> } prog_array_init SEC(".maps") = {
>>>          .values = {
>>>                  [0] = (void *)&handle_tp,
>>>          },
>>> };
>>>
>>> static __noinline int subprog(void *ctx)
>>> {
>>>          static int cnt;
>>>
>>>          cnt++;
>>>
>>>          bpf_printk("SUBPROG - BEFORE %d", cnt);
>>>
>>>          bpf_tail_call(ctx, &prog_array_init, 0);
>>>
>>>          bpf_printk("SUBPROG - AFTER %d", cnt);
>>>
>>>      return 0;
>>> }
>>>
>>> SEC("tp/syscalls/sys_enter_write")
>>> int handle_tp(void *ctx)
>>> {
>>>          static int cnt;
>>>          int pid = bpf_get_current_pid_tgid() >> 32;
>>>
>>>          if (pid != my_pid)
>>>                  return 0;
>>>
>>>          cnt++;
>>>
>>>          bpf_printk("ENTRY - BEFORE %d", cnt);
>>>
>>>          subprog(ctx);
>>>
>>>          bpf_printk("ENTRY - AFTER %d", cnt);
>>>
>>>          return 0;
>>> }
>>>
>>>
>>> And triggered one write syscall, getting the log above. You can see
>>> that only subprogs are replaced (we only get "SUBPROG - AFTER 34" due
>>> to the tail call limit). And we do indeed get lots of entry program
>>> recurrence.
>>>
>>> We *need to support recursion* is my main point.
>> Not quite.
>> It's not a recursion. The stack collapsed/gone/wiped out before 
>> tail_call.
>> static int cnt counts stuff because it's static.
>>
>> So we don't need to support recursion with private stack,
>> but tail_calls and private stack are buggy indeed.
>>
>> emit_bpf_tail_call*() shouldn't be adjusting 'rsp' when the private
>> stack is used.
>
> Right, stack_depth argument in 
> emit_bpf_tail_call_direct()/emit_bpf_tail_call_indirect()
> should be 0 if private stack is used. Will fix in next revision.

Actually, the current implementation is correct. We already set stack_depth to be 0
if private stack is used. So we should be fine here.