From 15b4a8b6210ef0241fb0f62133ad15ebdf48e9bd Mon Sep 17 00:00:00 2001
From: Yongrong Wang <wangyongrong@xiaomi.com>
Date: Tue, 23 Jan 2024 10:45:25 +0800
Subject: [PATCH] openamp: add notify_wait_cb to avoid looping in tx buffer get

Give users a chance to handle the no tx buffer situation when get tx
buffer, with this patch, user can set the notify_wait_cb in driver
and this callback function will be called to handle the wait when no
tx buffer in tx virtqueue.

Signed-off-by: Yongrong Wang <wangyongrong@xiaomi.com>
---
 lib/include/openamp/rpmsg.h        |  1 +
 lib/include/openamp/rpmsg_virtio.h |  9 +++++++++
 lib/rpmsg/rpmsg_virtio.c           | 26 ++++++++++++++++++++++++--
 3 files changed, 34 insertions(+), 2 deletions(-)

diff --git a/lib/include/openamp/rpmsg.h b/lib/include/openamp/rpmsg.h
index 17eaddec2..82371ea77 100644
--- a/lib/include/openamp/rpmsg.h
+++ b/lib/include/openamp/rpmsg.h
@@ -43,6 +43,7 @@ extern "C" {
 #define RPMSG_ERR_INIT			(RPMSG_ERROR_BASE - 6)
 #define RPMSG_ERR_ADDR			(RPMSG_ERROR_BASE - 7)
 #define RPMSG_ERR_PERM			(RPMSG_ERROR_BASE - 8)
+#define RPMSG_EOPNOTSUPP		(RPMSG_ERROR_BASE - 9)
 
 struct rpmsg_endpoint;
 struct rpmsg_device;
diff --git a/lib/include/openamp/rpmsg_virtio.h b/lib/include/openamp/rpmsg_virtio.h
index 58308f70a..752be7762 100644
--- a/lib/include/openamp/rpmsg_virtio.h
+++ b/lib/include/openamp/rpmsg_virtio.h
@@ -41,6 +41,9 @@ extern "C" {
 #define BUFFER_INVALIDATE(x, s)		do { } while (0)
 #endif /* VIRTIO_CACHED_BUFFERS || VIRTIO_USE_DCACHE */
 
+/* Callback handler for rpmsg virtio service */
+typedef int (*rpmsg_virtio_notify_wait_cb)(struct rpmsg_device *rdev, uint32_t id);
+
 /** @brief Shared memory pool used for RPMsg buffers */
 struct rpmsg_virtio_shm_pool {
 	/** Base address of the memory pool */
@@ -98,6 +101,12 @@ struct rpmsg_virtio_device {
 	 * \ref rpmsg_virtio_release_tx_buffer function
 	 */
 	struct metal_list reclaimer;
+
+	/**
+	 * Callback handler for rpmsg virtio service, called when service
+	 * can't get tx buffer
+	 */
+	rpmsg_virtio_notify_wait_cb notify_wait_cb;
 };
 
 #define RPMSG_REMOTE	VIRTIO_DEV_DEVICE
diff --git a/lib/rpmsg/rpmsg_virtio.c b/lib/rpmsg/rpmsg_virtio.c
index 592c89776..1a9d0bfaa 100644
--- a/lib/rpmsg/rpmsg_virtio.c
+++ b/lib/rpmsg/rpmsg_virtio.c
@@ -358,6 +358,18 @@ static void rpmsg_virtio_release_rx_buffer(struct rpmsg_device *rdev,
 	metal_mutex_release(&rdev->lock);
 }
 
+static int rpmsg_virtio_notify_wait(struct rpmsg_virtio_device *rvdev, struct virtqueue *vq)
+{
+	struct virtio_vring_info *vring_info;
+
+	vring_info = &rvdev->vdev->vrings_info[vq->vq_queue_index];
+
+	if (!rvdev->notify_wait_cb)
+		return RPMSG_EOPNOTSUPP;
+
+	return rvdev->notify_wait_cb(&rvdev->rdev, vring_info->notifyid);
+}
+
 static void *rpmsg_virtio_get_tx_payload_buffer(struct rpmsg_device *rdev,
 						uint32_t *len, int wait)
 {
@@ -387,8 +399,18 @@ static void *rpmsg_virtio_get_tx_payload_buffer(struct rpmsg_device *rdev,
 		metal_mutex_release(&rdev->lock);
 		if (rp_hdr || !tick_count)
 			break;
-		metal_sleep_usec(RPMSG_TICKS_PER_INTERVAL);
-		tick_count--;
+
+		/*
+		 * Try to use wait loop implemented in the virtio dispatcher and
+		 * use metal_sleep_usec() method by default.
+		 */
+		status = rpmsg_virtio_notify_wait(rvdev, rvdev->rvq);
+		if (status == RPMSG_EOPNOTSUPP) {
+			metal_sleep_usec(RPMSG_TICKS_PER_INTERVAL);
+			tick_count--;
+		} else if (status == RPMSG_SUCCESS) {
+			break;
+		}
 	}
 
 	if (!rp_hdr)