io_uring: make ASYNC_CANCEL work with poll and timeout

index a2548a6dd195..6db27bc9b2e3 100644
@@ -1957,6 +1957,20 @@ static void io_poll_remove_all(struct io_ring_ctx *ctx)
 	spin_unlock_irq(&ctx->completion_lock);
 }
 
+static int io_poll_cancel(struct io_ring_ctx *ctx, __u64 sqe_addr)
+{
+	struct io_kiocb *req;
+
+	list_for_each_entry(req, &ctx->cancel_list, list) {
+		if (req->user_data != sqe_addr)
+			continue;
+		io_poll_remove_one(req);
+		return 0;
+	}
+
+	return -ENOENT;
+}
+
 /*
  * Find a running poll command that matches one specified in sqe->addr,
  * and remove it if found.

@@ -1964,8 +1978,7 @@ static void io_poll_remove_all(struct io_ring_ctx *ctx)
 static int io_poll_remove(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 {
 	struct io_ring_ctx *ctx = req->ctx;
-	struct io_kiocb *poll_req, *next;
-	int ret = -ENOENT;
+	int ret;
 
 	if (unlikely(req->ctx->flags & IORING_SETUP_IOPOLL))
 		return -EINVAL;

@@ -1974,13 +1987,7 @@ static int io_poll_remove(struct io_kiocb *req, const struct io_uring_sqe *sqe)
 		return -EINVAL;
 
 	spin_lock_irq(&ctx->completion_lock);
-	list_for_each_entry_safe(poll_req, next, &ctx->cancel_list, list) {
-		if (READ_ONCE(sqe->addr) == poll_req->user_data) {
-			io_poll_remove_one(poll_req);
-			ret = 0;
-			break;
-		}
-	}
+	ret = io_poll_cancel(ctx, READ_ONCE(sqe->addr));
 	spin_unlock_irq(&ctx->completion_lock);
 
 	io_cqring_add_event(req, ret);

@@ -2200,6 +2207,31 @@ static enum hrtimer_restart io_timeout_fn(struct hrtimer *timer)
 	return HRTIMER_NORESTART;
 }
 
+static int io_timeout_cancel(struct io_ring_ctx *ctx, __u64 user_data)
+{
+	struct io_kiocb *req;
+	int ret = -ENOENT;
+
+	list_for_each_entry(req, &ctx->timeout_list, list) {
+		if (user_data == req->user_data) {
+			list_del_init(&req->list);
+			ret = 0;
+			break;
+		}
+	}
+
+	if (ret == -ENOENT)
+		return ret;
+
+	ret = hrtimer_try_to_cancel(&req->timeout.timer);
+	if (ret == -1)
+		return -EALREADY;
+
+	io_cqring_fill_event(req, -ECANCELED);
+	io_put_req(req);
+	return 0;
+}
+
 /*
  * Remove or update an existing timeout command
  */

@@ -2207,10 +2239,8 @@ static int io_timeout_remove(struct io_kiocb *req,
 			     const struct io_uring_sqe *sqe)
 {
 	struct io_ring_ctx *ctx = req->ctx;
-	struct io_kiocb *treq;
-	int ret = -ENOENT;
-	__u64 user_data;
 	unsigned flags;
+	int ret;
 
 	if (unlikely(ctx->flags & IORING_SETUP_IOPOLL))
 		return -EINVAL;

@@ -2220,42 +2250,15 @@ static int io_timeout_remove(struct io_kiocb *req,
 	if (flags)
 		return -EINVAL;
 
-	user_data = READ_ONCE(sqe->addr);
 	spin_lock_irq(&ctx->completion_lock);
-	list_for_each_entry(treq, &ctx->timeout_list, list) {
-		if (user_data == treq->user_data) {
-			list_del_init(&treq->list);
-			ret = 0;
-			break;
-		}
-	}
-
-	/* didn't find timeout */
-	if (ret) {
-fill_ev:
-		io_cqring_fill_event(req, ret);
-		io_commit_cqring(ctx);
-		spin_unlock_irq(&ctx->completion_lock);
-		io_cqring_ev_posted(ctx);
-		if (req->flags & REQ_F_LINK)
-			req->flags |= REQ_F_FAIL_LINK;
-		io_put_req(req);
-		return 0;
-	}
+	ret = io_timeout_cancel(ctx, READ_ONCE(sqe->addr));
 
-	ret = hrtimer_try_to_cancel(&treq->timeout.timer);
-	if (ret == -1) {
-		ret = -EBUSY;
-		goto fill_ev;
-	}
-
-	io_cqring_fill_event(req, 0);
-	io_cqring_fill_event(treq, -ECANCELED);
+	io_cqring_fill_event(req, ret);
 	io_commit_cqring(ctx);
 	spin_unlock_irq(&ctx->completion_lock);
 	io_cqring_ev_posted(ctx);
-
-	io_put_req(treq);
+	if (ret < 0 && req->flags & REQ_F_LINK)
+		req->flags |= REQ_F_FAIL_LINK;
 	io_put_req(req);
 	return 0;
 }

@@ -2387,10 +2390,24 @@ static int io_async_cancel(struct io_kiocb *req, const struct io_uring_sqe *sqe,
 
 	sqe_addr = (void *) (unsigned long) READ_ONCE(sqe->addr);
 	ret = io_async_cancel_one(ctx, sqe_addr);
+	if (ret != -ENOENT) {
+		spin_lock_irq(&ctx->completion_lock);
+		goto done;
+	}
+
+	spin_lock_irq(&ctx->completion_lock);
+	ret = io_timeout_cancel(ctx, (unsigned long) sqe_addr);
+	if (ret != -ENOENT)
+		goto done;
+	ret = io_poll_cancel(ctx, (unsigned long) sqe_addr);
+done:
+	io_cqring_fill_event(req, ret);
+	io_commit_cqring(ctx);
+	spin_unlock_irq(&ctx->completion_lock);
+	io_cqring_ev_posted(ctx);
 
 	if (ret < 0 && (req->flags & REQ_F_LINK))
 		req->flags |= REQ_F_FAIL_LINK;
-	io_cqring_add_event(req, ret);
 	io_put_req_find_next(req, nxt);
 	return 0;
 }