[PATCH 10/13] virtio_balloon: add higher-order page inflation

[Kevin Loughlin <kevinloughlin@google.com>]

In confidential computing guests requiring page state transitions,
unaccepted memory tracking operates at the EFI unaccepted memory unit
granularity (typically 2MB / order 9 on x86). Inflating the balloon
with 4KB pages would prevent unaccepting the backing 2MB unit.

Inflate balloon pages at unaccepted_memory_order() when
virtballoon_cc_needs_state_transition() is true, while retaining
order-0 allocations for non-CoCo guests and CoCo guests that do not
use unaccepted memory transitions. When a higher-order balloon page
spans more 4KB balloon PFNs than VIRTIO_BALLOON_ARRAY_PFNS_MAX (256),
transmit its PFNs to the host in VIRTIO_BALLOON_ARRAY_PFNS_MAX chunks.

Signed-off-by: Kevin Loughlin <kevinloughlin@google.com>
---
 drivers/virtio/virtio_balloon.c | 87 ++++++++++++++++++++++++++-------
 include/linux/balloon.h         |  1 +
 mm/balloon.c                    | 67 +++++++++++++++++++------
 3 files changed, 121 insertions(+), 34 deletions(-)

diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 11f74848fdc4..d162c162b496 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -182,6 +182,13 @@ static void balloon_ack(struct virtqueue *vq)
 	wake_up(&vb->acked);
 }
 
+static inline bool virtballoon_cc_needs_state_transition(void)
+{
+	return IS_ENABLED(CONFIG_UNACCEPTED_MEMORY) &&
+	       cc_platform_has(CC_ATTR_GUEST_MEM_ENCRYPT) &&
+	       has_unaccepted_memory();
+}
+
 static void tell_host(struct virtio_balloon *vb, struct virtqueue *vq)
 {
 	struct scatterlist sg;
@@ -243,24 +250,49 @@ static void set_page_pfns(struct virtio_balloon *vb,
 					  page_to_balloon_pfn(page) + i);
 }
 
+static void queue_page_pfns(struct virtio_balloon *vb, struct virtqueue *vq,
+			    struct page *page, unsigned int order)
+{
+	unsigned int nr_pages = 1 << order;
+	unsigned int i;
+
+	for (i = 0; i < nr_pages; i++) {
+		set_page_pfns(vb, vb->pfns + vb->num_pfns, page + i);
+		vb->num_pfns += VIRTIO_BALLOON_PAGES_PER_PAGE;
+		if (vb->num_pfns == ARRAY_SIZE(vb->pfns)) {
+			tell_host(vb, vq);
+			vb->num_pfns = 0;
+		}
+	}
+}
+
+static inline unsigned int virtballoon_preferred_order(void)
+{
+	if (virtballoon_cc_needs_state_transition())
+		return unaccepted_memory_order();
+	return 0;
+}
+
 static unsigned int fill_balloon(struct virtio_balloon *vb, size_t num)
 {
-	unsigned int num_allocated_pages;
+	unsigned int order = virtballoon_preferred_order();
+	unsigned int page_pfns = (1 << order) * VIRTIO_BALLOON_PAGES_PER_PAGE;
+	unsigned int num_allocated_pages = 0;
 	struct page *page, *next;
 	unsigned int num_pfns;
 	LIST_HEAD(pages);
 
-	/* We can only do one array worth at a time. */
-	num = min(num, ARRAY_SIZE(vb->pfns));
+	/* We can only do one array or preferred-order page worth at a time. */
+	num = min(num, max_t(size_t, ARRAY_SIZE(vb->pfns), page_pfns));
 
 	for (num_pfns = 0; num_pfns < num;
-	     num_pfns += VIRTIO_BALLOON_PAGES_PER_PAGE) {
-		page = balloon_page_alloc();
+	     num_pfns += page_pfns) {
+		page = balloon_pages_alloc(order);
 
 		if (!page) {
 			dev_info_ratelimited(&vb->vdev->dev,
 					     "Out of puff! Can't get %u pages\n",
-					     VIRTIO_BALLOON_PAGES_PER_PAGE);
+					     page_pfns);
 			/* Sleep for at least 1/5 of a second before retry. */
 			msleep(200);
 			break;
@@ -274,15 +306,17 @@ static unsigned int fill_balloon(struct virtio_balloon *vb, size_t num)
 	vb->num_pfns = 0;
 
 	list_for_each_entry_safe(page, next, &pages, lru) {
+		unsigned int page_order = compound_order(page);
+		unsigned int pfns = (1 << page_order) * VIRTIO_BALLOON_PAGES_PER_PAGE;
+
 		list_del(&page->lru);
 		balloon_page_enqueue(&vb->vb_dev_info, page);
 
-		set_page_pfns(vb, vb->pfns + vb->num_pfns, page);
-		vb->num_pages += VIRTIO_BALLOON_PAGES_PER_PAGE;
-		vb->num_pfns += VIRTIO_BALLOON_PAGES_PER_PAGE;
+		vb->num_pages += pfns;
+		num_allocated_pages += pfns;
+		queue_page_pfns(vb, vb->inflate_vq, page, page_order);
 	}
 
-	num_allocated_pages = vb->num_pfns;
 	/* Did we get any? */
 	if (vb->num_pfns != 0)
 		tell_host(vb, vb->inflate_vq);
@@ -304,28 +338,42 @@ static void release_pages_balloon(struct virtio_balloon *vb,
 
 static unsigned int leak_balloon(struct virtio_balloon *vb, size_t num)
 {
-	unsigned int num_freed_pages;
+	unsigned int order = virtballoon_preferred_order();
+	unsigned int page_pfns = (1 << order) * VIRTIO_BALLOON_PAGES_PER_PAGE;
+	unsigned int num_freed_pages = 0;
 	struct page *page;
 	struct balloon_dev_info *vb_dev_info = &vb->vb_dev_info;
 	LIST_HEAD(pages);
 
-	/* We can only do one array worth at a time. */
-	num = min(num, ARRAY_SIZE(vb->pfns));
+	/* We can only do one array or preferred-order page worth at a time. */
+	num = min(num, max_t(size_t, ARRAY_SIZE(vb->pfns), page_pfns));
 
 	mutex_lock(&vb->balloon_lock);
 	/* We can't release more pages than taken */
 	num = min(num, (size_t)vb->num_pages);
-	for (vb->num_pfns = 0; vb->num_pfns < num;
-	     vb->num_pfns += VIRTIO_BALLOON_PAGES_PER_PAGE) {
+	vb->num_pfns = 0;
+	while (num_freed_pages < num) {
+		unsigned int dequeued_order;
+		unsigned int dequeued_pfns;
+
 		page = balloon_page_dequeue(vb_dev_info);
 		if (!page)
 			break;
-		set_page_pfns(vb, vb->pfns + vb->num_pfns, page);
+		dequeued_order = compound_order(page);
+		dequeued_pfns = (1 << dequeued_order) * VIRTIO_BALLOON_PAGES_PER_PAGE;
+
+		if (num_freed_pages + dequeued_pfns > num &&
+		    num_freed_pages > 0) {
+			balloon_page_enqueue(vb_dev_info, page);
+			break;
+		}
+
+		queue_page_pfns(vb, vb->deflate_vq, page, dequeued_order);
 		list_add(&page->lru, &pages);
-		vb->num_pages -= VIRTIO_BALLOON_PAGES_PER_PAGE;
+		vb->num_pages -= dequeued_pfns;
+		num_freed_pages += dequeued_pfns;
 	}
 
-	num_freed_pages = vb->num_pfns;
 	/*
 	 * Note that if
 	 * virtio_has_feature(vdev, VIRTIO_BALLOON_F_MUST_TELL_HOST);
@@ -474,7 +522,8 @@ static inline s64 towards_target(struct virtio_balloon *vb)
 	 * Aligned up to guest page size to avoid inflating and deflating
 	 * balloon endlessly.
 	 */
-	target = ALIGN(num_pages, VIRTIO_BALLOON_PAGES_PER_PAGE);
+	target = ALIGN((s64)num_pages,
+		       VIRTIO_BALLOON_PAGES_PER_PAGE << virtballoon_preferred_order());
 	return target - vb->num_pages;
 }
 
diff --git a/include/linux/balloon.h b/include/linux/balloon.h
index ca5b15150f42..bc46e36b997f 100644
--- a/include/linux/balloon.h
+++ b/include/linux/balloon.h
@@ -58,6 +58,7 @@ struct balloon_dev_info {
 	bool adjust_managed_page_count;
 };
 
+struct page *balloon_pages_alloc(unsigned int order);
 struct page *balloon_page_alloc(void);
 void balloon_page_enqueue(struct balloon_dev_info *b_dev_info,
 		struct page *page);
diff --git a/mm/balloon.c b/mm/balloon.c
index 96a8f1e20bc6..e9941d08d8e9 100644
--- a/mm/balloon.c
+++ b/mm/balloon.c
@@ -27,9 +27,12 @@ static DEFINE_SPINLOCK(balloon_pages_lock);
 static void balloon_page_insert(struct balloon_dev_info *balloon,
 				       struct page *page)
 {
+	unsigned int i, nr_pages = 1U << compound_order(page);
+
 	lockdep_assert_held(&balloon_pages_lock);
-	__SetPageOffline(page);
-	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION)) {
+	for (i = 0; i < nr_pages; i++)
+		__SetPageOffline(page + i);
+	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION) && !PageCompound(page)) {
 		SetPageMovableOps(page);
 		set_page_private(page, (unsigned long)balloon);
 	}
@@ -45,20 +48,31 @@ static void balloon_page_insert(struct balloon_dev_info *balloon,
  */
 static void balloon_page_finalize(struct page *page)
 {
+	unsigned int i, nr_pages = 1U << compound_order(page);
+
 	lockdep_assert_held(&balloon_pages_lock);
-	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION))
+	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION) && !PageCompound(page))
 		set_page_private(page, 0);
-	/* PageOffline is sticky until the page is freed to the buddy. */
+	/*
+	 * PageOffline is sticky on the head page until freed to the buddy,
+	 * which resets the head page's page_type in __free_pages_prepare().
+	 * Clear PageOffline on tail pages here as __free_pages_prepare()
+	 * expects tail pages to have a default page_type.
+	 */
+	for (i = 1; i < nr_pages; i++)
+		__ClearPageOffline(page + i);
 }
 
 static void balloon_page_enqueue_one(struct balloon_dev_info *b_dev_info,
 				     struct page *page)
 {
+	unsigned int order = compound_order(page);
+
 	balloon_page_insert(b_dev_info, page);
 	if (b_dev_info->adjust_managed_page_count)
-		adjust_managed_page_count(page, -1);
-	__count_vm_event(BALLOON_INFLATE);
-	inc_node_page_state(page, NR_BALLOON_PAGES);
+		adjust_managed_page_count(page, -(1L << order));
+	__count_vm_events(BALLOON_INFLATE, 1UL << order);
+	mod_node_page_state(page_pgdat(page), NR_BALLOON_PAGES, 1L << order);
 }
 
 /**
@@ -117,15 +131,18 @@ size_t balloon_page_list_dequeue(struct balloon_dev_info *b_dev_info,
 
 	spin_lock_irqsave(&balloon_pages_lock, flags);
 	list_for_each_entry_safe(page, tmp, &b_dev_info->pages, lru) {
+		unsigned int order = compound_order(page);
+
 		if (n_pages == n_req_pages)
 			break;
 		list_del(&page->lru);
 		if (b_dev_info->adjust_managed_page_count)
-			adjust_managed_page_count(page, 1);
+			adjust_managed_page_count(page, 1L << order);
 		balloon_page_finalize(page);
-		__count_vm_event(BALLOON_DEFLATE);
+		__count_vm_events(BALLOON_DEFLATE, 1UL << order);
 		list_add(&page->lru, pages);
-		dec_node_page_state(page, NR_BALLOON_PAGES);
+		mod_node_page_state(page_pgdat(page), NR_BALLOON_PAGES,
+				    -(1L << order));
 		n_pages++;
 	}
 	spin_unlock_irqrestore(&balloon_pages_lock, flags);
@@ -135,8 +152,9 @@ size_t balloon_page_list_dequeue(struct balloon_dev_info *b_dev_info,
 EXPORT_SYMBOL_GPL(balloon_page_list_dequeue);
 
 /**
- * balloon_page_alloc - allocates a new page for insertion into the balloon
- *			page list.
+ * balloon_pages_alloc - allocates a new page of given order for insertion
+ *			 into the balloon page list.
+ * @order: allocation order
  *
  * Driver must call this function to properly allocate a new balloon page.
  * Driver must call balloon_page_enqueue before definitively removing the page
@@ -144,16 +162,35 @@ EXPORT_SYMBOL_GPL(balloon_page_list_dequeue);
  *
  * Return: struct page for the allocated page or NULL on allocation failure.
  */
-struct page *balloon_page_alloc(void)
+struct page *balloon_pages_alloc(unsigned int order)
 {
 	gfp_t gfp_flags = __GFP_NOMEMALLOC | __GFP_NORETRY | __GFP_NOWARN;
 
-	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION))
+	if (IS_ENABLED(CONFIG_BALLOON_MIGRATION) && !order)
 		gfp_flags |= GFP_HIGHUSER_MOVABLE;
 	else
 		gfp_flags |= GFP_HIGHUSER;
 
-	return alloc_page(gfp_flags);
+	if (order > 0)
+		gfp_flags |= __GFP_COMP;
+
+	return alloc_pages(gfp_flags, order);
+}
+EXPORT_SYMBOL_GPL(balloon_pages_alloc);
+
+/**
+ * balloon_page_alloc - allocates a new page for insertion into the balloon
+ *			page list.
+ *
+ * Driver must call this function to properly allocate a new balloon page.
+ * Driver must call balloon_page_enqueue before definitively removing the page
+ * from the guest system.
+ *
+ * Return: struct page for the allocated page or NULL on allocation failure.
+ */
+struct page *balloon_page_alloc(void)
+{
+	return balloon_pages_alloc(0);
 }
 EXPORT_SYMBOL_GPL(balloon_page_alloc);
 
-- 
2.56.0.307.g5fe7bd680c-goog