Re: [PATCH v3] io_uring/rsrc: add huge page accounting for registered buffers

Clément Léger
2026-05-04T07:31:58Z

On 5/2/26 04:31, Jens Axboe wrote:
> >
> Track huge page references in a per-ring xarray to prevent double
> accounting when the same huge page is used by multiple registered
> buffers, either within the same ring or across cloned rings.
>
> When registering buffers backed by huge pages, we need to account for
> RLIMIT_MEMLOCK. But if multiple buffers share the same huge page (common
> with cloned buffers), we must not account for the same page multiple
> times. Similarly, we must only unaccount when the last reference to a
> huge page is released.
>
> Maintain a per-ring xarray (hpage_acct) that tracks reference counts for
> each huge page. When registering a buffer, for each unique huge page,
> increment its accounting reference count, and only account pages that
> are newly added.
>
> When unregistering a buffer, for each unique huge page, decrement its
> refcount. Once the refcount hits zero, the page is unaccounted.
>
> Note: any account is done against the ctx->user that was assigned when
> the ring was setup. As before, if root is running the operation, no
> accounting is done.
>
> With these changes, any use of imu->acct_pages is also dead, hence kill
> it from struct io_mapped_ubuf. This shrinks it from 56b to 48b on a
> 64-bit arch. Additionally, hpage_already_acct() is gone, which was an
> O(M*M) scan over current + previous registrations.
>
> Signed-off-by: Jens Axboe <axboe@kernel.dk>
>
> ---
>
> See previous discussions here:
>
> https://lore.kernel.org/io-uring/20260119071039.2113739-1-danisjiang@gmail.com/
>
> diff --git a/include/linux/io_uring_types.h b/include/linux/io_uring_types.h
> index 244392026c6d..23b8891d5704 100644
> --- a/include/linux/io_uring_types.h
> +++ b/include/linux/io_uring_types.h
> @@ -446,6 +446,9 @@ struct io_ring_ctx {
> /* Stores zcrx object pointers of type struct io_zcrx_ifq */
> struct xarray zcrx_ctxs;
>
> + /* Used for accounting references on pages in registered buffers */
> + struct xarray hpage_acct;
> +
> u32 pers_next;
> struct xarray personalities;
>
> diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
> index 4ed998d60c09..fb6ed52bae61 100644
> --- a/io_uring/io_uring.c
> +++ b/io_uring/io_uring.c
> @@ -233,6 +233,7 @@ static __cold struct io_ring_ctx *io_ring_ctx_alloc(struct io_uring_params *p)
> return NULL;
>
> xa_init(&ctx->io_bl_xa);
> + xa_init(&ctx->hpage_acct);
>
> /*
> * Use 5 bits less than the max cq entries, that should give us around
> @@ -302,6 +303,7 @@ static __cold struct io_ring_ctx *io_ring_ctx_alloc(struct io_uring_params *p)
> io_free_alloc_caches(ctx);
> kvfree(ctx->cancel_table.hbs);
> xa_destroy(&ctx->io_bl_xa);
> + xa_destroy(&ctx->hpage_acct);
> kfree(ctx);
> return NULL;
> }
> @@ -2198,6 +2200,7 @@ static __cold void io_ring_ctx_free(struct io_ring_ctx *ctx)
> io_napi_free(ctx);
> kvfree(ctx->cancel_table.hbs);
> xa_destroy(&ctx->io_bl_xa);
> + xa_destroy(&ctx->hpage_acct);
> kfree(ctx);
> }
>
> diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
> index 650303626be6..ca22e07245c4 100644
> --- a/io_uring/rsrc.c
> +++ b/io_uring/rsrc.c
> @@ -28,7 +28,51 @@ struct io_rsrc_update {
> };
>
> static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
> - struct iovec *iov, struct page **last_hpage);
> + struct iovec *iov);
> +
> +static int hpage_acct_ref(struct io_ring_ctx *ctx, struct page *hpage,
> + bool *acct_new)
> +{
> + unsigned long key = (unsigned long) hpage;
> + unsigned long count;
> + void *entry;
> + int ret;
> +
> + lockdep_assert_held(&ctx->uring_lock);
> +
> + entry = xa_load(&ctx->hpage_acct, key);
> + if (!entry) {
> + ret = xa_reserve(&ctx->hpage_acct, key, GFP_KERNEL_ACCOUNT);
> + if (ret)
> + return ret;
> + }
> +
> + count = 1;
> + if (entry)
> + count = xa_to_value(entry) + 1;

Hi Jens,

Can't most of this be merged in the previous if/else ? ie:

entry = xa_load(&ctx->hpage_acct, key);>
count = 1;
if (!entry) {
ret = xa_reserve(&ctx->hpage_acct, key, GFP_KERNEL_ACCOUNT);
if (ret)
return ret;
*acct_new = true;
} else {
count = xa_to_value(entry) + 1;
*acct_new = false;
}

> + xa_store(&ctx->hpage_acct, key, xa_mk_value(count), GFP_KERNEL_ACCOUNT);
> + *acct_new = (count == 1);
> + return 0;
> +}
> +
> +static bool hpage_acct_unref(struct io_ring_ctx *ctx, struct page *hpage)
> +{
> + unsigned long key = (unsigned long) hpage;
> + unsigned long count;
> + void *entry;
> +
> + lockdep_assert_held(&ctx->uring_lock);
> +
> + entry = xa_load(&ctx->hpage_acct, key);
> + if (WARN_ON_ONCE(!entry))
> + return false;
> + count = xa_to_value(entry);
> + if (count == 1)
> + xa_erase(&ctx->hpage_acct, key);
> + else
> + xa_store(&ctx->hpage_acct, key, xa_mk_value(count - 1), GFP_KERNEL_ACCOUNT);
> + return count == 1;

Maybe something like this could easier to read ?:

if (count == 1) {
xa_erase(&ctx->hpage_acct, key);
return true;
}

xa_store(&ctx->hpage_acct, key, xa_mk_value(count - 1),
GFP_KERNEL_ACCOUNT);
return false;

> +}
>
> /* only define max */
> #define IORING_MAX_FIXED_FILES (1U << 20)
> @@ -124,15 +168,53 @@ static void io_free_imu(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
> kvfree(imu);
> }
>
> +static unsigned long io_buffer_unaccount_pages(struct io_ring_ctx *ctx,
> + struct io_mapped_ubuf *imu)
> +{
> + struct page *seen = NULL;
> + unsigned long acct = 0;
> + int i;
> +
> + if (imu->flags & IO_REGBUF_F_KBUF || !ctx->user)
> + return 0;
> +
> + for (i = 0; i < imu->nr_bvecs; i++) {
> + struct page *page = imu->bvec[i].bv_page;
> + struct page *hpage;
> +
> + if (!PageCompound(page)) {
> + acct++;
> + continue;
> + }
> +
> + hpage = compound_head(page);
> + if (hpage == seen)
> + continue;
> + seen = hpage;
> +
> + /* Unaccount on last reference */
> + if (hpage_acct_unref(ctx, hpage))
> + acct += page_size(hpage) >> PAGE_SHIFT;
> + cond_resched();
> + }
> +
> + return acct;
> +}
> +
> static void io_buffer_unmap(struct io_ring_ctx *ctx, struct io_mapped_ubuf *imu)
> {
> + unsigned long acct_pages = 0;
> +
> + /* Always decrement, so it works for cloned buffers too */
> + acct_pages = io_buffer_unaccount_pages(ctx, imu);
> +
> if (unlikely(refcount_read(&imu->refs) > 1)) {
> if (!refcount_dec_and_test(&imu->refs))
> return;
> }
>
> - if (imu->acct_pages)
> - io_unaccount_mem(ctx->user, ctx->mm_account, imu->acct_pages);
> + if (acct_pages)
> + io_unaccount_mem(ctx->user, ctx->mm_account, acct_pages);
> imu->release(imu->priv);
> io_free_imu(ctx, imu);
> }
> @@ -282,7 +364,6 @@ static int __io_sqe_buffers_update(struct io_ring_ctx *ctx,
> {
> u64 __user *tags = u64_to_user_ptr(up->tags);
> struct iovec fast_iov, *iov;
> - struct page *last_hpage = NULL;
> struct iovec __user *uvec;
> u64 user_data = up->data;
> __u32 done;
> @@ -307,7 +388,7 @@ static int __io_sqe_buffers_update(struct io_ring_ctx *ctx,
> err = -EFAULT;
> break;
> }
> - node = io_sqe_buffer_register(ctx, iov, &last_hpage);
> + node = io_sqe_buffer_register(ctx, iov);
> if (IS_ERR(node)) {
> err = PTR_ERR(node);
> break;
> @@ -605,76 +686,79 @@ int io_sqe_buffers_unregister(struct io_ring_ctx *ctx)
> }
>
> /*
> - * Not super efficient, but this is just a registration time. And we do cache
> - * the last compound head, so generally we'll only do a full search if we don't
> - * match that one.
> - *
> - * We check if the given compound head page has already been accounted, to
> - * avoid double accounting it. This allows us to account the full size of the
> - * page, not just the constituent pages of a huge page.
> + * Undo hpage_acct_ref() calls made during io_buffer_account_pin() on failure.
> + * This operates on the pages array since imu->bvec isn't populated yet.
> */
> -static bool headpage_already_acct(struct io_ring_ctx *ctx, struct page **pages,
> - int nr_pages, struct page *hpage)
> +static void io_buffer_unaccount_hpages(struct io_ring_ctx *ctx,
> + struct page **pages, int nr_pages)
> {
> - int i, j;
> + struct page *seen = NULL;
> + int i;
> +
> + if (!ctx->user)
> + return;
>
> - /* check current page array */
> for (i = 0; i < nr_pages; i++) {
> + struct page *hpage;
> +
> if (!PageCompound(pages[i]))
> continue;
> - if (compound_head(pages[i]) == hpage)
> - return true;
> - }
> -
> - /* check previously registered pages */
> - for (i = 0; i < ctx->buf_table.nr; i++) {
> - struct io_rsrc_node *node = ctx->buf_table.nodes[i];
> - struct io_mapped_ubuf *imu;
>
> - if (!node)
> + hpage = compound_head(pages[i]);
> + if (hpage == seen)
> continue;
> - imu = node->buf;
> - for (j = 0; j < imu->nr_bvecs; j++) {
> - if (!PageCompound(imu->bvec[j].bv_page))
> - continue;
> - if (compound_head(imu->bvec[j].bv_page) == hpage)
> - return true;
> - }
> - }
> + seen = hpage;
>
> - return false;
> + hpage_acct_unref(ctx, hpage);
> + cond_resched();
> + }
> }
>
> static int io_buffer_account_pin(struct io_ring_ctx *ctx, struct page **pages,
> - int nr_pages, struct io_mapped_ubuf *imu,
> - struct page **last_hpage)
> + int nr_pages)
> {
> + unsigned long acct_pages = 0;
> + struct page *seen = NULL;
> int i, ret;
>
> - imu->acct_pages = 0;
> + if (!ctx->user)
> + return 0;
> +
> for (i = 0; i < nr_pages; i++) {
> + struct page *hpage;
> + bool acct_new;
> +
> if (!PageCompound(pages[i])) {
> - imu->acct_pages++;
> - } else {
> - struct page *hpage;
> -
> - hpage = compound_head(pages[i]);
> - if (hpage == *last_hpage)
> - continue;
> - *last_hpage = hpage;
> - if (headpage_already_acct(ctx, pages, i, hpage))
> - continue;
> - imu->acct_pages += page_size(hpage) >> PAGE_SHIFT;
> + acct_pages++;
> + continue;
> }
> +
> + hpage = compound_head(pages[i]);
> + if (hpage == seen)
> + continue;
> + seen = hpage;
> +
> + ret = hpage_acct_ref(ctx, hpage, &acct_new);
> + if (ret) {
> + io_buffer_unaccount_hpages(ctx, pages, i);
> + return ret;
> + }
> + if (acct_new)
> + acct_pages += page_size(hpage) >> PAGE_SHIFT;
> + cond_resched();
> }
>
> - if (!imu->acct_pages)
> - return 0;
> + /* Try to account the memory */
> + if (acct_pages) {
> + ret = io_account_mem(ctx->user, ctx->mm_account, acct_pages);
> + if (ret) {
> + /* Undo the refs we just added */
> + io_buffer_unaccount_hpages(ctx, pages, nr_pages);
> + return ret;
> + }
> + }
>
> - ret = io_account_mem(ctx->user, ctx->mm_account, imu->acct_pages);
> - if (ret)
> - imu->acct_pages = 0;
> - return ret;
> + return 0;
> }
>
> static bool io_coalesce_buffer(struct page ***pages, int *nr_pages,
> @@ -763,8 +847,7 @@ bool io_check_coalesce_buffer(struct page **page_array, int nr_pages,
> }
>
> static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
> - struct iovec *iov,
> - struct page **last_hpage)
> + struct iovec *iov)
> {
> struct io_mapped_ubuf *imu = NULL;
> struct page **pages = NULL;
> @@ -811,7 +894,7 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
> goto done;
>
> imu->nr_bvecs = nr_pages;
> - ret = io_buffer_account_pin(ctx, pages, nr_pages, imu, last_hpage);
> + ret = io_buffer_account_pin(ctx, pages, nr_pages);
> if (ret)
> goto done;
>
> @@ -861,7 +944,6 @@ static struct io_rsrc_node *io_sqe_buffer_register(struct io_ring_ctx *ctx,
> int io_sqe_buffers_register(struct io_ring_ctx *ctx, void __user *arg,
> unsigned int nr_args, u64 __user *tags)
> {
> - struct page *last_hpage = NULL;
> struct io_rsrc_data data;
> struct iovec fast_iov, *iov = &fast_iov;
> const struct iovec __user *uvec;
> @@ -904,7 +986,7 @@ int io_sqe_buffers_register(struct io_ring_ctx *ctx, void __user *arg,
> }
> }
>
> - node = io_sqe_buffer_register(ctx, iov, &last_hpage);
> + node = io_sqe_buffer_register(ctx, iov);
> if (IS_ERR(node)) {
> ret = PTR_ERR(node);
> break;
> @@ -971,7 +1053,6 @@ int io_buffer_register_bvec(struct io_uring_cmd *cmd, struct request *rq,
>
> imu->ubuf = 0;
> imu->len = blk_rq_bytes(rq);
> - imu->acct_pages = 0;
> imu->folio_shift = PAGE_SHIFT;
> refcount_set(&imu->refs, 1);
> imu->release = release;
> @@ -1137,6 +1218,56 @@ int io_import_reg_buf(struct io_kiocb *req, struct iov_iter *iter,
> }
>
> /* Lock two rings at once. The rings must be different! */

This comment should be before lock_two_rings().

> +static int io_buffer_acct_cloned_hpages(struct io_ring_ctx *ctx,
> + struct io_mapped_ubuf *imu)
> +{
> + struct page *seen = NULL;
> + int i, ret = 0;
> +
> + if (imu->flags & IO_REGBUF_F_KBUF || !ctx->user)
> + return 0;
> +
> + for (i = 0; i < imu->nr_bvecs; i++) {
> + struct page *page = imu->bvec[i].bv_page;
> + struct page *hpage;
> + bool acct_new;
> +
> + if (!PageCompound(page))
> + continue;
> +
> + hpage = compound_head(page);
> + if (hpage == seen)
> + continue;
> + seen = hpage;
> +
> + /* Atomically add reference for cloned buffer */
> + ret = hpage_acct_ref(ctx, hpage, &acct_new);
> + if (ret)
> + break;
> +
> + cond_resched();
> + }
> +
> + if (!ret)
> + return 0;
> +
> + /* Undo refs we added for bvecs [0..i) */
> + seen = NULL;
> + for (int j = 0; j < i; j++) {
> + struct page *p = imu->bvec[j].bv_page;
> + struct page *hp;
> +
> + if (!PageCompound(p))
> + continue;
> + hp = compound_head(p);
> + if (hp == seen)
> + continue;
> + seen = hp;
> + hpage_acct_unref(ctx, hp);
> + }
> + return ret;
> +}
> +
> static void lock_two_rings(struct io_ring_ctx *ctx1, struct io_ring_ctx *ctx2)
> {
> if (ctx1 > ctx2)
> @@ -1218,6 +1349,14 @@ static int io_clone_buffers(struct io_ring_ctx *ctx, struct io_ring_ctx *src_ctx
>
> refcount_inc(&src_node->buf->refs);
> dst_node->buf = src_node->buf;
> + /* track compound references to clones */
> + ret = io_buffer_acct_cloned_hpages(ctx, src_node->buf);
> + if (ret) {
> + refcount_dec(&src_node->buf->refs);
> + io_cache_free(&ctx->node_cache, dst_node);
> + io_rsrc_data_free(ctx, &data);
> + return ret;
> + }
> }
> data.nodes[off++] = dst_node;
> i++;
> diff --git a/io_uring/rsrc.h b/io_uring/rsrc.h
> index 44e3386f7c1c..c0f8a18ec767 100644
> --- a/io_uring/rsrc.h
> +++ b/io_uring/rsrc.h
> @@ -38,7 +38,6 @@ struct io_mapped_ubuf {
> unsigned int nr_bvecs;
> unsigned int folio_shift;
> refcount_t refs;
> - unsigned long acct_pages;
> void (*release)(void *);
> void *priv;
> u8 flags;
>

Thanks,

Clément