[PATCH 1/3] s390/diag204: Introduce diag204 rework

[Mete Durlu <meted@linux.ibm.com>]

diag204 is an expensive call, limit its invocation to once per
second at most for all users. Use a caching mechanism to provide diag204
data at demand. Cache buffer gets freed by a delayed work function if
diag204 data is not fetched after a certain amount of time.

Signed-off-by: Mete Durlu <meted@linux.ibm.com>
---
 arch/s390/include/asm/diag.h    |   7 +
 arch/s390/kernel/diag/diag204.c | 339 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 346 insertions(+)

diff --git a/arch/s390/include/asm/diag.h b/arch/s390/include/asm/diag.h
index 8db8db3b1018..fe817a7c290d 100644
--- a/arch/s390/include/asm/diag.h
+++ b/arch/s390/include/asm/diag.h
@@ -337,6 +337,13 @@ static inline bool diag204_has_bif(void)
 	return sclp.has_diag204_bif;
 }
 
+int diag204_store(void *buf, int allow_busy, int allow_busy_wait);
+enum diag204_format diag204_get_info_type(void);
+enum diag204_sc diag204_get_subcode(void);
+void *diag204_allocate_buffer(void);
+void diag204_free_buffer(void *buf);
+bool diag204_is_available(void);
+int diag204_get_buf_size(void);
 int diag204(unsigned long subcode, unsigned long size, void *addr);
 int diag224(void *ptr);
 int diag26c(void *req, void *resp, enum diag26c_sc subcode);
diff --git a/arch/s390/kernel/diag/diag204.c b/arch/s390/kernel/diag/diag204.c
new file mode 100644
index 000000000000..8e1909e6caf9
--- /dev/null
+++ b/arch/s390/kernel/diag/diag204.c
@@ -0,0 +1,339 @@
+// SPDX-License-Identifier: GPL-2.0
+/*
+ * Request CPU utilization information from underlying diag 0x204
+ * facility. Diag 0x204 provides a variety of utilization values.
+ *
+ * Copyright IBM Corp. 2025
+ */
+
+#define pr_fmt(fmt) "diag204: " fmt
+
+#include <linux/jiffies.h>
+#include <linux/kernel.h>
+#include <linux/timer.h>
+#include <linux/vmalloc.h>
+
+#include <asm/diag.h>
+
+struct diag204_info {
+	int pages;			/* number of pages for diag204 data */
+	enum diag204_format type;	/* used diag204 data format */
+	enum diag204_sc subcode;	/* used diag204 subcode for store */
+	bool has_bif;
+};
+
+struct diag204_cache {
+	ktime_t timestamp;
+	void *buf;			/* page aligned diag204 buffer */
+	bool updating;
+};
+
+static struct diag204_cache g_cache;
+static struct diag204_info g_info;
+static DEFINE_MUTEX(diag204_mutex);
+
+#define DIAG204_CACHE_EXPIRY (NSEC_PER_SEC)
+#define DIAG204_WORK_DELAY (HZ * 10)
+#define DIAG204_BUSY_WAIT (HZ / 10)
+
+static void diag204_work_fn(struct work_struct *work);
+static DECLARE_DELAYED_WORK(diag204_work, diag204_work_fn);
+
+bool diag204_is_available(void)
+{
+	return g_info.pages > 0;
+}
+
+enum diag204_format diag204_get_info_type(void)
+{
+	return g_info.type;
+}
+
+enum diag204_sc diag204_get_subcode(void)
+{
+	return g_info.subcode;
+}
+
+int diag204_get_buf_size(void)
+{
+	return g_info.pages * PAGE_SIZE;
+}
+
+void diag204_free_buffer(void *buf)
+{
+	if (!buf)
+		return;
+	vfree(buf);
+}
+
+/*
+ * If we use subcode 6 or 7 with extended data format, we can (and should) use
+ * vmalloc, since we need a lot of memory in that case. Currently
+ * up to 423 pages!
+ */
+static void *__diag204_allocate_buffer(int pages)
+{
+	void *buf;
+
+	buf = __vmalloc_node(array_size(pages, PAGE_SIZE), PAGE_SIZE,
+			     GFP_KERNEL | __GFP_ZERO, NUMA_NO_NODE,
+			     __builtin_return_address(0));
+	return buf;
+}
+
+void *diag204_allocate_buffer(void)
+{
+	return __diag204_allocate_buffer(g_info.pages);
+}
+
+static void diag204_schedule_work(void)
+{
+	mod_delayed_work(system_dfl_wq, &diag204_work, DIAG204_WORK_DELAY);
+}
+
+static int diag204_is_cache_valid(unsigned long duration_ns)
+{
+	ktime_t expire;
+
+	lockdep_assert_held(&diag204_mutex);
+	if (!READ_ONCE(g_cache.buf))
+		return 0;
+	expire = ktime_add_ns(g_cache.timestamp, duration_ns);
+	if (ktime_after(ktime_get(), expire))
+		return 0;
+	return 1;
+}
+
+static void diag204_set_cache(void *buf)
+{
+	diag204_free_buffer(READ_ONCE(g_cache.buf));
+	WRITE_ONCE(g_cache.buf, buf);
+	if (!buf)
+		WRITE_ONCE(g_cache.timestamp, 0);
+	else
+		WRITE_ONCE(g_cache.timestamp, ktime_get());
+}
+
+static void diag204_work_fn(struct work_struct *work)
+{
+	mutex_lock(&diag204_mutex);
+	if (!READ_ONCE(g_cache.buf))
+		goto unlock;
+	if (diag204_is_cache_valid(DIAG204_CACHE_EXPIRY))
+		goto rearm;
+	if (READ_ONCE(g_cache.updating))
+		goto rearm;
+	diag204_set_cache(NULL);
+rearm:
+	diag204_schedule_work();
+unlock:
+	mutex_unlock(&diag204_mutex);
+}
+
+static int diag204_cache_new(void *buf, unsigned long subcode)
+{
+	int rc;
+
+	lockdep_assert_held(&diag204_mutex);
+	/*
+	 * Don't allow busy wait or busy looping when g_cache.buf has
+	 * no data yet. We first need a valid data cache before we
+	 * can do updates with busy looping.
+	 */
+	rc = diag204(subcode, g_info.pages, buf);
+	if (rc < 0)
+		return rc;
+	diag204_set_cache(buf);
+	return rc;
+}
+
+static int diag204_cache_replace(void *buf, unsigned long subcode, int allow_busy_wait)
+{
+	int rc;
+
+	lockdep_assert_held(&diag204_mutex);
+	if (!allow_busy_wait) {
+		rc = diag204(subcode, g_info.pages, buf);
+	} else {
+		/*
+		 * Unlock before BUSY loop so updater can sleep, we rely on
+		 * g_cache.updating flag here for synchronising with diag204
+		 * work function and other diag204 callers.
+		 */
+		mutex_unlock(&diag204_mutex);
+		do {
+			rc = diag204(subcode, g_info.pages, buf);
+			if (rc >= 0)
+				break;
+			if (signal_pending(current)) {
+				mutex_lock(&diag204_mutex);
+				return -ERESTARTSYS;
+			}
+			schedule_timeout_interruptible(DIAG204_BUSY_WAIT);
+		} while (rc == -EBUSY);
+		mutex_lock(&diag204_mutex);
+	}
+	if (rc < 0)
+		return rc;
+	diag204_set_cache(buf);
+	return rc;
+}
+
+static int diag204_update_cache(int allow_busy, int allow_busy_wait)
+{
+	unsigned long subcode;
+	void *buf;
+	int rc;
+
+	lockdep_assert_held(&diag204_mutex);
+	buf = diag204_allocate_buffer();
+	if (!buf)
+		return -ENOMEM;
+	subcode = g_info.type | g_info.subcode;
+	if (g_info.has_bif && allow_busy)
+		subcode |= DIAG204_BIF_BIT;
+	WRITE_ONCE(g_cache.updating, true);
+	if (!READ_ONCE(g_cache.buf))
+		rc = diag204_cache_new(buf, subcode);
+	else
+		rc = diag204_cache_replace(buf, subcode, allow_busy_wait);
+	WRITE_ONCE(g_cache.updating, false);
+	if (rc < 0)
+		diag204_free_buffer(buf);
+	return rc;
+}
+
+/*
+ * diag204_store() provides users data retrieved from the diag204 facility.
+ * Since diag204 is an expensive call to make, it is limited to a call per
+ * second for now. To circumvent this limitation a caching mechanism is used.
+ *
+ * diag204_store() first checks if diag204 cache is present and valid. If the
+ * cache is valid, it copies the contents to the users buffer.
+ * If the cache is expired, it first updates the cache and then proceeds with
+ * the same copy operation.
+ *
+ * @buf: Users have to manage their own buffers by using
+ * diag204_allocate_buffer() and diag204_free_buffer().
+ *
+ * @allow_busy: If diag204 supports busy indication facility(bif), users can
+ * specify diag204 to utilize it and return BUSY instead of blocking until
+ * data becomes ready.
+ *
+ * @allow_busy_wait: If users allow diag204 to use bif, they would also get
+ * the option to specify an interruptible wait behavior by allowing busy wait.
+ * This is useful since it gives users to interrupt diag204 waiting period
+ * via signals.
+ */
+int diag204_store(void *buf, int allow_busy, int allow_busy_wait)
+{
+	int rc;
+
+	if (!diag204_is_available())
+		return -EOPNOTSUPP;
+	if (!allow_busy && allow_busy_wait)
+		return -EOPNOTSUPP;
+	if (!buf)
+		return -ENOMEM;
+	rc = 0;
+	if (allow_busy && !allow_busy_wait) {
+		mutex_lock(&diag204_mutex);
+		/*
+		 * Drop the lock if an update is ongoing, user allows to return
+		 * BUSY without waiting
+		 */
+		if (READ_ONCE(g_cache.updating)) {
+			mutex_unlock(&diag204_mutex);
+			diag204_schedule_work();
+			return -EBUSY;
+		}
+	} else {
+		/* Wait for the ongoing update to end */
+		while (1) {
+			mutex_lock(&diag204_mutex);
+			if (!READ_ONCE(g_cache.updating))
+				break;
+			mutex_unlock(&diag204_mutex);
+			if (signal_pending(current))
+				return -ERESTARTSYS;
+			schedule_timeout_interruptible(DIAG204_BUSY_WAIT);
+		}
+	}
+	if (!diag204_is_cache_valid(DIAG204_CACHE_EXPIRY))
+		rc = diag204_update_cache(allow_busy, allow_busy_wait);
+	if (rc >= 0)
+		memcpy(buf, READ_ONCE(g_cache.buf), diag204_get_buf_size());
+	mutex_unlock(&diag204_mutex);
+	diag204_schedule_work();
+	return rc;
+}
+
+/*
+ * diag204_probe() has to find out, which type of diagnose 204 implementation
+ * we have on our machine. Currently there are three possible scanarios:
+ *   - subcode 4   + simple data format (only one page)
+ *   - subcode 4-6 + extended data format
+ *   - subcode 4-7 + extended data format
+ *
+ * Subcode 5 is used to retrieve the size of the data, provided by subcodes
+ * 6 and 7. Subcode 7 basically has the same function as subcode 6. In addition
+ * to subcode 6 it provides also information about secondary cpus.
+ * In order to get as much information as possible, we first try
+ * subcode 7, then 6 and if both fail, we use subcode 4.
+ */
+static int diag204_probe(void)
+{
+	unsigned long subcode, type;
+	int pages, rc;
+	void *buf;
+
+	pages = diag204((unsigned long)DIAG204_INFO_EXT |
+			(unsigned long)DIAG204_SUBC_RSI, 0, NULL);
+	if (pages > 0) {
+		buf = __diag204_allocate_buffer(pages);
+		if (!buf) {
+			rc = -ENOMEM;
+			goto fail_alloc;
+		}
+		type = DIAG204_INFO_EXT;
+		if (diag204(DIAG204_SUBC_STIB7 | type, pages, buf) >= 0) {
+			subcode = DIAG204_SUBC_STIB7;
+			goto out;
+		}
+		if (diag204(DIAG204_SUBC_STIB6 | type, pages, buf) >= 0) {
+			subcode = DIAG204_SUBC_STIB6;
+			goto out;
+		}
+	} else {
+		/* extended data format is not available, fallback to simple */
+		pages = 1;
+		buf = __diag204_allocate_buffer(pages);
+		if (!buf) {
+			rc = -ENOMEM;
+			goto fail_alloc;
+		}
+		type = DIAG204_INFO_SIMPLE;
+		if (diag204(DIAG204_SUBC_STIB4 | type, pages, buf) >= 0) {
+			subcode = DIAG204_SUBC_STIB4;
+			goto out;
+		}
+	}
+	rc = -EOPNOTSUPP;
+	goto fail_store;
+out:
+	rc = 0;
+	g_info.pages = pages;
+	g_info.subcode = subcode;
+	g_info.type = type;
+	g_info.has_bif = diag204_has_bif();
+fail_store:
+	vfree(buf);
+fail_alloc:
+	return rc;
+}
+
+static int __init diag204_init(void)
+{
+	return diag204_probe();
+}
+device_initcall(diag204_init);

-- 
2.56.0