From 87bcd3edf3ee7b9d5e76a8090b3f05ea1535026d Mon Sep 17 00:00:00 2001
From: Keir Fraser <keirf@google.com>
Date: Thu, 17 Nov 2022 15:24:40 +0000
Subject: [PATCH] ANDROID: virtio_balloon: New module parameter "pkvm"

This specifies that the driver is running on a PKVM hypervisor
and must use the memrelinquish service to cooperatively release
memory. If this service is unavailable, virtio_balloon cannot be
used.

Bug: 240239989
Change-Id: I8800c4435d8fae9df6f1ab108cc61c8f93020773
Signed-off-by: Keir Fraser <keirf@google.com>
---
 drivers/virtio/virtio_balloon.c | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 8e810041ee3f..7514c8cf7b5d 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -21,6 +21,11 @@
 #include <linux/magic.h>
 #include <linux/pseudo_fs.h>
 #include <linux/page_reporting.h>
+#include <linux/mem_relinquish.h>
+
+static bool pkvm;
+module_param(pkvm, bool, 0);
+MODULE_PARM_DESC(pkvm, "Running on PKVM. Must use MEM_RELINQUISH.");
 
 /*
  * Balloon device works in 4K page units.  So each page is pointed to by
@@ -887,6 +892,12 @@ static int virtballoon_probe(struct virtio_device *vdev)
 	struct virtio_balloon *vb;
 	int err;
 
+	if (pkvm && !kvm_has_memrelinquish_services()) {
+		dev_err(&vdev->dev, "%s failure: pkvm but no memrelinquish\n",
+			__func__);
+		return -EINVAL;
+	}
+
 	if (!vdev->config->get) {
 		dev_err(&vdev->dev, "%s failure: config access disabled\n",
 			__func__);