[PATCH 2/2] selftests/bpf: test a BPF CC writing app_limited

[Yixin Shen <bobankhshen@xxxxxxxxx>]

Test whether a TCP CC implemented in BPF is allowed to write
app_limited in struct tcp_sock. This is already allowed for
the built-in TCP CC.

Signed-off-by: Yixin Shen <bobankhshen@xxxxxxxxx>
---
 .../selftests/bpf/prog_tests/bpf_tcp_ca.c     | 19 +++++
 .../bpf/progs/tcp_ca_write_app_limited.c      | 71 +++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 tools/testing/selftests/bpf/progs/tcp_ca_write_app_limited.c

diff --git a/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c b/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
index a53c254c6058..1911ed04d4cf 100644
--- a/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
+++ b/tools/testing/selftests/bpf/prog_tests/bpf_tcp_ca.c
@@ -11,6 +11,7 @@
 #include "tcp_ca_update.skel.h"
 #include "bpf_dctcp_release.skel.h"
 #include "tcp_ca_write_sk_pacing.skel.h"
+#include "tcp_ca_write_app_limited.skel.h"
 #include "tcp_ca_incompl_cong_ops.skel.h"
 #include "tcp_ca_unsupp_cong_op.skel.h"
 
@@ -346,6 +347,22 @@ static void test_write_sk_pacing(void)
 	tcp_ca_write_sk_pacing__destroy(skel);
 }
 
+static void test_write_app_limited(void)
+{
+	struct tcp_ca_write_app_limited *skel;
+	struct bpf_link *link;
+
+	skel = tcp_ca_write_app_limited__open_and_load();
+	if (!ASSERT_OK_PTR(skel, "open_and_load"))
+		return;
+
+	link = bpf_map__attach_struct_ops(skel->maps.write_app_limited);
+	ASSERT_OK_PTR(link, "attach_struct_ops");
+
+	bpf_link__destroy(link);
+	tcp_ca_write_app_limited__destroy(skel);
+}
+
 static void test_incompl_cong_ops(void)
 {
 	struct tcp_ca_incompl_cong_ops *skel;
@@ -545,6 +562,8 @@ void test_bpf_tcp_ca(void)
 		test_rel_setsockopt();
 	if (test__start_subtest("write_sk_pacing"))
 		test_write_sk_pacing();
+	if (test__start_subtest("write_app_limited"))
+		test_write_app_limited();
 	if (test__start_subtest("incompl_cong_ops"))
 		test_incompl_cong_ops();
 	if (test__start_subtest("unsupp_cong_op"))
diff --git a/tools/testing/selftests/bpf/progs/tcp_ca_write_app_limited.c b/tools/testing/selftests/bpf/progs/tcp_ca_write_app_limited.c
new file mode 100644
index 000000000000..de5c9b5045a1
--- /dev/null
+++ b/tools/testing/selftests/bpf/progs/tcp_ca_write_app_limited.c
@@ -0,0 +1,71 @@
+// SPDX-License-Identifier: GPL-2.0
+
+#include "vmlinux.h"
+
+#include <bpf/bpf_helpers.h>
+#include <bpf/bpf_tracing.h>
+
+char _license[] SEC("license") = "GPL";
+
+#define USEC_PER_SEC 1000000UL
+
+#define min(a, b) ((a) < (b) ? (a) : (b))
+
+static inline struct tcp_sock *tcp_sk(const struct sock *sk)
+{
+	return (struct tcp_sock *)sk;
+}
+
+static inline unsigned int tcp_left_out(const struct tcp_sock *tp)
+{
+	return tp->sacked_out + tp->lost_out;
+}
+
+static inline unsigned int tcp_packets_in_flight(const struct tcp_sock *tp)
+{
+	return tp->packets_out - tcp_left_out(tp) + tp->retrans_out;
+}
+
+SEC("struct_ops/write_app_limited_init")
+void BPF_PROG(write_app_limited_init, struct sock *sk)
+{
+#ifdef ENABLE_ATOMICS_TESTS
+	__sync_bool_compare_and_swap(&sk->sk_pacing_status, SK_PACING_NONE,
+				     SK_PACING_NEEDED);
+#else
+	sk->sk_pacing_status = SK_PACING_NEEDED;
+#endif
+}
+
+SEC("struct_ops/write_app_limited_cong_control")
+void BPF_PROG(write_app_limited_cong_control, struct sock *sk,
+	      const struct rate_sample *rs)
+{
+	struct tcp_sock *tp = tcp_sk(sk);
+	unsigned long rate =
+		((tp->snd_cwnd * tp->mss_cache * USEC_PER_SEC) << 3) /
+		(tp->srtt_us ?: 1U << 3);
+	sk->sk_pacing_rate = min(rate, sk->sk_max_pacing_rate);
+	tp->app_limited = (tp->delivered + tcp_packets_in_flight(tp)) ?: 1;
+}
+
+SEC("struct_ops/write_app_limited_ssthresh")
+__u32 BPF_PROG(write_app_limited_ssthresh, struct sock *sk)
+{
+	return tcp_sk(sk)->snd_ssthresh;
+}
+
+SEC("struct_ops/write_app_limited_undo_cwnd")
+__u32 BPF_PROG(write_app_limited_undo_cwnd, struct sock *sk)
+{
+	return tcp_sk(sk)->snd_cwnd;
+}
+
+SEC(".struct_ops")
+struct tcp_congestion_ops write_app_limited = {
+	.init = (void *)write_app_limited_init,
+	.cong_control = (void *)write_app_limited_cong_control,
+	.ssthresh = (void *)write_app_limited_ssthresh,
+	.undo_cwnd = (void *)write_app_limited_undo_cwnd,
+	.name = "bpf_w_app_limit",
+};
-- 
2.25.1