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

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

[Jens Axboe]

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

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

[Jann Horn]

On Wed, Oct 30, 2024 at 5:58 PM Jens Axboe <axboe@kernel.dk> wrote:
> 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

You're saying array_index_nospec() can be quite costly - which
architecture is this on? Is this the cost of the compare+subtract+and
making the critical path longer?

> 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 :-)

I guess the overall approach should be safe as long as you make sure
that last_node is always valid or NULL, though I wonder if it wouldn't
be a more straightforward improvement to make sure the table has a
power-of-two size and then using a bitwise AND to truncate the
index... with that I think you'd maybe just have a single-cycle
lengthening of the critical path? Though we would need a new helper
for that in nospec.h.

> 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]) {

I guess I'm not sure if eliding the array_index_nospec() is worth
adding a new branch here that you could mispredict... probably depends
on your workload, I guess?

> +                               data->last_index = index;
> +                               data->last_node = data->nodes[index];

This seems a bit functionally broken - if data->nodes[index] is NULL,
you just leave data->last_node set to its previous value and return
that?


> +                       }
> +               }
> +               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
>