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

[Jann Horn <jannh@google.com>]

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
>