From 6f20431a00ca5600737ab7b072ed5fb54f722bfa Mon Sep 17 00:00:00 2001 From: Tianyi Gao Date: Wed, 21 Oct 2026 22:45:15 +0000 Subject: [PATCH] Wait for concurrent reapers in homa_sock_destroy * homa_rpc_reap takes RPCs off dead_rpcs under the socket lock and frees them after unlocking; since d00a93b6 that is also when it subtracts their frag_bytes from sk_wmem_alloc. homa_sock_destroy only waited for dead_rpcs to be empty (its comment wrongly said that all references were gone), so while homa_timer or softirq still freed the last batch it hit WARN_ON_ONCE(sk_wmem_alloc != 1) and freed the buffer pool under the reaper. * Count the reaps in progress in hsk->reapers. homa_sock_destroy now increments protect_count so that no new reap starts, and waits for hsk->reapers to reach zero before it checks sk_wmem_alloc and frees the pool. --- homa_rpc.c | 4 ++++ homa_sock.c | 18 ++++++++++++++++-- homa_sock.h | 9 +++++++++ test/unit_homa_rpc.c | 25 +++++++++++++++++++++++++ test/unit_homa_sock.c | 25 +++++++++++++++++++++++++ 5 files changed, 79 insertions(+), 2 deletions(-) diff --git a/homa_rpc.c b/homa_rpc.c index bdbe8752..9b3a40ba 100644 --- a/homa_rpc.c +++ b/homa_rpc.c @@ -514,6 +514,7 @@ int homa_rpc_reap(struct homa_sock *hsk) homa_sock_unlock(hsk); return 0; } + atomic_inc(&hsk->reapers); /* Collect freeable RPCs. */ list_for_each_entry_safe(rpc, tmp, &hsk->dead_rpcs, dead_links) { @@ -637,6 +638,9 @@ int homa_rpc_reap(struct homa_sock *hsk) homa_sock_wakeup_wmem(hsk); tt_record3("reaped %d rpcs; %d dead frags remain for port %d", num_rpcs, hsk->dead_frags, hsk->port); + + /* Pairs with atomic_read_acquire in homa_sock_destroy. */ + atomic_dec_return_release(&hsk->reapers); return !checked_all_rpcs; } diff --git a/homa_sock.c b/homa_sock.c index 85c2c9b9..7078b41d 100644 --- a/homa_sock.c +++ b/homa_sock.c @@ -215,6 +215,7 @@ int homa_sock_init(struct homa_sock *hsk) hsk->error_msg = "no error"; spin_lock_init(&hsk->lock); atomic_set(&hsk->protect_count, 0); + atomic_set(&hsk->reapers, 0); INIT_LIST_HEAD(&hsk->active_rpcs); INIT_LIST_HEAD(&hsk->dead_rpcs); INIT_LIST_HEAD(&hsk->ready_rpcs); @@ -402,8 +403,9 @@ void homa_sock_shutdown(struct homa_sock *hsk) /** * homa_sock_destroy() - Release all of the internal resources associated - * with a socket; is invoked at time when that is safe (i.e., all references - * on the socket have been dropped). + * with a socket. Invoked by sk_common_release before the socket is + * orphaned, so other threads (homa_timer, softirq) may still hold + * references and be reaping its RPCs. * @sk: Socket to destroy. */ void homa_sock_destroy(struct sock *sk) @@ -436,6 +438,18 @@ void homa_sock_destroy(struct sock *sk) #endif /* See strip.py */ } + /* Another thread (e.g. homa_timer) may have taken the last RPCs off + * dead_rpcs and still be freeing them, which uses sk_wmem_alloc and + * the buffer pool. Make any reap that starts from now on return + * without touching either of them, then wait for the ones in + * progress. + */ + homa_sock_lock(hsk); + atomic_inc(&hsk->protect_count); + homa_sock_unlock(hsk); + while (atomic_read_acquire(&hsk->reapers) != 0) + cpu_relax(); + WARN_ON_ONCE(refcount_read(&hsk->sock.sk_wmem_alloc) != 1); #ifdef __UNIT_TEST__ { diff --git a/homa_sock.h b/homa_sock.h index 6c68f77c..fd7ae36c 100644 --- a/homa_sock.h +++ b/homa_sock.h @@ -262,9 +262,18 @@ struct homa_sock { /** * @protect_count: counts the number of calls to homa_protect_rpcs * for which there have not yet been calls to homa_unprotect_rpcs. + * homa_sock_destroy increments it and never decrements it, so no + * reap starts after that. */ atomic_t protect_count; + /** + * @reapers: number of calls to homa_rpc_reap that are collecting or + * freeing dead RPCs for this socket. RPCs leave @dead_rpcs before + * they are freed, so homa_sock_destroy waits for this to reach zero. + */ + atomic_t reapers; + /** * @flags: Additional state information: an OR'ed combination of * various single-bit flags. See below for definitions. Must be diff --git a/test/unit_homa_rpc.c b/test/unit_homa_rpc.c index d201be6e..dbf788eb 100644 --- a/test/unit_homa_rpc.c +++ b/test/unit_homa_rpc.c @@ -17,6 +17,15 @@ #define n(x) htons(x) #define N(x) htonl(x) +static struct homa_sock *hook_hsk; +static int hook_reapers; +static void reapers_hook(char *id) +{ + if (strcmp(id, "spin_lock") == 0 && + atomic_read(&hook_hsk->reapers) > hook_reapers) + hook_reapers = atomic_read(&hook_hsk->reapers); +} + FIXTURE(homa_rpc) { struct in6_addr client_ip[1]; int client_port; @@ -799,6 +808,22 @@ TEST_F(homa_rpc, homa_rpc_reap__protected) EXPECT_EQ(0, homa_rpc_reap(&self->hsk)); homa_unprotect_rpcs(&self->hsk); EXPECT_STREQ("", unit_log_get()); + EXPECT_EQ(0, atomic_read(&self->hsk.reapers)); +} +TEST_F(homa_rpc, homa_rpc_reap__count_reapers) +{ + struct homa_rpc *crpc = unit_client_rpc(&self->hsk, UNIT_OUTGOING, + self->client_ip, self->server_ip, self->server_port, + self->client_id, 5000, 100); + + ASSERT_NE(NULL, crpc); + homa_rpc_end(crpc); + hook_hsk = &self->hsk; + hook_reapers = 0; + unit_hook_register(reapers_hook); + EXPECT_EQ(0, homa_rpc_reap(&self->hsk)); + EXPECT_EQ(1, hook_reapers); + EXPECT_EQ(0, atomic_read(&self->hsk.reapers)); } TEST_F(homa_rpc, homa_rpc_reap__stop_because_of_BATCH_MAX_RPCS) { diff --git a/test/unit_homa_sock.c b/test/unit_homa_sock.c index 8f15df59..e4f38c5a 100644 --- a/test/unit_homa_sock.c +++ b/test/unit_homa_sock.c @@ -41,6 +41,17 @@ static void nospace_hook(char *id) clear_bit(HOMA_SOCK_NOSPACE, &hook_hsk->flags); } +/* Counts the waits for reapers that happen before the buffer pool is freed. */ +static void reapers_hook(char *id) +{ + if (strcmp(id, "cpu_relax") != 0) + return; + if (hook_hsk->buffer_pool) + hook_count++; + if (atomic_read(&hook_hsk->reapers) > 0) + atomic_dec(&hook_hsk->reapers); +} + static void init_scan(struct homa_socktab_scan *scan, struct homa_socktab *socktab) { @@ -501,6 +512,20 @@ TEST_F(homa_sock, homa_sock_shutdown__delete_rpcs) EXPECT_EQ(0, unit_list_length(&self->hsk.active_rpcs)); unit_sock_destroy(&self->hsk); } +TEST_F(homa_sock, homa_sock_destroy__wait_for_reapers) +{ + homa_sock_shutdown(&self->hsk); + atomic_set(&self->hsk.reapers, 2); + hook_hsk = &self->hsk; + hook_count = 0; + unit_hook_register(reapers_hook); + + homa_sock_destroy(&self->hsk.sock); + EXPECT_EQ(2, hook_count); + EXPECT_EQ(0, atomic_read(&self->hsk.reapers)); + EXPECT_EQ(1, atomic_read(&self->hsk.protect_count)); + EXPECT_EQ(NULL, self->hsk.buffer_pool); +} TEST_F(homa_sock, homa_sock_shutdown__wakeup_interests_and_wmem) { struct homa_interest interest1, interest2;