From e0fb232671bd49fc47b9e47d760f4b8de3440b34 Mon Sep 17 00:00:00 2001
From: Stuart Ray <88628337+Stuartwastaken@users.noreply.github.com>
Date: Thu, 29 Oct 2026 18:50:04 +0000
Subject: [PATCH] Free tx frag arrays with kvfree

* homa_tx_pool_alloc allocates the frag array with kvmalloc_array, but
  it and homa_rpc_reap freed it with kfree. Since 6.19, kvmalloc falls
  back to vmalloc even for GFP_ATOMIC when kmalloc fails, which can
  happen once the array is bigger than a page (high-order pages not
  available and a message near 1 MB).
* The unit-test kvmalloc mock now marks its blocks so that kfree on
  them fails.
---
 homa_rpc.c               |  2 +-
 homa_tx_pool.c           |  4 ++--
 test/mock.c              | 12 +++++++++++-
 test/unit_homa_tx_pool.c | 10 +++++-----
 4 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/homa_rpc.c b/homa_rpc.c
index 7804577e..2d1c9107 100644
--- a/homa_rpc.c
+++ b/homa_rpc.c
@@ -598,7 +598,7 @@ int homa_rpc_reap(struct homa_sock *hsk)
 		WARN_ON(refcount_sub_and_test(rpc->msgout.frag_bytes,
 					      &hsk->sock.sk_wmem_alloc));
 		if (rpc->msgout.frags != &rpc->msgout.frag)
-			kfree(rpc->msgout.frags);
+			kvfree(rpc->msgout.frags);
 		tt_record2("homa_rpc_reap finished reaping id %d, port %d",
 				rpc->id, rpc->hsk->port);
 #ifndef __STRIP__ /* See strip.py */
diff --git a/homa_tx_pool.c b/homa_tx_pool.c
index c980a5dd..6a7e627a 100644
--- a/homa_tx_pool.c
+++ b/homa_tx_pool.c
@@ -202,7 +202,7 @@ int homa_tx_pool_alloc(struct homa *homa, int length, int *num_frags,
 			memcpy(new_frags, cur_frags,
 			       frags_allocated * sizeof(*new_frags));
 			if (cur_frags != *frags)
-				kfree(cur_frags);
+				kvfree(cur_frags);
 			cur_frags = new_frags;
 			max_frags = num_new;
 		}
@@ -224,7 +224,7 @@ int homa_tx_pool_alloc(struct homa *homa, int length, int *num_frags,
 error:
 	homa_tx_pool_free(homa, tx_core->nid, frags_allocated, cur_frags);
 	if (cur_frags != *frags)
-		kfree(cur_frags);
+		kvfree(cur_frags);
 	*num_frags = 0;
 	local_bh_enable();
 	return err;
diff --git a/test/mock.c b/test/mock.c
index 2756f4ef..8c61771b 100644
--- a/test/mock.c
+++ b/test/mock.c
@@ -1081,6 +1081,9 @@ void kfree(const void *block)
 		FAIL(" %s on unknown block %p", __func__, block);
 		return;
 	}
+	if (strcmp(unit_hash_get(kmallocs_in_use, block), "kvmalloc") == 0)
+		FAIL(" %s on block %p from kvmalloc (needs kvfree)", __func__,
+		     block);
 	unit_hash_erase(kmallocs_in_use, block);
 	free((void *) block);
 }
@@ -1202,6 +1205,8 @@ int kstrtoint(const char *s, unsigned int base, int *res)
 
 void kvfree(const void *addr)
 {
+	if (addr && kmallocs_in_use && unit_hash_get(kmallocs_in_use, addr))
+		unit_hash_set(kmallocs_in_use, addr, "used");
 	kfree(addr);
 }
 
@@ -1212,7 +1217,12 @@ void *__kvmalloc_node_noprof(DECL_BUCKET_PARAMS(size, b), unsigned long align,
 void *__kvmalloc_node_noprof(DECL_BUCKET_PARAMS(size, b), gfp_t flags, int node)
 #endif
 {
-	return mock_kmalloc(size, flags);
+	void *block = mock_kmalloc(size, flags);
+
+	/* Real kvmalloc may return vmalloc memory, which kfree can't free. */
+	if (block)
+		unit_hash_set(kmallocs_in_use, block, "kvmalloc");
+	return block;
 }
 
 struct task_struct *kthread_create_on_node(int (*threadfn)(void *data),
diff --git a/test/unit_homa_tx_pool.c b/test/unit_homa_tx_pool.c
index 89fe566b..3c7b4408 100644
--- a/test/unit_homa_tx_pool.c
+++ b/test/unit_homa_tx_pool.c
@@ -150,7 +150,7 @@ TEST_F(homa_tx_pool, homa_tx_pool_alloc__reallocate_frags_array_large_pages)
 	EXPECT_STREQ("kmalloc 3 units", unit_log_get());
 
 	homa_tx_pool_free(&self->homa, 0, num_frags, frags);
-	kfree(frags);
+	kvfree(frags);
 }
 TEST_F(homa_tx_pool, homa_tx_pool_alloc__reallocate_frags_array_small_pages)
 {
@@ -173,7 +173,7 @@ TEST_F(homa_tx_pool, homa_tx_pool_alloc__reallocate_frags_array_small_pages)
 	EXPECT_STREQ("kmalloc 10 units; kmalloc 160 units", unit_log_get());
 
 	homa_tx_pool_free(&self->homa, 0, num_frags, frags);
-	kfree(frags);
+	kvfree(frags);
 }
 TEST_F(homa_tx_pool, homa_tx_pool_alloc__cant_reallocate_frags_array)
 {
@@ -470,7 +470,7 @@ TEST_F(homa_tx_pool, homa_tx_pool_free__basics)
 	EXPECT_EQ(skb_frag_page(&frags[3]), tx_pool->pages[1]);
 
 	put_page(skb_frag_page(&frags[0]));
-	kfree(frags);
+	kvfree(frags);
 }
 TEST_F(homa_tx_pool, homa_tx_pool_free__pages_not_in_right_numa)
 {
@@ -489,7 +489,7 @@ TEST_F(homa_tx_pool, homa_tx_pool_free__pages_not_in_right_numa)
 	mock_page_nid_mask = 5;
 	homa_tx_pool_free(&self->homa, 0, num_frags, frags);
 	EXPECT_EQ(2, tx_pool->avail);
-	kfree(frags);
+	kvfree(frags);
 }
 TEST_F(homa_tx_pool, homa_tx_pool_free__pool_overflow)
 {
@@ -519,7 +519,7 @@ TEST_F(homa_tx_pool, homa_tx_pool_free__pool_overflow)
 	put_page(skb_frag_page(&frags[0]));
 	put_page(skb_frag_page(&frags[1]));
 	tx_pool->avail = 0;
-	kfree(frags);
+	kvfree(frags);
 }
 
 TEST_F(homa_tx_pool, homa_tx_pool_gc__basics)