io_uring: allow re-poll if we made progress

commit 10c873334febaeea9aa0c25c10b5ac0951b77a5f upstream.

We currently check REQ_F_POLLED before arming async poll for a
notification to retry. If it's set, then we don't allow poll and will
punt to io-wq instead. This is done to prevent a situation where a buggy
driver will repeatedly return that there's space/data available yet we
get -EAGAIN.

However, if we already transferred data, then it should be safe to rely
on poll again. Gate the check on whether or not REQ_F_PARTIAL_IO is
also set.

Signed-off-by: Jens Axboe <axboe@kernel.dk>
Signed-off-by: Sasha Levin <sashal@kernel.org>

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 3fb76863fed4d8fba97dd7175e465448b51bc3a2..997a7264e1d4b72bfd9133b49add9ef7f163e205 100644 (file)
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -5853,7 +5853,7 @@ static int io_arm_poll_handler(struct io_kiocb *req)
 
        if (!req->file || !file_can_poll(req->file))
                return IO_APOLL_ABORTED;
-       if (req->flags & REQ_F_POLLED)
+       if ((req->flags & (REQ_F_POLLED|REQ_F_PARTIAL_IO)) == REQ_F_POLLED)
                return IO_APOLL_ABORTED;
        if (!def->pollin && !def->pollout)
                return IO_APOLL_ABORTED;
@@ -5869,7 +5869,10 @@ static int io_arm_poll_handler(struct io_kiocb *req)
                mask |= POLLOUT | POLLWRNORM;
        }
 
-       apoll = kmalloc(sizeof(*apoll), GFP_ATOMIC);
+       if (req->flags & REQ_F_POLLED)
+               apoll = req->apoll;
+       else
+               apoll = kmalloc(sizeof(*apoll), GFP_ATOMIC);
        if (unlikely(!apoll))
                return IO_APOLL_ABORTED;
        apoll->double_poll = NULL;