[PATCH 09/19] io-wq: add IO_WORKER_F_SCHEDULED and its friends

[Hao Xu <hao.xu@linux.dev>]

From: Hao Xu <howeyxu@tencent.com>

Signed-off-by: Hao Xu <howeyxu@tencent.com>
---
 io_uring/io-wq.c | 29 ++---------------------------
 io_uring/io-wq.h | 43 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 45 insertions(+), 27 deletions(-)

diff --git a/io_uring/io-wq.c b/io_uring/io-wq.c
index 7e58bb5857ee..fe4faff79cf8 100644
--- a/io_uring/io-wq.c
+++ b/io_uring/io-wq.c
@@ -32,33 +32,6 @@ enum {
 	IO_ACCT_STALLED_BIT	= 0,	/* stalled on hash */
 };
 
-/*
- * One for each thread in a wqe pool
- */
-struct io_worker {
-	refcount_t ref;
-	unsigned flags;
-	struct hlist_nulls_node nulls_node;
-	struct list_head all_list;
-	struct task_struct *task;
-	struct io_wqe *wqe;
-
-	struct io_wq_work *cur_work;
-	struct io_wq_work *next_work;
-	raw_spinlock_t lock;
-
-	struct completion ref_done;
-
-	unsigned long create_state;
-	struct callback_head create_work;
-	int create_index;
-
-	union {
-		struct rcu_head rcu;
-		struct work_struct work;
-	};
-};
-
 #if BITS_PER_LONG == 64
 #define IO_WQ_HASH_ORDER	6
 #else
@@ -426,6 +399,7 @@ static void io_wqe_dec_running(struct io_worker *worker)
 		if (!io_worker_test_submit(worker))
 			return;
 
+		io_worker_set_scheduled(worker);
 		raw_spin_lock(&wqe->lock);
 		rcu_read_lock();
 		activated = io_wqe_activate_free_worker(wqe, acct);
@@ -706,6 +680,7 @@ static void io_wqe_worker_let(struct io_worker *worker)
 		do {
 			enum io_uringlet_state submit_state;
 
+			io_worker_clean_scheduled(worker);
 			io_worker_set_submit(worker);
 			submit_state = wq->do_work(wq->private);
 			io_worker_clean_submit(worker);
diff --git a/io_uring/io-wq.h b/io_uring/io-wq.h
index 1485e9009784..81146dba2ae6 100644
--- a/io_uring/io-wq.h
+++ b/io_uring/io-wq.h
@@ -3,9 +3,37 @@
 
 #include <linux/refcount.h>
 #include <linux/io_uring_types.h>
+#include <linux/list_nulls.h>
 
 struct io_wq;
 
+/*
+ * One for each thread in a wqe pool
+ */
+struct io_worker {
+	refcount_t ref;
+	unsigned flags;
+	struct hlist_nulls_node nulls_node;
+	struct list_head all_list;
+	struct task_struct *task;
+	struct io_wqe *wqe;
+
+	struct io_wq_work *cur_work;
+	struct io_wq_work *next_work;
+	raw_spinlock_t lock;
+
+	struct completion ref_done;
+
+	unsigned long create_state;
+	struct callback_head create_work;
+	int create_index;
+
+	union {
+		struct rcu_head rcu;
+		struct work_struct work;
+	};
+};
+
 enum {
 	IO_WQ_WORK_CANCEL	= 1,
 	IO_WQ_WORK_HASHED	= 2,
@@ -97,6 +125,21 @@ static inline bool io_wq_current_is_worker(void)
 		current->worker_private;
 }
 
+static inline void io_worker_set_scheduled(struct io_worker *worker)
+{
+	worker->flags |= IO_WORKER_F_SCHEDULED;
+}
+
+static inline void io_worker_clean_scheduled(struct io_worker *worker)
+{
+	worker->flags &= ~IO_WORKER_F_SCHEDULED;
+}
+
+static inline bool io_worker_test_scheduled(struct io_worker *worker)
+{
+	return worker->flags & IO_WORKER_F_SCHEDULED;
+}
+
 extern struct io_wq *io_init_wq_offload(struct io_ring_ctx *ctx,
 					struct task_struct *task);
 extern int io_uringlet_offload(struct io_wq *wq);
-- 
2.25.1