io_uring: add blkcg accounting to offloaded operationsio-uring/blkcg-acc

There are a few operations that are offloaded to the worker threads. In
this case, we lose process context and end up in kthread context. This
results in ios to be not accounted to the issuing cgroup and
consequently end up as issued by root. Just like others, adopt the
personality of the blkcg too when issuing via the workqueues.

Signed-off-by: Dennis Zhou <dennis@kernel.org>

3 files changed, 60 insertions, 0 deletions

diff --git a/fs/io-wq.c b/fs/io-wq.c
index 4023c984686086..969c46dd6d3361 100644
--- a/fs/io-wq.c
+++ b/fs/io-wq.c
@@ -58,6 +58,7 @@ struct io_worker {
 
 	struct rcu_head rcu;
 	struct mm_struct *mm;
+	struct cgroup_subsys_state *blkcg_css;
 	const struct cred *cur_creds;
 	const struct cred *saved_creds;
 	struct files_struct *restore_files;
@@ -176,6 +177,13 @@ static bool __io_worker_unuse(struct io_wqe *wqe, struct io_worker *worker)
 		worker->mm = NULL;
 	}
 
+#ifdef CONFIG_BLK_CGROUP
+	if (worker->blkcg_css) {
+		kthread_associate_blkcg(NULL);
+		worker->blkcg_css = NULL;
+	}
+#endif
+
 	return dropped_lock;
 }
 
@@ -440,6 +448,20 @@ static void io_wq_switch_mm(struct io_worker *worker, struct io_wq_work *work)
 	work->flags |= IO_WQ_WORK_CANCEL;
 }
 
+#ifdef CONFIG_BLK_CGROUP
+static inline void io_wq_switch_blkcg(struct io_worker *worker,
+				      struct io_wq_work *work)
+{
+	if (work->blkcg_css != worker->blkcg_css)
+		kthread_associate_blkcg(work->blkcg_css);
+}
+#else
+static inline void io_wq_switch_blkcg(struct io_worker *worker,
+				      struct io_wq_work *work)
+{
+}
+#endif
+
 static void io_wq_switch_creds(struct io_worker *worker,
 			       struct io_wq_work *work)
 {
@@ -466,6 +488,7 @@ static void io_impersonate_work(struct io_worker *worker,
 		io_wq_switch_mm(worker, work);
 	if (worker->cur_creds != work->creds)
 		io_wq_switch_creds(worker, work);
+	io_wq_switch_blkcg(worker, work);
 }
 
 static void io_assign_current_work(struct io_worker *worker,
diff --git a/fs/io-wq.h b/fs/io-wq.h
index 5ba12de7572f0f..00a5382c558bba 100644
--- a/fs/io-wq.h
+++ b/fs/io-wq.h
@@ -88,6 +88,7 @@ struct io_wq_work {
 	void (*func)(struct io_wq_work **);
 	struct files_struct *files;
 	struct mm_struct *mm;
+	struct cgroup_subsys_state *blkcg_css;
 	const struct cred *creds;
 	struct fs_struct *fs;
 	unsigned flags;
diff --git a/fs/io_uring.c b/fs/io_uring.c
index 5190bfb6a6657e..2f22c6a4bfdb29 100644
--- a/fs/io_uring.c
+++ b/fs/io_uring.c
@@ -79,6 +79,7 @@
 #include <linux/fs_struct.h>
 #include <linux/splice.h>
 #include <linux/task_work.h>
+#include <linux/blk-cgroup.h>
 
 #define CREATE_TRACE_POINTS
 #include <trace/events/io_uring.h>
@@ -673,6 +674,8 @@ struct io_op_def {
 	unsigned		async_ctx : 1;
 	/* needs current->mm setup, does mm access */
 	unsigned		needs_mm : 1;
+	/* needs blkcg context, issues async io */
+	unsigned		needs_blkcg : 1;
 	/* needs req->file assigned */
 	unsigned		needs_file : 1;
 	/* needs req->file assigned IFF fd is >= 0 */
@@ -699,6 +702,7 @@ static const struct io_op_def io_op_defs[] = {
 	[IORING_OP_READV] = {
 		.async_ctx		= 1,
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollin			= 1,
@@ -707,20 +711,24 @@ static const struct io_op_def io_op_defs[] = {
 	[IORING_OP_WRITEV] = {
 		.async_ctx		= 1,
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.hash_reg_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollout		= 1,
 	},
 	[IORING_OP_FSYNC] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 	},
 	[IORING_OP_READ_FIXED] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollin			= 1,
 	},
 	[IORING_OP_WRITE_FIXED] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.hash_reg_file		= 1,
 		.unbound_nonreg_file	= 1,
@@ -732,10 +740,12 @@ static const struct io_op_def io_op_defs[] = {
 	},
 	[IORING_OP_POLL_REMOVE] = {},
 	[IORING_OP_SYNC_FILE_RANGE] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 	},
 	[IORING_OP_SENDMSG] = {
 		.async_ctx		= 1,
+		.needs_blkcg		= 1,
 		.needs_mm		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
@@ -776,15 +786,18 @@ static const struct io_op_def io_op_defs[] = {
 		.pollout		= 1,
 	},
 	[IORING_OP_FALLOCATE] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 	},
 	[IORING_OP_OPENAT] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.fd_non_neg		= 1,
 		.file_table		= 1,
 		.needs_fs		= 1,
 	},
 	[IORING_OP_CLOSE] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.file_table		= 1,
 	},
@@ -794,12 +807,14 @@ static const struct io_op_def io_op_defs[] = {
 	},
 	[IORING_OP_STATX] = {
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.fd_non_neg		= 1,
 		.needs_fs		= 1,
 	},
 	[IORING_OP_READ] = {
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollin			= 1,
@@ -807,18 +822,22 @@ static const struct io_op_def io_op_defs[] = {
 	},
 	[IORING_OP_WRITE] = {
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollout		= 1,
 	},
 	[IORING_OP_FADVISE] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 	},
 	[IORING_OP_MADVISE] = {
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 	},
 	[IORING_OP_SEND] = {
 		.needs_mm		= 1,
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.unbound_nonreg_file	= 1,
 		.pollout		= 1,
@@ -831,6 +850,7 @@ static const struct io_op_def io_op_defs[] = {
 		.buffer_select		= 1,
 	},
 	[IORING_OP_OPENAT2] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.fd_non_neg		= 1,
 		.file_table		= 1,
@@ -841,6 +861,7 @@ static const struct io_op_def io_op_defs[] = {
 		.file_table		= 1,
 	},
 	[IORING_OP_SPLICE] = {
+		.needs_blkcg		= 1,
 		.needs_file		= 1,
 		.hash_reg_file		= 1,
 		.unbound_nonreg_file	= 1,
@@ -1017,6 +1038,17 @@ static inline void io_req_work_grab_env(struct io_kiocb *req,
 		mmgrab(current->mm);
 		req->work.mm = current->mm;
 	}
+#ifdef CONFIG_BLK_CGROUP
+	if (!req->work.blkcg_css && def->needs_blkcg) {
+		req->work.blkcg_css = blkcg_css();
+		/*
+		 * If we don't get a css, then the kthread will already
+		 * be in kernel context and grab the root css.
+		 */
+		if (!css_tryget_online(req->work.blkcg_css))
+			req->work.blkcg_css = NULL;
+	}
+#endif
 	if (!req->work.creds)
 		req->work.creds = get_current_cred();
 	if (!req->work.fs && def->needs_fs) {
@@ -1039,6 +1071,10 @@ static inline void io_req_work_drop_env(struct io_kiocb *req)
 		mmdrop(req->work.mm);
 		req->work.mm = NULL;
 	}
+#ifdef CONFIG_BLK_CGROUP
+	if (req->work.blkcg_css)
+		css_put(req->work.blkcg_css);
+#endif
 	if (req->work.creds) {
 		put_cred(req->work.creds);
 		req->work.creds = NULL;