[PATCH 06/15] io_uring: keep the tctx nodes on a list

[Jens Axboe <axboe@kernel.dk>]

tctx->xa is indexed by the ring pointer, which makes it a deep and
sparse xarray. Iterating it with xa_for_each() is expensive, 12 usec
for a single entry in testing. Keep the nodes on a list as well and
walk that instead, lookups by ring stay in the xarray.

Signed-off-by: Jens Axboe <axboe@kernel.dk>
---
 include/linux/io_uring_types.h |  2 ++
 io_uring/tctx.c                | 10 ++++++----
 io_uring/tctx.h                |  2 ++
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 39629ee77b91..4af3d579ead6 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -153,6 +153,8 @@ struct io_uring_task {
 	struct file			*registered_rings[IO_RINGFD_REG_MAX];
 
 	struct xarray			xa;
+	/* the nodes in ->xa, for walking without the xarray lookup cost */
+	struct list_head		node_list;
 	struct wait_queue_head		wait;
 	atomic_t			in_cancel;
 	atomic_t			inflight_tracked;
diff --git a/io_uring/tctx.c b/io_uring/tctx.c
index 466b7300e208..737dfad4a976 100644
--- a/io_uring/tctx.c
+++ b/io_uring/tctx.c
@@ -105,6 +105,7 @@ __cold struct io_uring_task *io_uring_alloc_task_context(struct task_struct *tas
 
 	tctx->task = task;
 	xa_init(&tctx->xa);
+	INIT_LIST_HEAD(&tctx->node_list);
 	init_waitqueue_head(&tctx->wait);
 	atomic_set(&tctx->in_cancel, 0);
 	atomic_set(&tctx->inflight_tracked, 0);
@@ -135,6 +136,7 @@ static int io_tctx_install_node(struct io_ring_ctx *ctx,
 		kfree(node);
 		return ret;
 	}
+	list_add(&node->tctx_link, &tctx->node_list);
 
 	mutex_lock(&ctx->tctx_lock);
 	list_add(&node->ctx_node, &ctx->tctx_list);
@@ -227,6 +229,7 @@ __cold void io_uring_del_tctx_node(unsigned long index)
 
 	WARN_ON_ONCE(current != node->task);
 	WARN_ON_ONCE(list_empty(&node->ctx_node));
+	list_del(&node->tctx_link);
 
 	mutex_lock(&node->ctx->tctx_lock);
 	list_del(&node->ctx_node);
@@ -243,11 +246,10 @@ __cold void io_uring_del_tctx_node(unsigned long index)
 __cold void io_uring_clean_tctx(struct io_uring_task *tctx)
 {
 	struct io_wq *wq = tctx->io_wq;
-	struct io_tctx_node *node;
-	unsigned long index;
+	struct io_tctx_node *node, *tmp;
 
-	xa_for_each(&tctx->xa, index, node) {
-		io_uring_del_tctx_node(index);
+	list_for_each_entry_safe(node, tmp, &tctx->node_list, tctx_link) {
+		io_uring_del_tctx_node((unsigned long)node->ctx);
 		cond_resched();
 	}
 	if (wq) {
diff --git a/io_uring/tctx.h b/io_uring/tctx.h
index 76ad1ad4594e..9a139816ce4f 100644
--- a/io_uring/tctx.h
+++ b/io_uring/tctx.h
@@ -2,6 +2,8 @@
 
 struct io_tctx_node {
 	struct list_head	ctx_node;
+	/* on tctx->node_list, only ever touched by the owning task */
+	struct list_head	tctx_link;
 	struct task_struct	*task;
 	struct io_ring_ctx	*ctx;
 };
-- 
2.55.0