[PATCH v2 05/14] hv: vmbus: add KUnit test for order-zero allocation fallback

@@ -60,6 +60,31 @@ bool vmbus_buffer_lower_order(unsigned int *order)
 	return true;
 }
 
+struct page *vmbus_alloc_pages_with_fallback(int nid, gfp_t gfp,
+					     unsigned int *order,
+					     vmbus_alloc_pages_fn alloc,
+					     void *context)
+{
+	struct page *page;
+
+	do {
+		gfp_t try_gfp = gfp;
+
+		/*
+		 * Use __GFP_NORETRY | __GFP_NOWARN to avoid OOM-killing,
+		 * but try harder at order 0 since that is the final
+		 * fallback.
+		 * __GFP_COMP stores order information in the page folio.
+		 */
+		if (*order)
+			try_gfp |= __GFP_COMP | __GFP_NORETRY | __GFP_NOWARN;
+
+		page = alloc(context, nid, try_gfp, *order);
+		if (page || !vmbus_buffer_lower_order(order))
+			return page;
+	} while (true);
+}
+
 bool vmbus_buffer_should_free(const struct vmbus_buffer *buffer)
 {
 	return !buffer->leak && !buffer->gpadl.leak &&

@@ -774,6 +799,14 @@ void vmbus_release_buffer(struct vmbus_buffer *buffer)
 }
 EXPORT_SYMBOL_GPL(vmbus_release_buffer);
 
+static struct page *vmbus_alloc_pages_node(void *context, int nid,
+					   gfp_t gfp,
+					   unsigned int order)
+{
+	(void)context;
+	return alloc_pages_node(nid, gfp, order);
+}
+
 /**
  * __vmbus_alloc_buffer - allocate host-visible, virtually-contiguous backing.
  *

@@ -837,26 +870,17 @@ static void *__vmbus_alloc_buffer(struct vmbus_channel *channel,
 	while (remaining) {
 		struct page *page;
 		gfp_t gfp;
+		int nid;
 
 		order = vmbus_buffer_order(remaining, order);
 
-		/*
-		 * Use __GFP_NORETRY | __GFP_NOWARN to avoid OOM-killing,
-		 * but try harder at order 0 since that is the final
-		 * fallback.
-		 * __GFP_COMP stores order information in the page folio.
-		 */
 		gfp = GFP_KERNEL | __GFP_ZERO;
-		if (order)
-			gfp |= __GFP_COMP | __GFP_NORETRY | __GFP_NOWARN;
-
-		page = alloc_pages_node(cpu_to_node(channel->target_cpu),
-					gfp, order);
-		if (!page) {
-			if (!vmbus_buffer_lower_order(&order))
-				goto err;
-			continue;
-		}
+
+		nid = cpu_to_node(channel->target_cpu);
+		page = vmbus_alloc_pages_with_fallback(nid, gfp, &order,
+						       vmbus_alloc_pages_node, NULL);
+		if (!page)
+			goto err;
 
 		ret = set_memory_decrypted((unsigned long)page_address(page),
 					   1U << order);

@@ -563,6 +563,19 @@ unsigned int vmbus_buffer_order(unsigned long remaining,
 bool vmbus_buffer_lower_order(unsigned int *order);
 bool vmbus_buffer_should_free(const struct vmbus_buffer *buffer);
 
+/*
+ * Order-descent allocation, shared with vmbus_buffer_test.c so the
+ * cases can inject per-order failures. Defined in channel.c, built
+ * into the hv_vmbus object, never exported.
+ */
+typedef struct page *(*vmbus_alloc_pages_fn)(void *context, int nid,
+					     gfp_t gfp,
+					     unsigned int order);
+struct page *vmbus_alloc_pages_with_fallback(int nid, gfp_t gfp,
+					     unsigned int *order,
+					     vmbus_alloc_pages_fn alloc,
+					     void *context);
+
 /*
  * GPADL post and response helpers, also shared with vmbus_buffer_test.c.
  * Same deal as the sizing helpers above: defined in channel.c, built into

@@ -257,12 +257,70 @@ static void vmbus_gpadl_teardown_post_failure_test(struct kunit *test)
 	KUNIT_EXPECT_EQ(test, context.call_count, 1U);
 }
 
+struct vmbus_buffer_page_alloc_test_context {
+	struct kunit *test;
+	unsigned int attempts;
+	unsigned int expected_order;
+	bool fail_order_zero;
+};
+
+static struct page *vmbus_buffer_test_alloc_page(void *context, int nid,
+						 gfp_t gfp,
+						 unsigned int order)
+{
+	struct vmbus_buffer_page_alloc_test_context *test_context = context;
+
+	(void)nid;
+	KUNIT_EXPECT_EQ(test_context->test, order,
+			test_context->expected_order);
+	test_context->attempts++;
+	if (test_context->expected_order)
+		test_context->expected_order--;
+
+	if (order)
+		return NULL;
+	if (test_context->fail_order_zero)
+		return NULL;
+
+	return alloc_pages(gfp, 0);
+}
+
+static void vmbus_buffer_order_zero_allocation_test(struct kunit *test)
+{
+	struct vmbus_buffer_page_alloc_test_context context = {
+		.test = test,
+		.expected_order = MAX_PAGE_ORDER,
+	};
+	struct page *page;
+	unsigned int order = MAX_PAGE_ORDER;
+
+	page = vmbus_alloc_pages_with_fallback(0, GFP_KERNEL, &order,
+					       vmbus_buffer_test_alloc_page, &context);
+	KUNIT_ASSERT_NOT_NULL(test, page);
+	KUNIT_EXPECT_EQ(test, order, 0U);
+	KUNIT_EXPECT_EQ(test, context.expected_order, 0U);
+	KUNIT_EXPECT_EQ(test, context.attempts, (unsigned int)MAX_PAGE_ORDER + 1);
+	__free_pages(page, 0);
+
+	context.attempts = 0;
+	context.expected_order = MAX_PAGE_ORDER;
+	context.fail_order_zero = true;
+	order = MAX_PAGE_ORDER;
+	page = vmbus_alloc_pages_with_fallback(0, GFP_KERNEL, &order,
+					       vmbus_buffer_test_alloc_page, &context);
+	KUNIT_EXPECT_PTR_EQ(test, page, NULL);
+	KUNIT_EXPECT_EQ(test, order, 0U);
+	KUNIT_EXPECT_EQ(test, context.expected_order, 0U);
+	KUNIT_EXPECT_EQ(test, context.attempts, (unsigned int)MAX_PAGE_ORDER + 1);
+}
+
 static struct kunit_case vmbus_buffer_test_cases[] = {
 	KUNIT_CASE(vmbus_buffer_size_rounding_test),
 	KUNIT_CASE(vmbus_buffer_size_overflow_test),
 	KUNIT_CASE(vmbus_ring_fallback_order_zero_test),
 	KUNIT_CASE(vmbus_buffer_failed_teardown_leaks_test),
 	KUNIT_CASE(vmbus_buffer_partial_allocation_cleanup_test),
+	KUNIT_CASE(vmbus_buffer_order_zero_allocation_test),
 	KUNIT_CASE(vmbus_gpadl_post_failure_test),
 	KUNIT_CASE(vmbus_gpadl_post_success_test),
 	KUNIT_CASE(vmbus_gpadl_response_state_test),