1 files changed, 9 insertions, 4 deletions

diff --git a/drivers/vhost/test.c b/drivers/vhost/test.c
index 9e90e969af55..7804869c6a31 100644
--- a/drivers/vhost/test.c
+++ b/drivers/vhost/test.c
@@ -22,6 +22,12 @@
  * Using this limit prevents one virtqueue from starving others. */
 #define VHOST_TEST_WEIGHT 0x80000
 
+/* Max number of packets transferred before requeueing the job.
+ * Using this limit prevents one virtqueue from starving others with
+ * pkts.
+ */
+#define VHOST_TEST_PKT_WEIGHT 256
+
 enum {
         VHOST_TEST_VQ = 0,
         VHOST_TEST_VQ_MAX = 1,
@@ -80,10 +86,8 @@ static void handle_vq(struct vhost_test *n)
                 }
                 vhost_add_used_and_signal(&n->dev, vq, head, 0);
                 total_len += len;
-                if (unlikely(total_len >= VHOST_TEST_WEIGHT)) {
-                        vhost_poll_queue(&vq->poll);
+                if (unlikely(vhost_exceeds_weight(vq, 0, total_len)))
                         break;
-                }
         }
 
         mutex_unlock(&vq->mutex);
@@ -115,7 +119,8 @@ static int vhost_test_open(struct inode *inode, struct file *f)
         dev = &n->dev;
         vqs[VHOST_TEST_VQ] = &n->vqs[VHOST_TEST_VQ];
         n->vqs[VHOST_TEST_VQ].handle_kick = handle_vq_kick;
-        vhost_dev_init(dev, vqs, VHOST_TEST_VQ_MAX);
+        vhost_dev_init(dev, vqs, VHOST_TEST_VQ_MAX, UIO_MAXIOV,
+                       VHOST_TEST_PKT_WEIGHT, VHOST_TEST_WEIGHT);
 
         f->private_data = n;