[PATCH 6/7] io_uring: account the ring context and SQPOLL thread data to memcg

[Tao Cui <cui.tao@linux.dev>]

From: Tao Cui <cuitao@kylinos.cn>

The io_ring_ctx and the SQPOLL io_sqd are the last per-ring allocations
made with plain GFP_KERNEL.  The larger surfaces around them are
already charged: io_mapped_ubuf and io_rsrc_node since commit
075d5837bafc ("io_uring/rsrc: charge buffer arrays and nodes to memcg"),
the tctx node and its XArrays since commit 0b34e7b20a48 ("io_uring:
charge tctx nodes and personalities").

Charge the context and the SQPOLL thread data the same way.

Signed-off-by: Tao Cui <cuitao@kylinos.cn>
---
 io_uring/io_uring.c | 2 +-
 io_uring/rsrc.c     | 3 ++-
 io_uring/sqpoll.c   | 2 +-
 io_uring/tctx.c     | 4 ++--
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/io_uring/io_uring.c b/io_uring/io_uring.c
index 03724419c4e1..d39cdde49a9f 100644
--- a/io_uring/io_uring.c
+++ b/io_uring/io_uring.c
@@ -229,7 +229,7 @@ static __cold struct io_ring_ctx *io_ring_ctx_alloc(struct io_uring_params *p)
 	int hash_bits;
 	bool ret;
 
-	ctx = kzalloc_obj(*ctx);
+	ctx = kzalloc_obj(*ctx, GFP_KERNEL_ACCOUNT);
 	if (!ctx)
 		return NULL;
 
diff --git a/io_uring/rsrc.c b/io_uring/rsrc.c
index e0345c3e764e..9a971c983686 100644
--- a/io_uring/rsrc.c
+++ b/io_uring/rsrc.c
@@ -825,7 +825,8 @@ static bool io_coalesce_buffer(struct page ***pages, int *nr_pages,
 	unsigned i, j;
 
 	/* Store head pages only*/
-	new_array = kvmalloc_objs(struct page *, nr_folios);
+	new_array = kvmalloc_objs(struct page *, nr_folios,
+				  GFP_KERNEL_ACCOUNT);
 	if (!new_array)
 		return false;
 
diff --git a/io_uring/sqpoll.c b/io_uring/sqpoll.c
index d387488db2de..b82d3a285337 100644
--- a/io_uring/sqpoll.c
+++ b/io_uring/sqpoll.c
@@ -172,7 +172,7 @@ static struct io_sq_data *io_get_sq_data(struct io_uring_params *p,
 			return sqd;
 	}
 
-	sqd = kzalloc_obj(*sqd);
+	sqd = kzalloc_obj(*sqd, GFP_KERNEL_ACCOUNT);
 	if (!sqd)
 		return ERR_PTR(-ENOMEM);
 
diff --git a/io_uring/tctx.c b/io_uring/tctx.c
index b4fecea3ea47..6673393100ec 100644
--- a/io_uring/tctx.c
+++ b/io_uring/tctx.c
@@ -23,7 +23,7 @@ static struct io_wq *io_init_wq_offload(struct io_ring_ctx *ctx,
 	mutex_lock(&ctx->uring_lock);
 	hash = ctx->hash_map;
 	if (!hash) {
-		hash = kzalloc_obj(*hash);
+		hash = kzalloc_obj(*hash, GFP_KERNEL_ACCOUNT);
 		if (!hash) {
 			mutex_unlock(&ctx->uring_lock);
 			return ERR_PTR(-ENOMEM);
@@ -85,7 +85,7 @@ __cold struct io_uring_task *io_uring_alloc_task_context(struct task_struct *tas
 	struct io_uring_task *tctx;
 	int ret;
 
-	tctx = kzalloc_obj(*tctx);
+	tctx = kzalloc_obj(*tctx, GFP_KERNEL_ACCOUNT);
 	if (unlikely(!tctx))
 		return ERR_PTR(-ENOMEM);
 
-- 
2.43.0