[PATCH 2/2] virtio_balloon: Treat stats requests as wakeup events

From: David Stevens
Date: Mon Mar 18 2024 - 05:11:30 EST


From: David Stevens <stevensd@xxxxxxxxxxxx>

Treat stats requests as wakeup events to ensure that the driver responds
to device requests in a timely manner.

Signed-off-by: David Stevens <stevensd@xxxxxxxxxxxx>
---
drivers/virtio/virtio_balloon.c | 75 ++++++++++++++++++++-------------
1 file changed, 46 insertions(+), 29 deletions(-)

diff --git a/drivers/virtio/virtio_balloon.c b/drivers/virtio/virtio_balloon.c
index 7fe7ef5f1c77..402dec98e08c 100644
--- a/drivers/virtio/virtio_balloon.c
+++ b/drivers/virtio/virtio_balloon.c
@@ -121,11 +121,14 @@ struct virtio_balloon {
struct page_reporting_dev_info pr_dev_info;

/* State for keeping the wakeup_source active while adjusting the balloon */
- spinlock_t adjustment_lock;
- bool adjustment_signal_pending;
- bool adjustment_in_progress;
+ spinlock_t wakeup_lock;
+ bool processing_wakeup_event;
+ u32 wakeup_signal_mask;
};

+#define ADJUSTMENT_WAKEUP_SIGNAL (1 << 0)
+#define STATS_WAKEUP_SIGNAL (1 << 1)
+
static const struct virtio_device_id id_table[] = {
{ VIRTIO_ID_BALLOON, VIRTIO_DEV_ANY_ID },
{ 0 },
@@ -140,6 +143,36 @@ static u32 page_to_balloon_pfn(struct page *page)
return pfn * VIRTIO_BALLOON_PAGES_PER_PAGE;
}

+static void start_wakeup_event(struct virtio_balloon *vb, u32 mask)
+{
+ unsigned long flags;
+
+ spin_lock_irqsave(&vb->wakeup_lock, flags);
+ vb->wakeup_signal_mask |= mask;
+ if (!vb->processing_wakeup_event) {
+ vb->processing_wakeup_event = true;
+ pm_stay_awake(&vb->vdev->dev);
+ }
+ spin_unlock_irqrestore(&vb->wakeup_lock, flags);
+}
+
+static void process_wakeup_event(struct virtio_balloon *vb, u32 mask)
+{
+ spin_lock_irq(&vb->wakeup_lock);
+ vb->wakeup_signal_mask &= ~mask;
+ spin_unlock_irq(&vb->wakeup_lock);
+}
+
+static void finish_wakeup_event(struct virtio_balloon *vb)
+{
+ spin_lock_irq(&vb->wakeup_lock);
+ if (!vb->wakeup_signal_mask && vb->processing_wakeup_event) {
+ vb->processing_wakeup_event = false;
+ pm_relax(&vb->vdev->dev);
+ }
+ spin_unlock_irq(&vb->wakeup_lock);
+}
+
static void balloon_ack(struct virtqueue *vq)
{
struct virtio_balloon *vb = vq->vdev->priv;
@@ -370,8 +403,10 @@ static void stats_request(struct virtqueue *vq)
struct virtio_balloon *vb = vq->vdev->priv;

spin_lock(&vb->stop_update_lock);
- if (!vb->stop_update)
+ if (!vb->stop_update) {
+ start_wakeup_event(vb, STATS_WAKEUP_SIGNAL);
queue_work(system_freezable_wq, &vb->update_balloon_stats_work);
+ }
spin_unlock(&vb->stop_update_lock);
}

@@ -444,29 +479,10 @@ static void virtio_balloon_queue_free_page_work(struct virtio_balloon *vb)

static void start_update_balloon_size(struct virtio_balloon *vb)
{
- unsigned long flags;
-
- spin_lock_irqsave(&vb->adjustment_lock, flags);
- vb->adjustment_signal_pending = true;
- if (!vb->adjustment_in_progress) {
- vb->adjustment_in_progress = true;
- pm_stay_awake(&vb->vdev->dev);
- }
- spin_unlock_irqrestore(&vb->adjustment_lock, flags);
-
+ start_wakeup_event(vb, ADJUSTMENT_WAKEUP_SIGNAL);
queue_work(system_freezable_wq, &vb->update_balloon_size_work);
}

-static void end_update_balloon_size(struct virtio_balloon *vb)
-{
- spin_lock_irq(&vb->adjustment_lock);
- if (!vb->adjustment_signal_pending && vb->adjustment_in_progress) {
- vb->adjustment_in_progress = false;
- pm_relax(&vb->vdev->dev);
- }
- spin_unlock_irq(&vb->adjustment_lock);
-}
-
static void virtballoon_changed(struct virtio_device *vdev)
{
struct virtio_balloon *vb = vdev->priv;
@@ -495,7 +511,10 @@ static void update_balloon_stats_func(struct work_struct *work)

vb = container_of(work, struct virtio_balloon,
update_balloon_stats_work);
+
+ process_wakeup_event(vb, STATS_WAKEUP_SIGNAL);
stats_handle_request(vb);
+ finish_wakeup_event(vb);
}

static void update_balloon_size_func(struct work_struct *work)
@@ -506,9 +525,7 @@ static void update_balloon_size_func(struct work_struct *work)
vb = container_of(work, struct virtio_balloon,
update_balloon_size_work);

- spin_lock_irq(&vb->adjustment_lock);
- vb->adjustment_signal_pending = false;
- spin_unlock_irq(&vb->adjustment_lock);
+ process_wakeup_event(vb, ADJUSTMENT_WAKEUP_SIGNAL);

diff = towards_target(vb);

@@ -523,7 +540,7 @@ static void update_balloon_size_func(struct work_struct *work)
if (diff)
queue_work(system_freezable_wq, work);
else
- end_update_balloon_size(vb);
+ finish_wakeup_event(vb);
}

static int init_vqs(struct virtio_balloon *vb)
@@ -1027,7 +1044,7 @@ static int virtballoon_probe(struct virtio_device *vdev)
goto out_unregister_oom;
}

- spin_lock_init(&vb->adjustment_lock);
+ spin_lock_init(&vb->wakeup_lock);
device_set_wakeup_capable(&vb->vdev->dev, true);

virtio_device_ready(vdev);
--
2.44.0.291.gc1ea87d7ee-goog