Move set_stale from get_subflow() to init().

Signed-off-by: Geliang Tang <geliang.tang@suse.com>
---
 .../selftests/bpf/progs/mptcp_bpf_stale.c     | 33 +++++++++++--------
 1 file changed, 19 insertions(+), 14 deletions(-)

diff --git a/tools/testing/selftests/bpf/progs/mptcp_bpf_stale.c b/tools/testing/selftests/bpf/progs/mptcp_bpf_stale.c
index 08c857f79221..4cd44f911c1f 100644
--- a/tools/testing/selftests/bpf/progs/mptcp_bpf_stale.c
+++ b/tools/testing/selftests/bpf/progs/mptcp_bpf_stale.c
@@ -90,7 +90,12 @@ void BPF_PROG(mptcp_sched_stale_init, struct mptcp_sock *msk)
 	if (!storage)
 		return;
 
+	for (int i = 0; i < MPTCP_SUBFLOWS_MAX; i++)
+		storage->ids[i] = 0;
 	storage->nr = 0;
+
+	mptcp_subflow_set_stale(storage, 2);
+	mptcp_subflow_set_stale(storage, 3);
 }
 
 SEC("struct_ops/mptcp_sched_stale_release")
@@ -102,28 +107,28 @@ void BPF_PROG(mptcp_sched_stale_release, struct mptcp_sock *msk)
 int BPF_STRUCT_OPS(bpf_stale_get_subflow, struct mptcp_sock *msk,
 		   struct mptcp_sched_data *data)
 {
-	struct mptcp_subflow_context *subflow;
 	struct mptcp_stale_storage *storage;
-	int nr = -1;
+	int nr = -1, i;
 
-	storage = bpf_sk_storage_get(&mptcp_stale_map, msk, 0,
-				     BPF_LOCAL_STORAGE_GET_F_CREATE);
+	mptcp_sched_data_set_contexts(msk, data);
+	if (data->subflows == 1) {
+		mptcp_subflow_set_scheduled(mptcp_subflow_ctx_by_pos(data, 0), true);
+		return 0;
+	}
+
+	storage = bpf_sk_storage_get(&mptcp_stale_map, msk, 0, 0);
 	if (!storage)
 		return -1;
 
-	mptcp_sched_data_set_contexts(msk, data);
-
 	/* Handle invalid subflow ids for subflows that have been closed */
-	for (int i = 0; i < storage->nr && i < MPTCP_SUBFLOWS_MAX; i++) {
-		if (!mptcp_subflow_is_active(data, storage->ids[i]))
-			mptcp_subflow_clear_stale(storage, storage->ids[i]);
+	if (data->subflows < storage->nr + 1) {
+		for (i = 0; i < storage->nr && i < MPTCP_SUBFLOWS_MAX; i++) {
+			if (!mptcp_subflow_is_active(data, storage->ids[i]))
+				mptcp_subflow_clear_stale(storage, storage->ids[i]);
+		}
 	}
 
-	subflow = mptcp_subflow_ctx_by_pos(data, 1);
-	if (subflow)
-		mptcp_subflow_set_stale(storage, subflow->subflow_id);
-
-	for (int i = 0; i < data->subflows && i < MPTCP_SUBFLOWS_MAX; i++) {
+	for (i = 0; i < data->subflows && i < MPTCP_SUBFLOWS_MAX; i++) {
 		struct mptcp_subflow_context *subflow;
 
 		subflow = mptcp_subflow_ctx_by_pos(data, i);
-- 
2.35.3