bpf: Add bpf_guard_preempt() convenience macro

Add bpf_guard_preempt() macro that uses newly introduced
bpf_preempt_disable/enable() kfuncs to guard a critical section.

Signed-off-by: Alexei Starovoitov <ast@kernel.org>
Acked-by: Kumar Kartikeya Dwivedi <memxor@gmail.com>
Link: https://lore.kernel.org/r/20240424225529.16782-1-alexei.starovoitov@gmail.com
Signed-off-by: Martin KaFai Lau <martin.lau@kernel.org>

2 files changed, 24 insertions, 5 deletions

diff --git a/tools/testing/selftests/bpf/bpf_experimental.h b/tools/testing/selftests/bpf/bpf_experimental.h
index 93c5a6c446b3..8b9cc87be4c4 100644
--- a/tools/testing/selftests/bpf/bpf_experimental.h
+++ b/tools/testing/selftests/bpf/bpf_experimental.h
@@ -397,6 +397,28 @@ l_true:												\
 		     , [as]"i"((dst_as << 16) | src_as));
 #endif
 
+void bpf_preempt_disable(void) __weak __ksym;
+void bpf_preempt_enable(void) __weak __ksym;
+
+typedef struct {
+} __bpf_preempt_t;
+
+static inline __bpf_preempt_t __bpf_preempt_constructor(void)
+{
+	__bpf_preempt_t ret = {};
+
+	bpf_preempt_disable();
+	return ret;
+}
+static inline void __bpf_preempt_destructor(__bpf_preempt_t *t)
+{
+	bpf_preempt_enable();
+}
+#define bpf_guard_preempt() \
+	__bpf_preempt_t ___bpf_apply(preempt, __COUNTER__)			\
+	__attribute__((__unused__, __cleanup__(__bpf_preempt_destructor))) =	\
+	__bpf_preempt_constructor()
+
 /* Description
  *	Assert that a conditional expression is true.
  * Returns
diff --git a/tools/testing/selftests/bpf/progs/preempt_lock.c b/tools/testing/selftests/bpf/progs/preempt_lock.c
index 6c637ee01ec4..672fc368d9c4 100644
--- a/tools/testing/selftests/bpf/progs/preempt_lock.c
+++ b/tools/testing/selftests/bpf/progs/preempt_lock.c
@@ -3,9 +3,7 @@
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_tracing.h>
 #include "bpf_misc.h"
-
-void bpf_preempt_disable(void) __ksym;
-void bpf_preempt_enable(void) __ksym;
+#include "bpf_experimental.h"
 
 SEC("?tc")
 __failure __msg("1 bpf_preempt_enable is missing")
@@ -92,8 +90,7 @@ static __noinline void preempt_balance_subprog(void)
 SEC("?tc")
 __success int preempt_balance(struct __sk_buff *ctx)
 {
-	bpf_preempt_disable();
-	bpf_preempt_enable();
+	bpf_guard_preempt();
 	return 0;
 }