selftests/bpf: Test preemption between bpf_obj_new() and bpf_obj_drop()

// SPDX-License-Identifier: GPL-2.0

/* Copyright (C) 2023. Huawei Technologies Co., Ltd */

#define _GNU_SOURCE

#include <sched.h>

#include <pthread.h>

#include <stdbool.h>

#include <test_progs.h>

#include "preempted_bpf_ma_op.skel.h"

#define ALLOC_THREAD_NR 4

#define ALLOC_LOOP_NR 512

struct alloc_ctx {

	/* output */

	int run_err;

	/* input */

	int fd;

	bool *nomem_err;

};

static void *run_alloc_prog(void *data)

{

	struct alloc_ctx *ctx = data;

	cpu_set_t cpu_set;

	int i;

	CPU_ZERO(&cpu_set);

	CPU_SET(0, &cpu_set);

	pthread_setaffinity_np(pthread_self(), sizeof(cpu_set), &cpu_set);

	for (i = 0; i < ALLOC_LOOP_NR && !*ctx->nomem_err; i++) {

		LIBBPF_OPTS(bpf_test_run_opts, topts);

		int err;

		err = bpf_prog_test_run_opts(ctx->fd, &topts);

		ctx->run_err |= err | topts.retval;

	}

	return NULL;

}

void test_preempted_bpf_ma_op(void)

{

	struct alloc_ctx ctx[ALLOC_THREAD_NR];

	struct preempted_bpf_ma_op *skel;

	pthread_t tid[ALLOC_THREAD_NR];

	int i, err;

	skel = preempted_bpf_ma_op__open_and_load();

	if (!ASSERT_OK_PTR(skel, "open_and_load"))

		return;

	err = preempted_bpf_ma_op__attach(skel);

	if (!ASSERT_OK(err, "attach"))

		goto out;

	for (i = 0; i < ARRAY_SIZE(ctx); i++) {

		struct bpf_program *prog;

		char name[8];

		snprintf(name, sizeof(name), "test%d", i);

		prog = bpf_object__find_program_by_name(skel->obj, name);

		if (!ASSERT_OK_PTR(prog, "no test prog"))

			goto out;

		ctx[i].run_err = 0;

		ctx[i].fd = bpf_program__fd(prog);

		ctx[i].nomem_err = &skel->bss->nomem_err;

	}

	memset(tid, 0, sizeof(tid));

	for (i = 0; i < ARRAY_SIZE(tid); i++) {

		err = pthread_create(&tid[i], NULL, run_alloc_prog, &ctx[i]);

		if (!ASSERT_OK(err, "pthread_create"))

			break;

	}

	for (i = 0; i < ARRAY_SIZE(tid); i++) {

		if (!tid[i])

			break;

		pthread_join(tid[i], NULL);

		ASSERT_EQ(ctx[i].run_err, 0, "run prog err");

	}

	ASSERT_FALSE(skel->bss->nomem_err, "ENOMEM");

out:

	preempted_bpf_ma_op__destroy(skel);

}

// SPDX-License-Identifier: GPL-2.0

/* Copyright (C) 2023. Huawei Technologies Co., Ltd */

#include <vmlinux.h>

#include <bpf/bpf_tracing.h>

#include <bpf/bpf_helpers.h>

#include "bpf_experimental.h"

struct bin_data {

	char data[256];

	struct bpf_spin_lock lock;

};

struct map_value {

	struct bin_data __kptr * data;

};

struct {

	__uint(type, BPF_MAP_TYPE_ARRAY);

	__type(key, int);

	__type(value, struct map_value);

	__uint(max_entries, 2048);

} array SEC(".maps");

char _license[] SEC("license") = "GPL";

bool nomem_err = false;

static int del_array(unsigned int i, int *from)

{

	struct map_value *value;

	struct bin_data *old;

	value = bpf_map_lookup_elem(&array, from);

	if (!value)

		return 1;

	old = bpf_kptr_xchg(&value->data, NULL);

	if (old)

		bpf_obj_drop(old);

	(*from)++;

	return 0;

}

static int add_array(unsigned int i, int *from)

{

	struct bin_data *old, *new;

	struct map_value *value;

	value = bpf_map_lookup_elem(&array, from);

	if (!value)

		return 1;

	new = bpf_obj_new(typeof(*new));

	if (!new) {

		nomem_err = true;

		return 1;

	}

	old = bpf_kptr_xchg(&value->data, new);

	if (old)

		bpf_obj_drop(old);

	(*from)++;

	return 0;

}

static void del_then_add_array(int from)

{

	int i;

	i = from;

	bpf_loop(512, del_array, &i, 0);

	i = from;

	bpf_loop(512, add_array, &i, 0);

}

SEC("fentry/bpf_fentry_test1")

int BPF_PROG2(test0, int, a)

{

	del_then_add_array(0);

	return 0;

}

SEC("fentry/bpf_fentry_test2")

int BPF_PROG2(test1, int, a, u64, b)

{

	del_then_add_array(512);

	return 0;

}

SEC("fentry/bpf_fentry_test3")

int BPF_PROG2(test2, char, a, int, b, u64, c)

{

	del_then_add_array(1024);

	return 0;

}

SEC("fentry/bpf_fentry_test4")

int BPF_PROG2(test3, void *, a, char, b, int, c, u64, d)

{

	del_then_add_array(1536);

	return 0;

}