Re: [RFC PATCH bpf-next 11/11] selftests/bpf: Add rbtree map tests

[Dave Marchevsky <davemarchevsky@fb.com>]

On 7/28/22 3:18 AM, Yonghong Song wrote:   
> 
> 
> On 7/22/22 11:34 AM, Dave Marchevsky wrote:
>> Add tests demonstrating happy path of rbtree map usage as well as
>> exercising numerous failure paths and conditions. Structure of failing
>> test runner is based on dynptr tests.
>>
>> Signed-off-by: Dave Marchevsky <davemarchevsky@fb.com>
>> ---
>>   .../selftests/bpf/prog_tests/rbtree_map.c     | 164 ++++++++++++
>>   .../testing/selftests/bpf/progs/rbtree_map.c  | 111 ++++++++
>>   .../selftests/bpf/progs/rbtree_map_fail.c     | 236 ++++++++++++++++++
>>   .../bpf/progs/rbtree_map_load_fail.c          |  24 ++
>>   4 files changed, 535 insertions(+)
>>   create mode 100644 tools/testing/selftests/bpf/prog_tests/rbtree_map.c
>>   create mode 100644 tools/testing/selftests/bpf/progs/rbtree_map.c
>>   create mode 100644 tools/testing/selftests/bpf/progs/rbtree_map_fail.c
>>   create mode 100644 tools/testing/selftests/bpf/progs/rbtree_map_load_fail.c
>>
>> diff --git a/tools/testing/selftests/bpf/prog_tests/rbtree_map.c b/tools/testing/selftests/bpf/prog_tests/rbtree_map.c
>> new file mode 100644
>> index 000000000000..17cadcd05ee4
>> --- /dev/null
>> +++ b/tools/testing/selftests/bpf/prog_tests/rbtree_map.c
>> @@ -0,0 +1,164 @@
>> +// SPDX-License-Identifier: GPL-2.0
>> +/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
>> +
>> +#include <sys/syscall.h>
>> +#include <test_progs.h>
>> +#include "rbtree_map.skel.h"
>> +#include "rbtree_map_fail.skel.h"
>> +#include "rbtree_map_load_fail.skel.h"
>> +
>> +static size_t log_buf_sz = 1048576; /* 1 MB */
>> +static char obj_log_buf[1048576];
>> +
>> +static struct {
>> +    const char *prog_name;
>> +    const char *expected_err_msg;
>> +} rbtree_prog_load_fail_tests[] = {
>> +    {"rb_node__field_store", "only read is supported"},
>> +    {"rb_node__alloc_no_add", "Unreleased reference id=2 alloc_insn=3"},
>> +    {"rb_node__two_alloc_one_add", "Unreleased reference id=2 alloc_insn=3"},
>> +    {"rb_node__remove_no_free", "Unreleased reference id=5 alloc_insn=28"},
>> +    {"rb_tree__add_wrong_type", "rbtree: R2 is of type task_struct but node_data is expected"},
>> +    {"rb_tree__conditional_release_helper_usage",
>> +        "R2 type=ptr_cond_rel_ expected=ptr_"},
>> +};
>> +
>> +void test_rbtree_map_load_fail(void)
>> +{
>> +    struct rbtree_map_load_fail *skel;
>> +
>> +    skel = rbtree_map_load_fail__open_and_load();
>> +    if (!ASSERT_ERR_PTR(skel, "rbtree_map_load_fail__open_and_load"))
>> +        rbtree_map_load_fail__destroy(skel);
>> +}
>> +
>> +static void verify_fail(const char *prog_name, const char *expected_err_msg)
>> +{
>> +    LIBBPF_OPTS(bpf_object_open_opts, opts);
>> +    struct rbtree_map_fail *skel;
>> +    struct bpf_program *prog;
>> +    int err;
>> +
>> +    opts.kernel_log_buf = obj_log_buf;
>> +    opts.kernel_log_size = log_buf_sz;
>> +    opts.kernel_log_level = 1;
>> +
>> +    skel = rbtree_map_fail__open_opts(&opts);
>> +    if (!ASSERT_OK_PTR(skel, "rbtree_map_fail__open_opts"))
>> +        goto cleanup;
>> +
>> +    prog = bpf_object__find_program_by_name(skel->obj, prog_name);
>> +    if (!ASSERT_OK_PTR(prog, "bpf_object__find_program_by_name"))
>> +        goto cleanup;
>> +
>> +    bpf_program__set_autoload(prog, true);
>> +    err = rbtree_map_fail__load(skel);
>> +    if (!ASSERT_ERR(err, "unexpected load success"))
>> +        goto cleanup;
>> +
>> +    if (!ASSERT_OK_PTR(strstr(obj_log_buf, expected_err_msg), "expected_err_msg")) {
>> +        fprintf(stderr, "Expected err_msg: %s\n", expected_err_msg);
>> +        fprintf(stderr, "Verifier output: %s\n", obj_log_buf);
>> +    }
>> +
>> +cleanup:
>> +    rbtree_map_fail__destroy(skel);
>> +}
>> +
>> +void test_rbtree_map_alloc_node__size_too_small(void)
>> +{
>> +    struct rbtree_map_fail *skel;
>> +    struct bpf_program *prog;
>> +    struct bpf_link *link;
>> +    int err;
>> +
>> +    skel = rbtree_map_fail__open();
>> +    if (!ASSERT_OK_PTR(skel, "rbtree_map_fail__open"))
>> +        goto cleanup;
>> +
>> +    prog = skel->progs.alloc_node__size_too_small;
>> +    bpf_program__set_autoload(prog, true);
>> +
>> +    err = rbtree_map_fail__load(skel);
>> +    if (!ASSERT_OK(err, "unexpected load fail"))
>> +        goto cleanup;
>> +
>> +    link = bpf_program__attach(skel->progs.alloc_node__size_too_small);
>> +    if (!ASSERT_OK_PTR(link, "link"))
>> +        goto cleanup;
>> +
>> +    syscall(SYS_getpgid);
>> +
>> +    ASSERT_EQ(skel->bss->size_too_small__alloc_fail, 1, "alloc_fail");
>> +
>> +    bpf_link__destroy(link);
>> +cleanup:
>> +    rbtree_map_fail__destroy(skel);
>> +}
>> +
>> +void test_rbtree_map_add_node__no_lock(void)
>> +{
>> +    struct rbtree_map_fail *skel;
>> +    struct bpf_program *prog;
>> +    struct bpf_link *link;
>> +    int err;
>> +
>> +    skel = rbtree_map_fail__open();
>> +    if (!ASSERT_OK_PTR(skel, "rbtree_map_fail__open"))
>> +        goto cleanup;
>> +
>> +    prog = skel->progs.add_node__no_lock;
>> +    bpf_program__set_autoload(prog, true);
>> +
>> +    err = rbtree_map_fail__load(skel);
>> +    if (!ASSERT_OK(err, "unexpected load fail"))
>> +        goto cleanup;
>> +
>> +    link = bpf_program__attach(skel->progs.add_node__no_lock);
>> +    if (!ASSERT_OK_PTR(link, "link"))
>> +        goto cleanup;
>> +
>> +    syscall(SYS_getpgid);
>> +
>> +    ASSERT_EQ(skel->bss->no_lock_add__fail, 1, "alloc_fail");
>> +
>> +    bpf_link__destroy(link);
>> +cleanup:
>> +    rbtree_map_fail__destroy(skel);
>> +}
>> +
>> +void test_rbtree_map_prog_load_fail(void)
>> +{
>> +    int i;
>> +
>> +    for (i = 0; i < ARRAY_SIZE(rbtree_prog_load_fail_tests); i++) {
>> +        if (!test__start_subtest(rbtree_prog_load_fail_tests[i].prog_name))
>> +            continue;
>> +
>> +        verify_fail(rbtree_prog_load_fail_tests[i].prog_name,
>> +                rbtree_prog_load_fail_tests[i].expected_err_msg);
>> +    }
>> +}
>> +
>> +void test_rbtree_map(void)
>> +{
>> +    struct rbtree_map *skel;
>> +    struct bpf_link *link;
>> +
>> +    skel = rbtree_map__open_and_load();
>> +    if (!ASSERT_OK_PTR(skel, "rbtree_map__open_and_load"))
>> +        goto cleanup;
>> +
>> +    link = bpf_program__attach(skel->progs.check_rbtree);
>> +    if (!ASSERT_OK_PTR(link, "link"))
>> +        goto cleanup;
>> +
>> +    for (int i = 0; i < 100; i++)
>> +        syscall(SYS_getpgid);
>> +
>> +    ASSERT_EQ(skel->bss->calls, 100, "calls_equal");
>> +
>> +    bpf_link__destroy(link);
>> +cleanup:
>> +    rbtree_map__destroy(skel);
>> +}
>> diff --git a/tools/testing/selftests/bpf/progs/rbtree_map.c b/tools/testing/selftests/bpf/progs/rbtree_map.c
>> new file mode 100644
>> index 000000000000..0cd467838f6e
>> --- /dev/null
>> +++ b/tools/testing/selftests/bpf/progs/rbtree_map.c
>> @@ -0,0 +1,111 @@
>> +// SPDX-License-Identifier: GPL-2.0
>> +/* Copyright (c) 2022 Meta Platforms, Inc. and affiliates. */
>> +
>> +#include "vmlinux.h"
>> +#include <bpf/bpf_helpers.h>
>> +#include "bpf_misc.h"
>> +
>> +struct node_data {
>> +    struct rb_node node;
>> +    __u32 one;
>> +    __u32 two;
>> +};
>> +
>> +struct {
>> +    __uint(type, BPF_MAP_TYPE_RBTREE);
>> +    __type(value, struct node_data);
>> +} rbtree SEC(".maps");
>> +
>> +long calls;
>> +
>> +static bool less(struct rb_node *a, const struct rb_node *b)
>> +{
>> +    struct node_data *node_a;
>> +    struct node_data *node_b;
>> +
>> +    node_a = container_of(a, struct node_data, node);
>> +    node_b = container_of(b, struct node_data, node);
>> +
>> +    return node_a->one < node_b->one;
>> +}
>> +
>> +// Key = node_datq
>> +static int cmp(const void *key, const struct rb_node *b)
>> +{
>> +    struct node_data *node_a;
>> +    struct node_data *node_b;
>> +
>> +    node_a = container_of(key, struct node_data, node);
>> +    node_b = container_of(b, struct node_data, node);
>> +
>> +    return node_b->one - node_a->one;
>> +}
>> +
>> +// Key = just node_data.one
>> +static int cmp2(const void *key, const struct rb_node *b)
>> +{
>> +    __u32 one;
>> +    struct node_data *node_b;
>> +
>> +    one = *(__u32 *)key;
>> +    node_b = container_of(b, struct node_data, node);
>> +
>> +    return node_b->one - one;
>> +}
>> +
>> +SEC("fentry/" SYS_PREFIX "sys_getpgid")
>> +int check_rbtree(void *ctx)
>> +{
>> +    struct node_data *node, *found, *ret;
>> +    struct node_data popped;
>> +    struct node_data search;
>> +    __u32 search2;
>> +
>> +    node = bpf_rbtree_alloc_node(&rbtree, sizeof(struct node_data));
> 
> If I understand correctly, bpf_rtbree_alloc_node() may cause reschedule
> inside the function. So, the program should be sleepable, right?
> 

My mistake, will change alloc flag to GFP_NOWAIT.

>> +    if (!node)
>> +        return 0;
>> +
>> +    node->one = calls;
>> +    node->two = 6;
>> +    bpf_rbtree_lock(bpf_rbtree_get_lock(&rbtree));
>> +
>> +    ret = (struct node_data *)bpf_rbtree_add(&rbtree, node, less);
>> +    if (!ret) {
>> +        bpf_rbtree_free_node(&rbtree, node);
>> +        goto unlock_ret;
>> +    }
>> +
>> +    bpf_rbtree_unlock(bpf_rbtree_get_lock(&rbtree));
>> +
>> +    bpf_rbtree_lock(bpf_rbtree_get_lock(&rbtree));
>> +
>> +    search.one = calls;
>> +    found = (struct node_data *)bpf_rbtree_find(&rbtree, &search, cmp);
>> +    if (!found)
>> +        goto unlock_ret;
>> +
>> +    int node_ct = 0;
>> +    struct node_data *iter = (struct node_data *)bpf_rbtree_first(&rbtree);
>> +
>> +    while (iter) {
>> +        node_ct++;
>> +        iter = (struct node_data *)bpf_rbtree_next(&rbtree, iter);
>> +    }
>> +
>> +    ret = (struct node_data *)bpf_rbtree_remove(&rbtree, found);
>> +    if (!ret)
>> +        goto unlock_ret;
>> +
>> +    bpf_rbtree_unlock(bpf_rbtree_get_lock(&rbtree));
>> +
>> +    bpf_rbtree_free_node(&rbtree, ret);
>> +
>> +    __sync_fetch_and_add(&calls, 1);
>> +    return 0;
>> +
>> +unlock_ret:
>> +    bpf_rbtree_unlock(bpf_rbtree_get_lock(&rbtree));
>> +    return 0;
>> +}
>> +
> [...]