io-wq.c revision 05c5f4ee
1// SPDX-License-Identifier: GPL-2.0
2/*
3 * Basic worker thread pool for io_uring
4 *
5 * Copyright (C) 2019 Jens Axboe
6 *
7 */
8#include <linux/kernel.h>
9#include <linux/init.h>
10#include <linux/errno.h>
11#include <linux/sched/signal.h>
12#include <linux/percpu.h>
13#include <linux/slab.h>
14#include <linux/rculist_nulls.h>
15#include <linux/cpu.h>
16#include <linux/tracehook.h>
17
18#include "io-wq.h"
19
20#define WORKER_IDLE_TIMEOUT	(5 * HZ)
21
22enum {
23	IO_WORKER_F_UP		= 1,	/* up and active */
24	IO_WORKER_F_RUNNING	= 2,	/* account as running */
25	IO_WORKER_F_FREE	= 4,	/* worker on free list */
26	IO_WORKER_F_BOUND	= 8,	/* is doing bounded work */
27};
28
29enum {
30	IO_WQ_BIT_EXIT		= 0,	/* wq exiting */
31};
32
33enum {
34	IO_ACCT_STALLED_BIT	= 0,	/* stalled on hash */
35};
36
37/*
38 * One for each thread in a wqe pool
39 */
40struct io_worker {
41	refcount_t ref;
42	unsigned flags;
43	struct hlist_nulls_node nulls_node;
44	struct list_head all_list;
45	struct task_struct *task;
46	struct io_wqe *wqe;
47
48	struct io_wq_work *cur_work;
49	spinlock_t lock;
50
51	struct completion ref_done;
52
53	unsigned long create_state;
54	struct callback_head create_work;
55	int create_index;
56
57	struct rcu_head rcu;
58};
59
60#if BITS_PER_LONG == 64
61#define IO_WQ_HASH_ORDER	6
62#else
63#define IO_WQ_HASH_ORDER	5
64#endif
65
66#define IO_WQ_NR_HASH_BUCKETS	(1u << IO_WQ_HASH_ORDER)
67
68struct io_wqe_acct {
69	unsigned nr_workers;
70	unsigned max_workers;
71	int index;
72	atomic_t nr_running;
73	struct io_wq_work_list work_list;
74	unsigned long flags;
75};
76
77enum {
78	IO_WQ_ACCT_BOUND,
79	IO_WQ_ACCT_UNBOUND,
80	IO_WQ_ACCT_NR,
81};
82
83/*
84 * Per-node worker thread pool
85 */
86struct io_wqe {
87	raw_spinlock_t lock;
88	struct io_wqe_acct acct[2];
89
90	int node;
91
92	struct hlist_nulls_head free_list;
93	struct list_head all_list;
94
95	struct wait_queue_entry wait;
96
97	struct io_wq *wq;
98	struct io_wq_work *hash_tail[IO_WQ_NR_HASH_BUCKETS];
99
100	cpumask_var_t cpu_mask;
101};
102
103/*
104 * Per io_wq state
105  */
106struct io_wq {
107	unsigned long state;
108
109	free_work_fn *free_work;
110	io_wq_work_fn *do_work;
111
112	struct io_wq_hash *hash;
113
114	atomic_t worker_refs;
115	struct completion worker_done;
116
117	struct hlist_node cpuhp_node;
118
119	struct task_struct *task;
120
121	struct io_wqe *wqes[];
122};
123
124static enum cpuhp_state io_wq_online;
125
126struct io_cb_cancel_data {
127	work_cancel_fn *fn;
128	void *data;
129	int nr_running;
130	int nr_pending;
131	bool cancel_all;
132};
133
134static void create_io_worker(struct io_wq *wq, struct io_wqe *wqe, int index);
135static void io_wqe_dec_running(struct io_worker *worker);
136
137static bool io_worker_get(struct io_worker *worker)
138{
139	return refcount_inc_not_zero(&worker->ref);
140}
141
142static void io_worker_release(struct io_worker *worker)
143{
144	if (refcount_dec_and_test(&worker->ref))
145		complete(&worker->ref_done);
146}
147
148static inline struct io_wqe_acct *io_get_acct(struct io_wqe *wqe, bool bound)
149{
150	return &wqe->acct[bound ? IO_WQ_ACCT_BOUND : IO_WQ_ACCT_UNBOUND];
151}
152
153static inline struct io_wqe_acct *io_work_get_acct(struct io_wqe *wqe,
154						   struct io_wq_work *work)
155{
156	return io_get_acct(wqe, !(work->flags & IO_WQ_WORK_UNBOUND));
157}
158
159static inline struct io_wqe_acct *io_wqe_get_acct(struct io_worker *worker)
160{
161	return io_get_acct(worker->wqe, worker->flags & IO_WORKER_F_BOUND);
162}
163
164static void io_worker_ref_put(struct io_wq *wq)
165{
166	if (atomic_dec_and_test(&wq->worker_refs))
167		complete(&wq->worker_done);
168}
169
170static void io_worker_exit(struct io_worker *worker)
171{
172	struct io_wqe *wqe = worker->wqe;
173	struct io_wqe_acct *acct = io_wqe_get_acct(worker);
174
175	if (refcount_dec_and_test(&worker->ref))
176		complete(&worker->ref_done);
177	wait_for_completion(&worker->ref_done);
178
179	raw_spin_lock(&wqe->lock);
180	if (worker->flags & IO_WORKER_F_FREE)
181		hlist_nulls_del_rcu(&worker->nulls_node);
182	list_del_rcu(&worker->all_list);
183	acct->nr_workers--;
184	preempt_disable();
185	io_wqe_dec_running(worker);
186	worker->flags = 0;
187	current->flags &= ~PF_IO_WORKER;
188	preempt_enable();
189	raw_spin_unlock(&wqe->lock);
190
191	kfree_rcu(worker, rcu);
192	io_worker_ref_put(wqe->wq);
193	do_exit(0);
194}
195
196static inline bool io_acct_run_queue(struct io_wqe_acct *acct)
197{
198	if (!wq_list_empty(&acct->work_list) &&
199	    !test_bit(IO_ACCT_STALLED_BIT, &acct->flags))
200		return true;
201	return false;
202}
203
204/*
205 * Check head of free list for an available worker. If one isn't available,
206 * caller must create one.
207 */
208static bool io_wqe_activate_free_worker(struct io_wqe *wqe,
209					struct io_wqe_acct *acct)
210	__must_hold(RCU)
211{
212	struct hlist_nulls_node *n;
213	struct io_worker *worker;
214
215	/*
216	 * Iterate free_list and see if we can find an idle worker to
217	 * activate. If a given worker is on the free_list but in the process
218	 * of exiting, keep trying.
219	 */
220	hlist_nulls_for_each_entry_rcu(worker, n, &wqe->free_list, nulls_node) {
221		if (!io_worker_get(worker))
222			continue;
223		if (io_wqe_get_acct(worker) != acct) {
224			io_worker_release(worker);
225			continue;
226		}
227		if (wake_up_process(worker->task)) {
228			io_worker_release(worker);
229			return true;
230		}
231		io_worker_release(worker);
232	}
233
234	return false;
235}
236
237/*
238 * We need a worker. If we find a free one, we're good. If not, and we're
239 * below the max number of workers, create one.
240 */
241static void io_wqe_create_worker(struct io_wqe *wqe, struct io_wqe_acct *acct)
242{
243	bool do_create = false;
244
245	/*
246	 * Most likely an attempt to queue unbounded work on an io_wq that
247	 * wasn't setup with any unbounded workers.
248	 */
249	if (unlikely(!acct->max_workers))
250		pr_warn_once("io-wq is not configured for unbound workers");
251
252	raw_spin_lock(&wqe->lock);
253	if (acct->nr_workers < acct->max_workers) {
254		acct->nr_workers++;
255		do_create = true;
256	}
257	raw_spin_unlock(&wqe->lock);
258	if (do_create) {
259		atomic_inc(&acct->nr_running);
260		atomic_inc(&wqe->wq->worker_refs);
261		create_io_worker(wqe->wq, wqe, acct->index);
262	}
263}
264
265static void io_wqe_inc_running(struct io_worker *worker)
266{
267	struct io_wqe_acct *acct = io_wqe_get_acct(worker);
268
269	atomic_inc(&acct->nr_running);
270}
271
272static void create_worker_cb(struct callback_head *cb)
273{
274	struct io_worker *worker;
275	struct io_wq *wq;
276	struct io_wqe *wqe;
277	struct io_wqe_acct *acct;
278	bool do_create = false;
279
280	worker = container_of(cb, struct io_worker, create_work);
281	wqe = worker->wqe;
282	wq = wqe->wq;
283	acct = &wqe->acct[worker->create_index];
284	raw_spin_lock(&wqe->lock);
285	if (acct->nr_workers < acct->max_workers) {
286		acct->nr_workers++;
287		do_create = true;
288	}
289	raw_spin_unlock(&wqe->lock);
290	if (do_create) {
291		create_io_worker(wq, wqe, worker->create_index);
292	} else {
293		atomic_dec(&acct->nr_running);
294		io_worker_ref_put(wq);
295	}
296	clear_bit_unlock(0, &worker->create_state);
297	io_worker_release(worker);
298}
299
300static void io_queue_worker_create(struct io_wqe *wqe, struct io_worker *worker,
301				   struct io_wqe_acct *acct)
302{
303	struct io_wq *wq = wqe->wq;
304
305	/* raced with exit, just ignore create call */
306	if (test_bit(IO_WQ_BIT_EXIT, &wq->state))
307		goto fail;
308	if (!io_worker_get(worker))
309		goto fail;
310	/*
311	 * create_state manages ownership of create_work/index. We should
312	 * only need one entry per worker, as the worker going to sleep
313	 * will trigger the condition, and waking will clear it once it
314	 * runs the task_work.
315	 */
316	if (test_bit(0, &worker->create_state) ||
317	    test_and_set_bit_lock(0, &worker->create_state))
318		goto fail_release;
319
320	init_task_work(&worker->create_work, create_worker_cb);
321	worker->create_index = acct->index;
322	if (!task_work_add(wq->task, &worker->create_work, TWA_SIGNAL))
323		return;
324	clear_bit_unlock(0, &worker->create_state);
325fail_release:
326	io_worker_release(worker);
327fail:
328	atomic_dec(&acct->nr_running);
329	io_worker_ref_put(wq);
330}
331
332static void io_wqe_dec_running(struct io_worker *worker)
333	__must_hold(wqe->lock)
334{
335	struct io_wqe_acct *acct = io_wqe_get_acct(worker);
336	struct io_wqe *wqe = worker->wqe;
337
338	if (!(worker->flags & IO_WORKER_F_UP))
339		return;
340
341	if (atomic_dec_and_test(&acct->nr_running) && io_acct_run_queue(acct)) {
342		atomic_inc(&acct->nr_running);
343		atomic_inc(&wqe->wq->worker_refs);
344		io_queue_worker_create(wqe, worker, acct);
345	}
346}
347
348/*
349 * Worker will start processing some work. Move it to the busy list, if
350 * it's currently on the freelist
351 */
352static void __io_worker_busy(struct io_wqe *wqe, struct io_worker *worker,
353			     struct io_wq_work *work)
354	__must_hold(wqe->lock)
355{
356	if (worker->flags & IO_WORKER_F_FREE) {
357		worker->flags &= ~IO_WORKER_F_FREE;
358		hlist_nulls_del_init_rcu(&worker->nulls_node);
359	}
360}
361
362/*
363 * No work, worker going to sleep. Move to freelist, and unuse mm if we
364 * have one attached. Dropping the mm may potentially sleep, so we drop
365 * the lock in that case and return success. Since the caller has to
366 * retry the loop in that case (we changed task state), we don't regrab
367 * the lock if we return success.
368 */
369static void __io_worker_idle(struct io_wqe *wqe, struct io_worker *worker)
370	__must_hold(wqe->lock)
371{
372	if (!(worker->flags & IO_WORKER_F_FREE)) {
373		worker->flags |= IO_WORKER_F_FREE;
374		hlist_nulls_add_head_rcu(&worker->nulls_node, &wqe->free_list);
375	}
376}
377
378static inline unsigned int io_get_work_hash(struct io_wq_work *work)
379{
380	return work->flags >> IO_WQ_HASH_SHIFT;
381}
382
383static void io_wait_on_hash(struct io_wqe *wqe, unsigned int hash)
384{
385	struct io_wq *wq = wqe->wq;
386
387	spin_lock_irq(&wq->hash->wait.lock);
388	if (list_empty(&wqe->wait.entry)) {
389		__add_wait_queue(&wq->hash->wait, &wqe->wait);
390		if (!test_bit(hash, &wq->hash->map)) {
391			__set_current_state(TASK_RUNNING);
392			list_del_init(&wqe->wait.entry);
393		}
394	}
395	spin_unlock_irq(&wq->hash->wait.lock);
396}
397
398static struct io_wq_work *io_get_next_work(struct io_wqe_acct *acct,
399					   struct io_worker *worker)
400	__must_hold(wqe->lock)
401{
402	struct io_wq_work_node *node, *prev;
403	struct io_wq_work *work, *tail;
404	unsigned int stall_hash = -1U;
405	struct io_wqe *wqe = worker->wqe;
406
407	wq_list_for_each(node, prev, &acct->work_list) {
408		unsigned int hash;
409
410		work = container_of(node, struct io_wq_work, list);
411
412		/* not hashed, can run anytime */
413		if (!io_wq_is_hashed(work)) {
414			wq_list_del(&acct->work_list, node, prev);
415			return work;
416		}
417
418		hash = io_get_work_hash(work);
419		/* all items with this hash lie in [work, tail] */
420		tail = wqe->hash_tail[hash];
421
422		/* hashed, can run if not already running */
423		if (!test_and_set_bit(hash, &wqe->wq->hash->map)) {
424			wqe->hash_tail[hash] = NULL;
425			wq_list_cut(&acct->work_list, &tail->list, prev);
426			return work;
427		}
428		if (stall_hash == -1U)
429			stall_hash = hash;
430		/* fast forward to a next hash, for-each will fix up @prev */
431		node = &tail->list;
432	}
433
434	if (stall_hash != -1U) {
435		/*
436		 * Set this before dropping the lock to avoid racing with new
437		 * work being added and clearing the stalled bit.
438		 */
439		set_bit(IO_ACCT_STALLED_BIT, &acct->flags);
440		raw_spin_unlock(&wqe->lock);
441		io_wait_on_hash(wqe, stall_hash);
442		raw_spin_lock(&wqe->lock);
443	}
444
445	return NULL;
446}
447
448static bool io_flush_signals(void)
449{
450	if (unlikely(test_thread_flag(TIF_NOTIFY_SIGNAL))) {
451		__set_current_state(TASK_RUNNING);
452		tracehook_notify_signal();
453		return true;
454	}
455	return false;
456}
457
458static void io_assign_current_work(struct io_worker *worker,
459				   struct io_wq_work *work)
460{
461	if (work) {
462		io_flush_signals();
463		cond_resched();
464	}
465
466	spin_lock(&worker->lock);
467	worker->cur_work = work;
468	spin_unlock(&worker->lock);
469}
470
471static void io_wqe_enqueue(struct io_wqe *wqe, struct io_wq_work *work);
472
473static void io_worker_handle_work(struct io_worker *worker)
474	__releases(wqe->lock)
475{
476	struct io_wqe_acct *acct = io_wqe_get_acct(worker);
477	struct io_wqe *wqe = worker->wqe;
478	struct io_wq *wq = wqe->wq;
479	bool do_kill = test_bit(IO_WQ_BIT_EXIT, &wq->state);
480
481	do {
482		struct io_wq_work *work;
483get_next:
484		/*
485		 * If we got some work, mark us as busy. If we didn't, but
486		 * the list isn't empty, it means we stalled on hashed work.
487		 * Mark us stalled so we don't keep looking for work when we
488		 * can't make progress, any work completion or insertion will
489		 * clear the stalled flag.
490		 */
491		work = io_get_next_work(acct, worker);
492		if (work)
493			__io_worker_busy(wqe, worker, work);
494
495		raw_spin_unlock(&wqe->lock);
496		if (!work)
497			break;
498		io_assign_current_work(worker, work);
499		__set_current_state(TASK_RUNNING);
500
501		/* handle a whole dependent link */
502		do {
503			struct io_wq_work *next_hashed, *linked;
504			unsigned int hash = io_get_work_hash(work);
505
506			next_hashed = wq_next_work(work);
507
508			if (unlikely(do_kill) && (work->flags & IO_WQ_WORK_UNBOUND))
509				work->flags |= IO_WQ_WORK_CANCEL;
510			wq->do_work(work);
511			io_assign_current_work(worker, NULL);
512
513			linked = wq->free_work(work);
514			work = next_hashed;
515			if (!work && linked && !io_wq_is_hashed(linked)) {
516				work = linked;
517				linked = NULL;
518			}
519			io_assign_current_work(worker, work);
520			if (linked)
521				io_wqe_enqueue(wqe, linked);
522
523			if (hash != -1U && !next_hashed) {
524				clear_bit(hash, &wq->hash->map);
525				clear_bit(IO_ACCT_STALLED_BIT, &acct->flags);
526				if (wq_has_sleeper(&wq->hash->wait))
527					wake_up(&wq->hash->wait);
528				raw_spin_lock(&wqe->lock);
529				/* skip unnecessary unlock-lock wqe->lock */
530				if (!work)
531					goto get_next;
532				raw_spin_unlock(&wqe->lock);
533			}
534		} while (work);
535
536		raw_spin_lock(&wqe->lock);
537	} while (1);
538}
539
540static int io_wqe_worker(void *data)
541{
542	struct io_worker *worker = data;
543	struct io_wqe_acct *acct = io_wqe_get_acct(worker);
544	struct io_wqe *wqe = worker->wqe;
545	struct io_wq *wq = wqe->wq;
546	bool last_timeout = false;
547	char buf[TASK_COMM_LEN];
548
549	worker->flags |= (IO_WORKER_F_UP | IO_WORKER_F_RUNNING);
550
551	snprintf(buf, sizeof(buf), "iou-wrk-%d", wq->task->pid);
552	set_task_comm(current, buf);
553
554	while (!test_bit(IO_WQ_BIT_EXIT, &wq->state)) {
555		long ret;
556
557		set_current_state(TASK_INTERRUPTIBLE);
558loop:
559		raw_spin_lock(&wqe->lock);
560		if (io_acct_run_queue(acct)) {
561			io_worker_handle_work(worker);
562			goto loop;
563		}
564		/* timed out, exit unless we're the last worker */
565		if (last_timeout && acct->nr_workers > 1) {
566			raw_spin_unlock(&wqe->lock);
567			__set_current_state(TASK_RUNNING);
568			break;
569		}
570		last_timeout = false;
571		__io_worker_idle(wqe, worker);
572		raw_spin_unlock(&wqe->lock);
573		if (io_flush_signals())
574			continue;
575		ret = schedule_timeout(WORKER_IDLE_TIMEOUT);
576		if (signal_pending(current)) {
577			struct ksignal ksig;
578
579			if (!get_signal(&ksig))
580				continue;
581			if (fatal_signal_pending(current))
582				break;
583			continue;
584		}
585		last_timeout = !ret;
586	}
587
588	if (test_bit(IO_WQ_BIT_EXIT, &wq->state)) {
589		raw_spin_lock(&wqe->lock);
590		io_worker_handle_work(worker);
591	}
592
593	io_worker_exit(worker);
594	return 0;
595}
596
597/*
598 * Called when a worker is scheduled in. Mark us as currently running.
599 */
600void io_wq_worker_running(struct task_struct *tsk)
601{
602	struct io_worker *worker = tsk->pf_io_worker;
603
604	if (!worker)
605		return;
606	if (!(worker->flags & IO_WORKER_F_UP))
607		return;
608	if (worker->flags & IO_WORKER_F_RUNNING)
609		return;
610	worker->flags |= IO_WORKER_F_RUNNING;
611	io_wqe_inc_running(worker);
612}
613
614/*
615 * Called when worker is going to sleep. If there are no workers currently
616 * running and we have work pending, wake up a free one or create a new one.
617 */
618void io_wq_worker_sleeping(struct task_struct *tsk)
619{
620	struct io_worker *worker = tsk->pf_io_worker;
621
622	if (!worker)
623		return;
624	if (!(worker->flags & IO_WORKER_F_UP))
625		return;
626	if (!(worker->flags & IO_WORKER_F_RUNNING))
627		return;
628
629	worker->flags &= ~IO_WORKER_F_RUNNING;
630
631	raw_spin_lock(&worker->wqe->lock);
632	io_wqe_dec_running(worker);
633	raw_spin_unlock(&worker->wqe->lock);
634}
635
636static void create_io_worker(struct io_wq *wq, struct io_wqe *wqe, int index)
637{
638	struct io_wqe_acct *acct = &wqe->acct[index];
639	struct io_worker *worker;
640	struct task_struct *tsk;
641
642	__set_current_state(TASK_RUNNING);
643
644	worker = kzalloc_node(sizeof(*worker), GFP_KERNEL, wqe->node);
645	if (!worker)
646		goto fail;
647
648	refcount_set(&worker->ref, 1);
649	worker->nulls_node.pprev = NULL;
650	worker->wqe = wqe;
651	spin_lock_init(&worker->lock);
652	init_completion(&worker->ref_done);
653
654	tsk = create_io_thread(io_wqe_worker, worker, wqe->node);
655	if (IS_ERR(tsk)) {
656		kfree(worker);
657fail:
658		atomic_dec(&acct->nr_running);
659		raw_spin_lock(&wqe->lock);
660		acct->nr_workers--;
661		raw_spin_unlock(&wqe->lock);
662		io_worker_ref_put(wq);
663		return;
664	}
665
666	tsk->pf_io_worker = worker;
667	worker->task = tsk;
668	set_cpus_allowed_ptr(tsk, wqe->cpu_mask);
669	tsk->flags |= PF_NO_SETAFFINITY;
670
671	raw_spin_lock(&wqe->lock);
672	hlist_nulls_add_head_rcu(&worker->nulls_node, &wqe->free_list);
673	list_add_tail_rcu(&worker->all_list, &wqe->all_list);
674	worker->flags |= IO_WORKER_F_FREE;
675	if (index == IO_WQ_ACCT_BOUND)
676		worker->flags |= IO_WORKER_F_BOUND;
677	raw_spin_unlock(&wqe->lock);
678	wake_up_new_task(tsk);
679}
680
681/*
682 * Iterate the passed in list and call the specific function for each
683 * worker that isn't exiting
684 */
685static bool io_wq_for_each_worker(struct io_wqe *wqe,
686				  bool (*func)(struct io_worker *, void *),
687				  void *data)
688{
689	struct io_worker *worker;
690	bool ret = false;
691
692	list_for_each_entry_rcu(worker, &wqe->all_list, all_list) {
693		if (io_worker_get(worker)) {
694			/* no task if node is/was offline */
695			if (worker->task)
696				ret = func(worker, data);
697			io_worker_release(worker);
698			if (ret)
699				break;
700		}
701	}
702
703	return ret;
704}
705
706static bool io_wq_worker_wake(struct io_worker *worker, void *data)
707{
708	set_notify_signal(worker->task);
709	wake_up_process(worker->task);
710	return false;
711}
712
713static bool io_wq_work_match_all(struct io_wq_work *work, void *data)
714{
715	return true;
716}
717
718static void io_run_cancel(struct io_wq_work *work, struct io_wqe *wqe)
719{
720	struct io_wq *wq = wqe->wq;
721
722	do {
723		work->flags |= IO_WQ_WORK_CANCEL;
724		wq->do_work(work);
725		work = wq->free_work(work);
726	} while (work);
727}
728
729static void io_wqe_insert_work(struct io_wqe *wqe, struct io_wq_work *work)
730{
731	struct io_wqe_acct *acct = io_work_get_acct(wqe, work);
732	unsigned int hash;
733	struct io_wq_work *tail;
734
735	if (!io_wq_is_hashed(work)) {
736append:
737		wq_list_add_tail(&work->list, &acct->work_list);
738		return;
739	}
740
741	hash = io_get_work_hash(work);
742	tail = wqe->hash_tail[hash];
743	wqe->hash_tail[hash] = work;
744	if (!tail)
745		goto append;
746
747	wq_list_add_after(&work->list, &tail->list, &acct->work_list);
748}
749
750static void io_wqe_enqueue(struct io_wqe *wqe, struct io_wq_work *work)
751{
752	struct io_wqe_acct *acct = io_work_get_acct(wqe, work);
753	unsigned work_flags = work->flags;
754	bool do_create;
755
756	/*
757	 * If io-wq is exiting for this task, or if the request has explicitly
758	 * been marked as one that should not get executed, cancel it here.
759	 */
760	if (test_bit(IO_WQ_BIT_EXIT, &wqe->wq->state) ||
761	    (work->flags & IO_WQ_WORK_CANCEL)) {
762		io_run_cancel(work, wqe);
763		return;
764	}
765
766	raw_spin_lock(&wqe->lock);
767	io_wqe_insert_work(wqe, work);
768	clear_bit(IO_ACCT_STALLED_BIT, &acct->flags);
769
770	rcu_read_lock();
771	do_create = !io_wqe_activate_free_worker(wqe, acct);
772	rcu_read_unlock();
773
774	raw_spin_unlock(&wqe->lock);
775
776	if (do_create && ((work_flags & IO_WQ_WORK_CONCURRENT) ||
777	    !atomic_read(&acct->nr_running)))
778		io_wqe_create_worker(wqe, acct);
779}
780
781void io_wq_enqueue(struct io_wq *wq, struct io_wq_work *work)
782{
783	struct io_wqe *wqe = wq->wqes[numa_node_id()];
784
785	io_wqe_enqueue(wqe, work);
786}
787
788/*
789 * Work items that hash to the same value will not be done in parallel.
790 * Used to limit concurrent writes, generally hashed by inode.
791 */
792void io_wq_hash_work(struct io_wq_work *work, void *val)
793{
794	unsigned int bit;
795
796	bit = hash_ptr(val, IO_WQ_HASH_ORDER);
797	work->flags |= (IO_WQ_WORK_HASHED | (bit << IO_WQ_HASH_SHIFT));
798}
799
800static bool io_wq_worker_cancel(struct io_worker *worker, void *data)
801{
802	struct io_cb_cancel_data *match = data;
803
804	/*
805	 * Hold the lock to avoid ->cur_work going out of scope, caller
806	 * may dereference the passed in work.
807	 */
808	spin_lock(&worker->lock);
809	if (worker->cur_work &&
810	    match->fn(worker->cur_work, match->data)) {
811		set_notify_signal(worker->task);
812		match->nr_running++;
813	}
814	spin_unlock(&worker->lock);
815
816	return match->nr_running && !match->cancel_all;
817}
818
819static inline void io_wqe_remove_pending(struct io_wqe *wqe,
820					 struct io_wq_work *work,
821					 struct io_wq_work_node *prev)
822{
823	struct io_wqe_acct *acct = io_work_get_acct(wqe, work);
824	unsigned int hash = io_get_work_hash(work);
825	struct io_wq_work *prev_work = NULL;
826
827	if (io_wq_is_hashed(work) && work == wqe->hash_tail[hash]) {
828		if (prev)
829			prev_work = container_of(prev, struct io_wq_work, list);
830		if (prev_work && io_get_work_hash(prev_work) == hash)
831			wqe->hash_tail[hash] = prev_work;
832		else
833			wqe->hash_tail[hash] = NULL;
834	}
835	wq_list_del(&acct->work_list, &work->list, prev);
836}
837
838static void io_wqe_cancel_pending_work(struct io_wqe *wqe,
839				       struct io_cb_cancel_data *match)
840{
841	struct io_wq_work_node *node, *prev;
842	struct io_wq_work *work;
843	int i;
844
845retry:
846	raw_spin_lock(&wqe->lock);
847	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
848		struct io_wqe_acct *acct = io_get_acct(wqe, i == 0);
849
850		wq_list_for_each(node, prev, &acct->work_list) {
851			work = container_of(node, struct io_wq_work, list);
852			if (!match->fn(work, match->data))
853				continue;
854			io_wqe_remove_pending(wqe, work, prev);
855			raw_spin_unlock(&wqe->lock);
856			io_run_cancel(work, wqe);
857			match->nr_pending++;
858			if (!match->cancel_all)
859				return;
860
861			/* not safe to continue after unlock */
862			goto retry;
863		}
864	}
865	raw_spin_unlock(&wqe->lock);
866}
867
868static void io_wqe_cancel_running_work(struct io_wqe *wqe,
869				       struct io_cb_cancel_data *match)
870{
871	rcu_read_lock();
872	io_wq_for_each_worker(wqe, io_wq_worker_cancel, match);
873	rcu_read_unlock();
874}
875
876enum io_wq_cancel io_wq_cancel_cb(struct io_wq *wq, work_cancel_fn *cancel,
877				  void *data, bool cancel_all)
878{
879	struct io_cb_cancel_data match = {
880		.fn		= cancel,
881		.data		= data,
882		.cancel_all	= cancel_all,
883	};
884	int node;
885
886	/*
887	 * First check pending list, if we're lucky we can just remove it
888	 * from there. CANCEL_OK means that the work is returned as-new,
889	 * no completion will be posted for it.
890	 */
891	for_each_node(node) {
892		struct io_wqe *wqe = wq->wqes[node];
893
894		io_wqe_cancel_pending_work(wqe, &match);
895		if (match.nr_pending && !match.cancel_all)
896			return IO_WQ_CANCEL_OK;
897	}
898
899	/*
900	 * Now check if a free (going busy) or busy worker has the work
901	 * currently running. If we find it there, we'll return CANCEL_RUNNING
902	 * as an indication that we attempt to signal cancellation. The
903	 * completion will run normally in this case.
904	 */
905	for_each_node(node) {
906		struct io_wqe *wqe = wq->wqes[node];
907
908		io_wqe_cancel_running_work(wqe, &match);
909		if (match.nr_running && !match.cancel_all)
910			return IO_WQ_CANCEL_RUNNING;
911	}
912
913	if (match.nr_running)
914		return IO_WQ_CANCEL_RUNNING;
915	if (match.nr_pending)
916		return IO_WQ_CANCEL_OK;
917	return IO_WQ_CANCEL_NOTFOUND;
918}
919
920static int io_wqe_hash_wake(struct wait_queue_entry *wait, unsigned mode,
921			    int sync, void *key)
922{
923	struct io_wqe *wqe = container_of(wait, struct io_wqe, wait);
924	int i;
925
926	list_del_init(&wait->entry);
927
928	rcu_read_lock();
929	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
930		struct io_wqe_acct *acct = &wqe->acct[i];
931
932		if (test_and_clear_bit(IO_ACCT_STALLED_BIT, &acct->flags))
933			io_wqe_activate_free_worker(wqe, acct);
934	}
935	rcu_read_unlock();
936	return 1;
937}
938
939struct io_wq *io_wq_create(unsigned bounded, struct io_wq_data *data)
940{
941	int ret, node, i;
942	struct io_wq *wq;
943
944	if (WARN_ON_ONCE(!data->free_work || !data->do_work))
945		return ERR_PTR(-EINVAL);
946	if (WARN_ON_ONCE(!bounded))
947		return ERR_PTR(-EINVAL);
948
949	wq = kzalloc(struct_size(wq, wqes, nr_node_ids), GFP_KERNEL);
950	if (!wq)
951		return ERR_PTR(-ENOMEM);
952	ret = cpuhp_state_add_instance_nocalls(io_wq_online, &wq->cpuhp_node);
953	if (ret)
954		goto err_wq;
955
956	refcount_inc(&data->hash->refs);
957	wq->hash = data->hash;
958	wq->free_work = data->free_work;
959	wq->do_work = data->do_work;
960
961	ret = -ENOMEM;
962	for_each_node(node) {
963		struct io_wqe *wqe;
964		int alloc_node = node;
965
966		if (!node_online(alloc_node))
967			alloc_node = NUMA_NO_NODE;
968		wqe = kzalloc_node(sizeof(struct io_wqe), GFP_KERNEL, alloc_node);
969		if (!wqe)
970			goto err;
971		if (!alloc_cpumask_var(&wqe->cpu_mask, GFP_KERNEL))
972			goto err;
973		cpumask_copy(wqe->cpu_mask, cpumask_of_node(node));
974		wq->wqes[node] = wqe;
975		wqe->node = alloc_node;
976		wqe->acct[IO_WQ_ACCT_BOUND].max_workers = bounded;
977		wqe->acct[IO_WQ_ACCT_UNBOUND].max_workers =
978					task_rlimit(current, RLIMIT_NPROC);
979		INIT_LIST_HEAD(&wqe->wait.entry);
980		wqe->wait.func = io_wqe_hash_wake;
981		for (i = 0; i < IO_WQ_ACCT_NR; i++) {
982			struct io_wqe_acct *acct = &wqe->acct[i];
983
984			acct->index = i;
985			atomic_set(&acct->nr_running, 0);
986			INIT_WQ_LIST(&acct->work_list);
987		}
988		wqe->wq = wq;
989		raw_spin_lock_init(&wqe->lock);
990		INIT_HLIST_NULLS_HEAD(&wqe->free_list, 0);
991		INIT_LIST_HEAD(&wqe->all_list);
992	}
993
994	wq->task = get_task_struct(data->task);
995	atomic_set(&wq->worker_refs, 1);
996	init_completion(&wq->worker_done);
997	return wq;
998err:
999	io_wq_put_hash(data->hash);
1000	cpuhp_state_remove_instance_nocalls(io_wq_online, &wq->cpuhp_node);
1001	for_each_node(node) {
1002		if (!wq->wqes[node])
1003			continue;
1004		free_cpumask_var(wq->wqes[node]->cpu_mask);
1005		kfree(wq->wqes[node]);
1006	}
1007err_wq:
1008	kfree(wq);
1009	return ERR_PTR(ret);
1010}
1011
1012static bool io_task_work_match(struct callback_head *cb, void *data)
1013{
1014	struct io_worker *worker;
1015
1016	if (cb->func != create_worker_cb)
1017		return false;
1018	worker = container_of(cb, struct io_worker, create_work);
1019	return worker->wqe->wq == data;
1020}
1021
1022void io_wq_exit_start(struct io_wq *wq)
1023{
1024	set_bit(IO_WQ_BIT_EXIT, &wq->state);
1025}
1026
1027static void io_wq_exit_workers(struct io_wq *wq)
1028{
1029	struct callback_head *cb;
1030	int node;
1031
1032	if (!wq->task)
1033		return;
1034
1035	while ((cb = task_work_cancel_match(wq->task, io_task_work_match, wq)) != NULL) {
1036		struct io_worker *worker;
1037
1038		worker = container_of(cb, struct io_worker, create_work);
1039		atomic_dec(&worker->wqe->acct[worker->create_index].nr_running);
1040		io_worker_ref_put(wq);
1041		clear_bit_unlock(0, &worker->create_state);
1042		io_worker_release(worker);
1043	}
1044
1045	rcu_read_lock();
1046	for_each_node(node) {
1047		struct io_wqe *wqe = wq->wqes[node];
1048
1049		io_wq_for_each_worker(wqe, io_wq_worker_wake, NULL);
1050	}
1051	rcu_read_unlock();
1052	io_worker_ref_put(wq);
1053	wait_for_completion(&wq->worker_done);
1054
1055	for_each_node(node) {
1056		spin_lock_irq(&wq->hash->wait.lock);
1057		list_del_init(&wq->wqes[node]->wait.entry);
1058		spin_unlock_irq(&wq->hash->wait.lock);
1059	}
1060	put_task_struct(wq->task);
1061	wq->task = NULL;
1062}
1063
1064static void io_wq_destroy(struct io_wq *wq)
1065{
1066	int node;
1067
1068	cpuhp_state_remove_instance_nocalls(io_wq_online, &wq->cpuhp_node);
1069
1070	for_each_node(node) {
1071		struct io_wqe *wqe = wq->wqes[node];
1072		struct io_cb_cancel_data match = {
1073			.fn		= io_wq_work_match_all,
1074			.cancel_all	= true,
1075		};
1076		io_wqe_cancel_pending_work(wqe, &match);
1077		free_cpumask_var(wqe->cpu_mask);
1078		kfree(wqe);
1079	}
1080	io_wq_put_hash(wq->hash);
1081	kfree(wq);
1082}
1083
1084void io_wq_put_and_exit(struct io_wq *wq)
1085{
1086	WARN_ON_ONCE(!test_bit(IO_WQ_BIT_EXIT, &wq->state));
1087
1088	io_wq_exit_workers(wq);
1089	io_wq_destroy(wq);
1090}
1091
1092struct online_data {
1093	unsigned int cpu;
1094	bool online;
1095};
1096
1097static bool io_wq_worker_affinity(struct io_worker *worker, void *data)
1098{
1099	struct online_data *od = data;
1100
1101	if (od->online)
1102		cpumask_set_cpu(od->cpu, worker->wqe->cpu_mask);
1103	else
1104		cpumask_clear_cpu(od->cpu, worker->wqe->cpu_mask);
1105	return false;
1106}
1107
1108static int __io_wq_cpu_online(struct io_wq *wq, unsigned int cpu, bool online)
1109{
1110	struct online_data od = {
1111		.cpu = cpu,
1112		.online = online
1113	};
1114	int i;
1115
1116	rcu_read_lock();
1117	for_each_node(i)
1118		io_wq_for_each_worker(wq->wqes[i], io_wq_worker_affinity, &od);
1119	rcu_read_unlock();
1120	return 0;
1121}
1122
1123static int io_wq_cpu_online(unsigned int cpu, struct hlist_node *node)
1124{
1125	struct io_wq *wq = hlist_entry_safe(node, struct io_wq, cpuhp_node);
1126
1127	return __io_wq_cpu_online(wq, cpu, true);
1128}
1129
1130static int io_wq_cpu_offline(unsigned int cpu, struct hlist_node *node)
1131{
1132	struct io_wq *wq = hlist_entry_safe(node, struct io_wq, cpuhp_node);
1133
1134	return __io_wq_cpu_online(wq, cpu, false);
1135}
1136
1137int io_wq_cpu_affinity(struct io_wq *wq, cpumask_var_t mask)
1138{
1139	int i;
1140
1141	rcu_read_lock();
1142	for_each_node(i) {
1143		struct io_wqe *wqe = wq->wqes[i];
1144
1145		if (mask)
1146			cpumask_copy(wqe->cpu_mask, mask);
1147		else
1148			cpumask_copy(wqe->cpu_mask, cpumask_of_node(i));
1149	}
1150	rcu_read_unlock();
1151	return 0;
1152}
1153
1154/*
1155 * Set max number of unbounded workers, returns old value. If new_count is 0,
1156 * then just return the old value.
1157 */
1158int io_wq_max_workers(struct io_wq *wq, int *new_count)
1159{
1160	int i, node, prev = 0;
1161
1162	for (i = 0; i < 2; i++) {
1163		if (new_count[i] > task_rlimit(current, RLIMIT_NPROC))
1164			new_count[i] = task_rlimit(current, RLIMIT_NPROC);
1165	}
1166
1167	rcu_read_lock();
1168	for_each_node(node) {
1169		struct io_wqe_acct *acct;
1170
1171		for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1172			acct = &wq->wqes[node]->acct[i];
1173			prev = max_t(int, acct->max_workers, prev);
1174			if (new_count[i])
1175				acct->max_workers = new_count[i];
1176			new_count[i] = prev;
1177		}
1178	}
1179	rcu_read_unlock();
1180	return 0;
1181}
1182
1183static __init int io_wq_init(void)
1184{
1185	int ret;
1186
1187	ret = cpuhp_setup_state_multi(CPUHP_AP_ONLINE_DYN, "io-wq/online",
1188					io_wq_cpu_online, io_wq_cpu_offline);
1189	if (ret < 0)
1190		return ret;
1191	io_wq_online = ret;
1192	return 0;
1193}
1194subsys_initcall(io_wq_init);
1195