[PATCH RFC] io_uring/rsrc: add last-lookup cache hit to io_rsrc_node_lookup()

[Jens Axboe <axboe@kernel.dk>]

This avoids array_index_nospec() for repeated lookups on the same node,
which can be quite common (and costly). If a cached node is removed from
the given table, it'll get cleared in the cache as well.
io_reset_rsrc_node() takes care of that, which is used in the spots
that's replacing a node.

Note: need to double check this is 100% safe wrt speculation, but I
believe it should be as we're not using the passed in value to index
any arrays (directly).

Signed-off-by: Jens Axboe <axboe@kernel.dk>

---

Sending this out as an RFC, as array_index_nospec() can cause stalls for
frequent lookups. For buffers, it's not unusual to have larger regions
registered, which means hitting the same resource node lookup all the
time.

At the same time, I'm not 100% certain on the sanity of this. Before
you'd always do:

index = array_index_nospec(index, max_nr);
node = some_table[index];

and now you can do:

if (index == last_index)
	return last_node;
last_node = some_table[array_index_nospec(index, max_nr)];
last_index = index;
return last_node;

which _seems_ like it should be safe as no array indexing occurs. Hence
the Jann CC :-)

diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
index 77fd508d043a..c283179b0c89 100644
--- a/include/linux/io_uring_types.h
+++ b/include/linux/io_uring_types.h
@@ -57,6 +57,8 @@ struct io_wq_work {
 
 struct io_rsrc_data {
 	unsigned int			nr;
+	unsigned int			last_index;
+	struct io_rsrc_node		*last_node;
 	struct io_rsrc_node		**nodes;
 };
 
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index 9829c51105ed..413d003bc5d7 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -139,6 +139,8 @@ __cold void io_rsrc_data_free(struct io_rsrc_data *data)
 		if (data->nodes[data->nr])
 			io_put_rsrc_node(data->nodes[data->nr]);
 	}
+	data->last_node = NULL;
+	data->last_index = -1U;
 	kvfree(data->nodes);
 	data->nodes = NULL;
 	data->nr = 0;
@@ -150,6 +152,7 @@ __cold int io_rsrc_data_alloc(struct io_rsrc_data *data, unsigned nr)
 					GFP_KERNEL_ACCOUNT | __GFP_ZERO);
 	if (data->nodes) {
 		data->nr = nr;
+		data->last_index = -1U;
 		return 0;
 	}
 	return -ENOMEM;
diff --git a/io_uring/rsrc.h b/io_uring/rsrc.h
index a40fad783a69..e2795daa877d 100644
--- a/io_uring/rsrc.h
+++ b/io_uring/rsrc.h
@@ -70,8 +70,16 @@ int io_register_rsrc(struct io_ring_ctx *ctx, void __user *arg,
 static inline struct io_rsrc_node *io_rsrc_node_lookup(struct io_rsrc_data *data,
 						       int index)
 {
-	if (index < data->nr)
-		return data->nodes[array_index_nospec(index, data->nr)];
+	if (index < data->nr) {
+		if (index != data->last_index) {
+			index = array_index_nospec(index, data->nr);
+			if (data->nodes[index]) {
+				data->last_index = index;
+				data->last_node = data->nodes[index];
+			}
+		}
+		return data->last_node;
+	}
 	return NULL;
 }
 
@@ -85,8 +93,14 @@ static inline bool io_reset_rsrc_node(struct io_rsrc_data *data, int index)
 {
 	struct io_rsrc_node *node = data->nodes[index];
 
-	if (!node)
+	if (!node) {
+		WARN_ON_ONCE(index == data->last_index);
 		return false;
+	}
+	if (index == data->last_index) {
+		data->last_node = NULL;
+		data->last_index = -1U;
+	}
 	io_put_rsrc_node(node);
 	data->nodes[index] = NULL;
 	return true;

-- 
Jens Axboe