xdpsock: Reuse immediate exit logic from xdp_sample

Use the immediate exit logic from xdp_sample in xdpsock as
well (exporting the function that checks the envvar), so we can use it
in tests.

Signed-off-by: Toke Høiland-Jørgensen <[email protected]>

Author: tohojo

lib/util/xdp_sample.c

@@ -1559,21 +1559,26 @@ static int sample_timer_cb(int timerfd, struct stats_record **rec,
 	return print_stats(rec, prev);
 }
 
-int sample_run(unsigned int interval, void (*post_cb)(void *), void *ctx)
+bool sample_immediate_exit(void)
 {
-	struct timespec ts = { interval, 0 };
-	struct itimerspec its = { ts, ts };
-	struct stats_record *rec, *prev;
-	struct pollfd pfd[2] = {};
-	bool imm_exit = false;
 	const char *envval;
-	int timerfd, ret;
 
 	envval = secure_getenv("XDP_SAMPLE_IMMEDIATE_EXIT");
 	if (envval && envval[0] == '1' && envval[1] == '\0') {
 		pr_debug("XDP_SAMPLE_IMMEDIATE_EXIT envvar set, exiting immediately after setup\n");
-		imm_exit = true;
+		return true;
 	}
+	return false;
+}
+
+int sample_run(unsigned int interval, void (*post_cb)(void *), void *ctx)
+{
+	bool imm_exit = sample_immediate_exit();
+	struct timespec ts = { interval, 0 };
+	struct itimerspec its = { ts, ts };
+	struct stats_record *rec, *prev;
+	struct pollfd pfd[2] = {};
+	int timerfd, ret;
 
 	if (!interval) {
 		pr_warn("Incorrect interval 0\n");

lib/util/xdp_sample.h

@@ -44,6 +44,7 @@ int sample_run(unsigned int interval, void (*post_cb)(void *), void *ctx);
 bool sample_is_compat(enum sample_compat compat_value);
 bool sample_probe_cpumap_compat(void);
 bool sample_probe_xdp_load_bytes(void);
+bool sample_immediate_exit(void);
 void sample_check_cpumap_compat(struct bpf_program *prog,
 				struct bpf_program *prog_compat);
 

lib/util/xdpsock.c

@@ -41,6 +41,7 @@
 #include <bpf/libbpf.h>
 #include <bpf/bpf.h>
 #include "xdpsock.h"
+#include "xdp_sample.h"
 #include "logging.h"
 #include "util.h"
 
@@ -529,6 +530,10 @@ static bool is_benchmark_done(struct xsk_ctx *ctx)
 		if (dt >= ctx->duration)
 			ctx->benchmark_done = true;
 	}
+
+	if (sample_immediate_exit())
+		ctx->benchmark_done = true;
+
 	return ctx->benchmark_done;
 }