[PATCH bpf-next v1 3/3] selftest/bpf/benchs: Add cpu-affinity for sockmap bench

[Jiayuan Chen <jiayuan.chen@linux.dev>]

Add cpu-affinity for sockmap bench. Also add no-verify args to avoid
validating data for performance enhancements.

Signed-off-by: Jiayuan Chen <jiayuan.chen@linux.dev>
---
 .../selftests/bpf/benchs/bench_sockmap.c      | 35 +++++++++++++++++--
 tools/testing/selftests/bpf/bpf_kfuncs.h      |  6 ++++
 .../selftests/bpf/progs/bench_sockmap_prog.c  |  7 ++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/tools/testing/selftests/bpf/benchs/bench_sockmap.c b/tools/testing/selftests/bpf/benchs/bench_sockmap.c
index 8ebf563a67a2..e004a618822a 100644
--- a/tools/testing/selftests/bpf/benchs/bench_sockmap.c
+++ b/tools/testing/selftests/bpf/benchs/bench_sockmap.c
@@ -43,6 +43,8 @@ enum SOCKMAP_ARG_FLAG {
 	ARG_FW_TX_VERDICT_INGRESS,
 	ARG_FW_TX_VERDICT_EGRESS,
 	ARG_CTL_RX_STRP,
+	ARG_CTL_CPU_AFFINITY,
+	ARG_CTL_NO_VERIFY,
 	ARG_CONSUMER_DELAY_TIME,
 	ARG_PRODUCER_DURATION,
 };
@@ -109,6 +111,8 @@ static struct socmap_ctx {
 	int		delay_consumer;
 	int		prod_run_time;
 	int		strp_size;
+	int		cpu_affinity;
+	int		skip_verify;
 } ctx = {
 	.prod_send	= 0,
 	.user_read	= 0,
@@ -118,6 +122,8 @@ static struct socmap_ctx {
 	.delay_consumer = 0,
 	.prod_run_time	= 0,
 	.strp_size	= 0,
+	.cpu_affinity	= 0,
+	.skip_verify	= 0,
 };
 
 static void bench_sockmap_prog_destroy(void)
@@ -235,11 +241,18 @@ static int create_sockets(void)
 static void validate(void)
 {
 	if (env.consumer_cnt != 2 || env.producer_cnt != 1 ||
-	    !env.affinity)
+	    !env.affinity) {
+		fprintf(stderr, "argument '-c 2 -p 1 -a' is necessary\n");
 		goto err;
+	}
+
+	if (!ctx.cpu_affinity && env.nr_cpus < 4) {
+		fprintf(stderr, "4 CPU are needed to test cpu-affinity\n");
+		goto err;
+	}
+
 	return;
 err:
-	fprintf(stderr, "argument '-c 2 -p 1 -a' is necessary");
 	exit(1);
 }
 
@@ -327,6 +340,9 @@ static void setup(void)
 		exit(1);
 	}
 
+	if (ctx.cpu_affinity)
+		ctx.skel->data->redir_cpu = 3;
+
 	if (create_sockets()) {
 		fprintf(stderr, "create_net_mode error\n");
 		goto err;
@@ -367,9 +383,12 @@ static void measure(struct bench_res *res)
 
 static void verify_data(int *check_pos, char *buf, int rcv)
 {
+	if (ctx.skip_verify)
+		return;
+
 	for (int i = 0 ; i < rcv; i++) {
 		if (buf[i] != snd_data[(*check_pos) % DATA_REPEAT_SIZE]) {
-			fprintf(stderr, "verify data fail");
+			fprintf(stderr, "verify data fail\n");
 			exit(1);
 		}
 		(*check_pos)++;
@@ -553,6 +572,10 @@ static const struct argp_option opts[] = {
 		"delay consumer start"},
 	{ "producer-duration", ARG_PRODUCER_DURATION, "SEC", 0,
 		"producer duration"},
+	{ "cpu-affinity", ARG_CTL_CPU_AFFINITY, NULL, 0,
+		"set cpu-affinity for sockmap backlog thread"},
+	{ "no-verify", ARG_CTL_NO_VERIFY, NULL, 0,
+		"skip data validation for performance enhancements"},
 	{},
 };
 
@@ -571,6 +594,12 @@ static error_t parse_arg(int key, char *arg, struct argp_state *state)
 	case ARG_CTL_RX_STRP:
 		ctx.strp_size = strtol(arg, NULL, 10);
 		break;
+	case ARG_CTL_CPU_AFFINITY:
+		ctx.cpu_affinity = 1;
+		break;
+	case ARG_CTL_NO_VERIFY:
+		ctx.skip_verify = 1;
+		break;
 	default:
 		return ARGP_ERR_UNKNOWN;
 	}
diff --git a/tools/testing/selftests/bpf/bpf_kfuncs.h b/tools/testing/selftests/bpf/bpf_kfuncs.h
index 8215c9b3115e..173329c5d034 100644
--- a/tools/testing/selftests/bpf/bpf_kfuncs.h
+++ b/tools/testing/selftests/bpf/bpf_kfuncs.h
@@ -92,4 +92,10 @@ extern int bpf_set_dentry_xattr(struct dentry *dentry, const char *name__str,
 				const struct bpf_dynptr *value_p, int flags) __ksym __weak;
 extern int bpf_remove_dentry_xattr(struct dentry *dentry, const char *name__str) __ksym __weak;
 
+/* Description
+ *  Set sockmap redir cpu
+ * Returns
+ *  Error code
+ */
+extern int bpf_sk_skb_set_redirect_cpu(struct __sk_buff *skb, int redir_cpu) __ksym;
 #endif
diff --git a/tools/testing/selftests/bpf/progs/bench_sockmap_prog.c b/tools/testing/selftests/bpf/progs/bench_sockmap_prog.c
index 079bf3794b3a..dd1a11cb4f48 100644
--- a/tools/testing/selftests/bpf/progs/bench_sockmap_prog.c
+++ b/tools/testing/selftests/bpf/progs/bench_sockmap_prog.c
@@ -2,11 +2,15 @@
 #include <linux/bpf.h>
 #include <bpf/bpf_helpers.h>
 #include <bpf/bpf_endian.h>
+#include <stdbool.h>
+#include "bpf_kfuncs.h"
 
 long process_byte = 0;
 int  verdict_dir = 0;
 int  dropped = 0;
 int  pkt_size = 0;
+int  redir_cpu = -1;
+
 struct {
 	__uint(type, BPF_MAP_TYPE_SOCKMAP);
 	__uint(max_entries, 20);
@@ -33,6 +37,9 @@ int prog_skb_verdict(struct __sk_buff *skb)
 	int one = 1;
 	int ret =  bpf_sk_redirect_map(skb, &sock_map_rx, one, verdict_dir);
 
+	if (redir_cpu != -1)
+		bpf_sk_skb_set_redirect_cpu(skb, redir_cpu);
+
 	if (ret == SK_DROP)
 		dropped++;
 	__sync_fetch_and_add(&process_byte, skb->len);
-- 
2.47.1