Re: [PATCH for-next v2 4/6] io_uring: add IORING_SETUP_DEFER_TASKRUN

[Pavel Begunkov <asml.silence@gmail.com>]

On 8/16/22 16:37, Dylan Yudaken wrote:
> Allow deferring async tasks until the user calls io_uring_enter(2) with
> the IORING_ENTER_GETEVENTS flag. Enable this mode with a flag at
> io_uring_setup time. This functionality requires that the later
> io_uring_enter will be called from the same submission task, and therefore
> restrict this flag to work only when IORING_SETUP_SINGLE_ISSUER is also
> set.
> 
> Being able to hand pick when tasks are run prevents the problem where
> there is current work to be done, however task work runs anyway.
> 
> For example, a common workload would obtain a batch of CQEs, and process
> each one. Interrupting this to additional taskwork would add latency but
> not gain anything. If instead task work is deferred to just before more
> CQEs are obtained then no additional latency is added.
> 
> The way this is implemented is by trying to keep task work local to a
> io_ring_ctx, rather than to the submission task. This is required, as the
> application will want to wake up only a single io_ring_ctx at a time to
> process work, and so the lists of work have to be kept separate.
> 
> This has some other benefits like not having to check the task continually
> in handle_tw_list (and potentially unlocking/locking those), and reducing
> locks in the submit & process completions path.
> 
> There are networking cases where using this option can reduce request
> latency by 50%. For example a contrived example using [1] where the client
> sends 2k data and receives the same data back while doing some system
> calls (to trigger task work) shows this reduction. The reason ends up
> being that if sending responses is delayed by processing task work, then
> the client side sits idle. Whereas reordering the sends first means that
> the client runs it's workload in parallel with the local task work.
> 
> [1]:
> Using https://github.com/DylanZA/netbench/tree/defer_run
> Client:
> ./netbench  --client_only 1 --control_port 10000 --host <host> --tx "epoll --threads 16 --per_thread 1 --size 2048 --resp 2048 --workload 1000"
> Server:
> ./netbench  --server_only 1 --control_port 10000  --rx "io_uring --defer_taskrun 0 --workload 100"   --rx "io_uring  --defer_taskrun 1 --workload 100"
> 
> Signed-off-by: Dylan Yudaken <dylany@fb.com>
> ---
>   include/linux/io_uring_types.h |   2 +
>   include/uapi/linux/io_uring.h  |   7 ++
>   io_uring/cancel.c              |   2 +-
>   io_uring/io_uring.c            | 119 +++++++++++++++++++++++++++++----
>   io_uring/io_uring.h            |  30 ++++++++-
>   io_uring/rsrc.c                |   2 +-
>   6 files changed, 146 insertions(+), 16 deletions(-)
> 
> diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
> index 677a25d44d7f..d56ff2185168 100644
> --- a/include/linux/io_uring_types.h
> +++ b/include/linux/io_uring_types.h
> @@ -301,6 +301,8 @@ struct io_ring_ctx {
>   		struct io_hash_table	cancel_table;
>   		bool			poll_multi_queue;
>   
> +		struct llist_head	work_llist;
> +
>   		struct list_head	io_buffers_comp;
>   	} ____cacheline_aligned_in_smp;
>   
> diff --git a/include/uapi/linux/io_uring.h b/include/uapi/linux/io_uring.h
> index 1463cfecb56b..be8d1801bf4a 100644
> --- a/include/uapi/linux/io_uring.h
> +++ b/include/uapi/linux/io_uring.h
> @@ -153,6 +153,13 @@ enum {
>    */
>   #define IORING_SETUP_SINGLE_ISSUER	(1U << 12)
>   
> +/*
> + * Defer running task work to get events.
> + * Rather than running bits of task work whenever the task transitions
> + * try to do it just before it is needed.
> + */
> +#define IORING_SETUP_DEFER_TASKRUN	(1U << 13)
> +
>   enum io_uring_op {
>   	IORING_OP_NOP,
>   	IORING_OP_READV,
> diff --git a/io_uring/cancel.c b/io_uring/cancel.c
> index e4e1dc0325f0..db6180b62e41 100644
> --- a/io_uring/cancel.c
> +++ b/io_uring/cancel.c
> @@ -292,7 +292,7 @@ int io_sync_cancel(struct io_ring_ctx *ctx, void __user *arg)
>   			break;
>   
>   		mutex_unlock(&ctx->uring_lock);
> -		ret = io_run_task_work_sig();
> +		ret = io_run_task_work_sig(ctx);
>   		if (ret < 0) {
>   			mutex_lock(&ctx->uring_lock);
>   			break;
> diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
> index 53696dd90626..6b16da0be712 100644
> --- a/io_uring/io_uring.c
> +++ b/io_uring/io_uring.c
> @@ -316,6 +316,7 @@ static __cold struct io_ring_ctx *io_ring_ctx_alloc(struct io_uring_params *p)
>   	INIT_LIST_HEAD(&ctx->rsrc_ref_list);
>   	INIT_DELAYED_WORK(&ctx->rsrc_put_work, io_rsrc_put_work);
>   	init_llist_head(&ctx->rsrc_put_llist);
> +	init_llist_head(&ctx->work_llist);
>   	INIT_LIST_HEAD(&ctx->tctx_list);
>   	ctx->submit_state.free_list.next = NULL;
>   	INIT_WQ_LIST(&ctx->locked_free_list);
> @@ -1047,12 +1048,30 @@ void tctx_task_work(struct callback_head *cb)
>   	trace_io_uring_task_work_run(tctx, count, loops);
>   }
>   
> -void io_req_task_work_add(struct io_kiocb *req)
> +static void io_req_local_work_add(struct io_kiocb *req)
> +{
> +	struct io_ring_ctx *ctx = req->ctx;
> +
> +	if (!llist_add(&req->io_task_work.node, &ctx->work_llist))
> +		return;
> +
> +	if (ctx->flags & IORING_SETUP_TASKRUN_FLAG)
> +		atomic_or(IORING_SQ_TASKRUN, &ctx->rings->sq_flags);
> +
> +	io_cqring_wake(ctx);
> +}
> +
> +static inline void __io_req_task_work_add(struct io_kiocb *req, bool allow_local)
>   {
>   	struct io_uring_task *tctx = req->task->io_uring;
>   	struct io_ring_ctx *ctx = req->ctx;
>   	struct llist_node *node;
>   
> +	if (allow_local && ctx->flags & IORING_SETUP_DEFER_TASKRUN) {
> +		io_req_local_work_add(req);
> +		return;
> +	}
> +
>   	/* task_work already pending, we're done */
>   	if (!llist_add(&req->io_task_work.node, &tctx->task_list))
>   		return;
> @@ -1074,6 +1093,55 @@ void io_req_task_work_add(struct io_kiocb *req)
>   	}
>   }
>   
> +void io_req_task_work_add(struct io_kiocb *req)
> +{
> +	__io_req_task_work_add(req, true);
> +}
> +
> +bool io_run_local_work(struct io_ring_ctx *ctx, bool locked)
> +{
> +	struct llist_node *node;
> +	struct llist_node fake;
> +	struct llist_node *current_final = NULL;
> +	unsigned int count;
> +
> +	if (!locked)
> +		locked = mutex_trylock(&ctx->uring_lock);
> +
> +	node = io_llist_xchg(&ctx->work_llist, &fake);
> +	count = 0;
> +again:
> +	while (node != current_final) {
> +		struct llist_node *next = node->next;
> +		struct io_kiocb *req = container_of(node, struct io_kiocb,
> +						    io_task_work.node);
> +		prefetch(container_of(next, struct io_kiocb, io_task_work.node));
> +		if (unlikely(req->task != current)) {
> +			__io_req_task_work_add(req, false);
> +		} else {
> +			req->io_task_work.func(req, &locked);
> +			count++;
> +		}
> +		node = next;
> +	}
> +
> +	if (ctx->flags & IORING_SETUP_TASKRUN_FLAG)
> +		atomic_andnot(IORING_SQ_TASKRUN, &ctx->rings->sq_flags);
> +
> +	node = io_llist_cmpxchg(&ctx->work_llist, &fake, NULL);
> +	if (node != &fake) {
> +		current_final = &fake;
> +		node = io_llist_xchg(&ctx->work_llist, &fake);
> +		goto again;
> +	}
> +
> +	if (locked) {
> +		io_submit_flush_completions(ctx);
> +		mutex_unlock(&ctx->uring_lock);
> +	}
> +	return count > 0;
> +}
> +
>   static void io_req_tw_post(struct io_kiocb *req, bool *locked)
>   {
>   	io_req_complete_post(req);
> @@ -1284,8 +1352,7 @@ static int io_iopoll_check(struct io_ring_ctx *ctx, long min)
>   		if (wq_list_empty(&ctx->iopoll_list)) {
>   			u32 tail = ctx->cached_cq_tail;
>   
> -			mutex_unlock(&ctx->uring_lock);
> -			io_run_task_work();
> +			io_run_task_work_unlock_ctx(ctx);
>   			mutex_lock(&ctx->uring_lock);
>   
>   			/* some requests don't go through iopoll_list */
> @@ -2146,7 +2213,9 @@ struct io_wait_queue {
>   
>   static inline bool io_has_work(struct io_ring_ctx *ctx)
>   {
> -	return test_bit(IO_CHECK_CQ_OVERFLOW_BIT, &ctx->check_cq);
> +	return test_bit(IO_CHECK_CQ_OVERFLOW_BIT, &ctx->check_cq) ||
> +	       ((ctx->flags & IORING_SETUP_DEFER_TASKRUN) &&
> +		!llist_empty(&ctx->work_llist));
>   }
>   
>   static inline bool io_should_wake(struct io_wait_queue *iowq)
> @@ -2178,9 +2247,9 @@ static int io_wake_function(struct wait_queue_entry *curr, unsigned int mode,
>   	return -1;
>   }
>   
> -int io_run_task_work_sig(void)
> +int io_run_task_work_sig(struct io_ring_ctx *ctx)
>   {
> -	if (io_run_task_work())
> +	if (io_run_task_work_ctx(ctx, true))
>   		return 1;
>   	if (task_sigpending(current))
>   		return -EINTR;
> @@ -2196,7 +2265,7 @@ static inline int io_cqring_wait_schedule(struct io_ring_ctx *ctx,
>   	unsigned long check_cq;
>   
>   	/* make sure we run task_work before checking for signals */
> -	ret = io_run_task_work_sig();
> +	ret = io_run_task_work_sig(ctx);
>   	if (ret || io_should_wake(iowq))
>   		return ret;
>   
> @@ -2230,7 +2299,7 @@ static int io_cqring_wait(struct io_ring_ctx *ctx, int min_events,
>   		io_cqring_overflow_flush(ctx);
>   		if (io_cqring_events(ctx) >= min_events)
>   			return 0;
> -		if (!io_run_task_work())
> +		if (!io_run_task_work_ctx(ctx, false))

Sounds like there should be true, but see comments
around io_run_task_work_ctx().

Also, should we do

if (DEFER_TW) {
	if (current != ctx->submitter_task) // + SQPOLL handling
		return -E<error>;
}

And on top we can remove moving local tw -> normal tw in
io_run_local_work(), which probably will only be needed in
exit_work().


>   			break;
>   	} while (1);
>   
> @@ -2768,6 +2837,8 @@ static __cold bool io_uring_try_cancel_requests(struct io_ring_ctx *ctx,
>   		}
>   	}
>   
> +	if (ctx->flags & IORING_SETUP_DEFER_TASKRUN)
> +		ret |= io_run_local_work(ctx, false);
>   	ret |= io_cancel_defer_files(ctx, task, cancel_all);
>   	mutex_lock(&ctx->uring_lock);
>   	ret |= io_poll_remove_all(ctx, task, cancel_all);
> @@ -2837,7 +2908,7 @@ __cold void io_uring_cancel_generic(bool cancel_all, struct io_sq_data *sqd)
>   		}
>   
>   		prepare_to_wait(&tctx->wait, &wait, TASK_INTERRUPTIBLE);
> -		io_run_task_work();
> +		io_run_task_work_ctx(ctx, true);

Here we rely on normal tw waking up the task, which is not done
by the new local tw

>   		io_uring_drop_tctx_refs(current);
>   
>   		/*
> @@ -3057,10 +3128,17 @@ SYSCALL_DEFINE6(io_uring_enter, unsigned int, fd, u32, to_submit,
>   		}
>   		if ((flags & IORING_ENTER_GETEVENTS) && ctx->syscall_iopoll)
>   			goto iopoll_locked;
> +		if ((flags & IORING_ENTER_GETEVENTS) &&
> +			(ctx->flags & IORING_SETUP_DEFER_TASKRUN)) {
> +			io_run_local_work(ctx, true);
> +			goto getevents_ran_local;
> +		}
>   		mutex_unlock(&ctx->uring_lock);
>   	}
> +
>   	if (flags & IORING_ENTER_GETEVENTS) {
>   		int ret2;
> +
>   		if (ctx->syscall_iopoll) {
>   			/*
>   			 * We disallow the app entering submit/complete with
> @@ -3081,6 +3159,9 @@ SYSCALL_DEFINE6(io_uring_enter, unsigned int, fd, u32, to_submit,
>   			const sigset_t __user *sig;
>   			struct __kernel_timespec __user *ts;
>   
> +			if (ctx->flags & IORING_SETUP_DEFER_TASKRUN)
> +				io_run_local_work(ctx, false);
> +getevents_ran_local:
>   			ret2 = io_get_ext_arg(flags, argp, &argsz, &ts, &sig);
>   			if (likely(!ret2)) {
>   				min_complete = min(min_complete,
> @@ -3289,17 +3370,29 @@ static __cold int io_uring_create(unsigned entries, struct io_uring_params *p,
>   	if (ctx->flags & IORING_SETUP_SQPOLL) {
>   		/* IPI related flags don't make sense with SQPOLL */
>   		if (ctx->flags & (IORING_SETUP_COOP_TASKRUN |
> -				  IORING_SETUP_TASKRUN_FLAG))
> +				  IORING_SETUP_TASKRUN_FLAG |
> +				  IORING_SETUP_DEFER_TASKRUN))
>   			goto err;
>   		ctx->notify_method = TWA_SIGNAL_NO_IPI;
>   	} else if (ctx->flags & IORING_SETUP_COOP_TASKRUN) {
>   		ctx->notify_method = TWA_SIGNAL_NO_IPI;
>   	} else {
> -		if (ctx->flags & IORING_SETUP_TASKRUN_FLAG)
> +		if (ctx->flags & IORING_SETUP_TASKRUN_FLAG &&
> +		    !(ctx->flags & IORING_SETUP_DEFER_TASKRUN))
>   			goto err;
>   		ctx->notify_method = TWA_SIGNAL;
>   	}
>   
> +	/*
> +	 * For DEFER_TASKRUN we require the completion task to be the same as the
> +	 * submission task. This implies that there is only one submitter, so enforce
> +	 * that.
> +	 */
> +	if (ctx->flags & IORING_SETUP_DEFER_TASKRUN &&
> +	    !(ctx->flags & IORING_SETUP_SINGLE_ISSUER)) {
> +		goto err;
> +	}
> +
>   	/*
>   	 * This is just grabbed for accounting purposes. When a process exits,
>   	 * the mm is exited and dropped before the files, hence we need to hang
> @@ -3400,7 +3493,7 @@ static long io_uring_setup(u32 entries, struct io_uring_params __user *params)
>   			IORING_SETUP_R_DISABLED | IORING_SETUP_SUBMIT_ALL |
>   			IORING_SETUP_COOP_TASKRUN | IORING_SETUP_TASKRUN_FLAG |
>   			IORING_SETUP_SQE128 | IORING_SETUP_CQE32 |
> -			IORING_SETUP_SINGLE_ISSUER))
> +			IORING_SETUP_SINGLE_ISSUER | IORING_SETUP_DEFER_TASKRUN))
>   		return -EINVAL;
>   
>   	return io_uring_create(entries, &p, params);
> @@ -3872,7 +3965,7 @@ SYSCALL_DEFINE4(io_uring_register, unsigned int, fd, unsigned int, opcode,
>   
>   	ctx = f.file->private_data;
>   
> -	io_run_task_work();
> +	io_run_task_work_ctx(ctx, true);
>   
>   	mutex_lock(&ctx->uring_lock);
>   	ret = __io_uring_register(ctx, opcode, arg, nr_args);
> diff --git a/io_uring/io_uring.h b/io_uring/io_uring.h
> index 2f73f83af960..4d0e9c2e00d0 100644
> --- a/io_uring/io_uring.h
> +++ b/io_uring/io_uring.h
> @@ -26,7 +26,8 @@ enum {
>   
>   struct io_uring_cqe *__io_get_cqe(struct io_ring_ctx *ctx);
>   bool io_req_cqe_overflow(struct io_kiocb *req);
> -int io_run_task_work_sig(void);
> +int io_run_task_work_sig(struct io_ring_ctx *ctx);
> +bool io_run_local_work(struct io_ring_ctx *ctx, bool locked);
>   void io_req_complete_failed(struct io_kiocb *req, s32 res);
>   void __io_req_complete(struct io_kiocb *req, unsigned issue_flags);
>   void io_req_complete_post(struct io_kiocb *req);
> @@ -234,6 +235,33 @@ static inline bool io_run_task_work(void)
>   	return false;
>   }
>   
> +static inline bool io_run_task_work_ctx(struct io_ring_ctx *ctx, bool all)
> +{
> +	bool ret = false;

I'd rather prefer to have two inlined back to back calls like

io_run_local_work();
io_run_task_work();

instead of passing true/false, which are always confusing.


> +
> +	if (ctx->flags & IORING_SETUP_DEFER_TASKRUN) {
> +		ret = io_run_local_work(ctx, false);
> +		if (!all)
> +			return ret;
> +	}
> +	ret  |= io_run_task_work();
> +	return ret;
> +}
> +
> +static inline bool io_run_task_work_unlock_ctx(struct io_ring_ctx *ctx)
> +{
> +	bool ret;
> +
> +	if (ctx->flags & IORING_SETUP_DEFER_TASKRUN) {
> +		ret = io_run_local_work(ctx, true);
> +	} else {
> +		mutex_unlock(&ctx->uring_lock);
> +		ret = io_run_task_work();
> +	}
> +
> +	return ret;
> +}
> +
>   static inline void io_tw_lock(struct io_ring_ctx *ctx, bool *locked)
>   {
>   	if (!*locked) {
> diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
> index 71359a4d0bd4..80cda6e2067f 100644
> --- a/io_uring/rsrc.c
> +++ b/io_uring/rsrc.c
> @@ -343,7 +343,7 @@ __cold static int io_rsrc_ref_quiesce(struct io_rsrc_data *data,
>   		flush_delayed_work(&ctx->rsrc_put_work);
>   		reinit_completion(&data->done);
>   
> -		ret = io_run_task_work_sig();
> +		ret = io_run_task_work_sig(ctx);
>   		mutex_lock(&ctx->uring_lock);
>   	} while (ret >= 0);
>   	data->quiesce = false;

-- 
Pavel Begunkov