Wait for concurrent reapers in homa_sock_destroy

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;
 }
 

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__
 	{

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

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)
 {

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;