io-wq.c

   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/cpuset.h>
  17#include <linux/task_work.h>
  18#include <linux/audit.h>
  19#include <linux/mmu_context.h>
  20#include <uapi/linux/io_uring.h>
  21
  22#include "io-wq.h"
  23#include "slist.h"
  24#include "io_uring.h"
  25
  26#define WORKER_IDLE_TIMEOUT	(5 * HZ)
  27#define WORKER_INIT_LIMIT	3
  28
  29enum {
  30	IO_WORKER_F_UP		= 0,	/* up and active */
  31	IO_WORKER_F_RUNNING	= 1,	/* account as running */
  32	IO_WORKER_F_FREE	= 2,	/* worker on free list */
  33};
  34
  35enum {
  36	IO_WQ_BIT_EXIT		= 0,	/* wq exiting */
  37	IO_WQ_BIT_EXIT_ON_IDLE	= 1,	/* allow all workers to exit on idle */
  38};
  39
  40enum {
  41	IO_ACCT_STALLED_BIT	= 0,	/* stalled on hash */
  42};
  43
  44/*
  45 * One for each thread in a wq pool
  46 */
  47struct io_worker {
  48	refcount_t ref;
  49	unsigned long flags;
  50	struct hlist_nulls_node nulls_node;
  51	struct list_head all_list;
  52	struct task_struct *task;
  53	struct io_wq *wq;
  54	struct io_wq_acct *acct;
  55
  56	struct io_wq_work *cur_work;
  57	raw_spinlock_t lock;
  58
  59	struct completion ref_done;
  60
  61	unsigned long create_state;
  62	struct callback_head create_work;
  63	int init_retries;
  64
  65	union {
  66		struct rcu_head rcu;
  67		struct delayed_work work;
  68	};
  69};
  70
  71#if BITS_PER_LONG == 64
  72#define IO_WQ_HASH_ORDER	6
  73#else
  74#define IO_WQ_HASH_ORDER	5
  75#endif
  76
  77#define IO_WQ_NR_HASH_BUCKETS	(1u << IO_WQ_HASH_ORDER)
  78
  79struct io_wq_acct {
  80	/**
  81	 * Protects access to the worker lists.
  82	 */
  83	raw_spinlock_t workers_lock;
  84
  85	unsigned nr_workers;
  86	unsigned max_workers;
  87	atomic_t nr_running;
  88
  89	/**
  90	 * The list of free workers.  Protected by #workers_lock
  91	 * (write) and RCU (read).
  92	 */
  93	struct hlist_nulls_head free_list;
  94
  95	/**
  96	 * The list of all workers.  Protected by #workers_lock
  97	 * (write) and RCU (read).
  98	 */
  99	struct list_head all_list;
 100
 101	raw_spinlock_t lock;
 102	struct io_wq_work_list work_list;
 103	unsigned long flags;
 104};
 105
 106enum {
 107	IO_WQ_ACCT_BOUND,
 108	IO_WQ_ACCT_UNBOUND,
 109	IO_WQ_ACCT_NR,
 110};
 111
 112/*
 113 * Per io_wq state
 114  */
 115struct io_wq {
 116	unsigned long state;
 117
 118	struct io_wq_hash *hash;
 119
 120	atomic_t worker_refs;
 121	struct completion worker_done;
 122
 123	struct hlist_node cpuhp_node;
 124
 125	struct task_struct *task;
 126
 127	struct io_wq_acct acct[IO_WQ_ACCT_NR];
 128
 129	struct wait_queue_entry wait;
 130
 131	struct io_wq_work *hash_tail[IO_WQ_NR_HASH_BUCKETS];
 132
 133	cpumask_var_t cpu_mask;
 134};
 135
 136static enum cpuhp_state io_wq_online;
 137
 138struct io_cb_cancel_data {
 139	work_cancel_fn *fn;
 140	void *data;
 141	int nr_running;
 142	int nr_pending;
 143	bool cancel_all;
 144};
 145
 146static bool create_io_worker(struct io_wq *wq, struct io_wq_acct *acct);
 147static void io_wq_dec_running(struct io_worker *worker);
 148static bool io_acct_cancel_pending_work(struct io_wq *wq,
 149					struct io_wq_acct *acct,
 150					struct io_cb_cancel_data *match);
 151static void create_worker_cb(struct callback_head *cb);
 152static void io_wq_cancel_tw_create(struct io_wq *wq);
 153
 154static inline unsigned int __io_get_work_hash(unsigned int work_flags)
 155{
 156	return work_flags >> IO_WQ_HASH_SHIFT;
 157}
 158
 159static inline unsigned int io_get_work_hash(struct io_wq_work *work)
 160{
 161	return __io_get_work_hash(atomic_read(&work->flags));
 162}
 163
 164static bool io_worker_get(struct io_worker *worker)
 165{
 166	return refcount_inc_not_zero(&worker->ref);
 167}
 168
 169static void io_worker_release(struct io_worker *worker)
 170{
 171	if (refcount_dec_and_test(&worker->ref))
 172		complete(&worker->ref_done);
 173}
 174
 175static inline struct io_wq_acct *io_get_acct(struct io_wq *wq, bool bound)
 176{
 177	return &wq->acct[bound ? IO_WQ_ACCT_BOUND : IO_WQ_ACCT_UNBOUND];
 178}
 179
 180static inline struct io_wq_acct *io_work_get_acct(struct io_wq *wq,
 181						  unsigned int work_flags)
 182{
 183	return io_get_acct(wq, !(work_flags & IO_WQ_WORK_UNBOUND));
 184}
 185
 186static inline struct io_wq_acct *io_wq_get_acct(struct io_worker *worker)
 187{
 188	return worker->acct;
 189}
 190
 191static void io_worker_ref_put(struct io_wq *wq)
 192{
 193	if (atomic_dec_and_test(&wq->worker_refs))
 194		complete(&wq->worker_done);
 195}
 196
 197bool io_wq_worker_stopped(void)
 198{
 199	struct io_worker *worker = current->worker_private;
 200
 201	if (WARN_ON_ONCE(!io_wq_current_is_worker()))
 202		return true;
 203
 204	return test_bit(IO_WQ_BIT_EXIT, &worker->wq->state);
 205}
 206
 207static void io_worker_cancel_cb(struct io_worker *worker)
 208{
 209	struct io_wq_acct *acct = io_wq_get_acct(worker);
 210	struct io_wq *wq = worker->wq;
 211
 212	atomic_dec(&acct->nr_running);
 213	raw_spin_lock(&acct->workers_lock);
 214	acct->nr_workers--;
 215	raw_spin_unlock(&acct->workers_lock);
 216	io_worker_ref_put(wq);
 217	clear_bit_unlock(0, &worker->create_state);
 218	io_worker_release(worker);
 219}
 220
 221static bool io_task_worker_match(struct callback_head *cb, void *data)
 222{
 223	struct io_worker *worker;
 224
 225	if (cb->func != create_worker_cb)
 226		return false;
 227	worker = container_of(cb, struct io_worker, create_work);
 228	return worker == data;
 229}
 230
 231static void io_worker_exit(struct io_worker *worker)
 232{
 233	struct io_wq *wq = worker->wq;
 234	struct io_wq_acct *acct = io_wq_get_acct(worker);
 235
 236	while (1) {
 237		struct callback_head *cb = task_work_cancel_match(wq->task,
 238						io_task_worker_match, worker);
 239
 240		if (!cb)
 241			break;
 242		io_worker_cancel_cb(worker);
 243	}
 244
 245	io_worker_release(worker);
 246	wait_for_completion(&worker->ref_done);
 247
 248	raw_spin_lock(&acct->workers_lock);
 249	if (test_bit(IO_WORKER_F_FREE, &worker->flags))
 250		hlist_nulls_del_rcu(&worker->nulls_node);
 251	list_del_rcu(&worker->all_list);
 252	raw_spin_unlock(&acct->workers_lock);
 253	io_wq_dec_running(worker);
 254	/*
 255	 * this worker is a goner, clear ->worker_private to avoid any
 256	 * inc/dec running calls that could happen as part of exit from
 257	 * touching 'worker'.
 258	 */
 259	current->worker_private = NULL;
 260
 261	kfree_rcu(worker, rcu);
 262	io_worker_ref_put(wq);
 263	do_exit(0);
 264}
 265
 266static inline bool __io_acct_run_queue(struct io_wq_acct *acct)
 267{
 268	return !test_bit(IO_ACCT_STALLED_BIT, &acct->flags) &&
 269		!wq_list_empty(&acct->work_list);
 270}
 271
 272/*
 273 * If there's work to do, returns true with acct->lock acquired. If not,
 274 * returns false with no lock held.
 275 */
 276static inline bool io_acct_run_queue(struct io_wq_acct *acct)
 277	__acquires(&acct->lock)
 278{
 279	raw_spin_lock(&acct->lock);
 280	if (__io_acct_run_queue(acct))
 281		return true;
 282
 283	raw_spin_unlock(&acct->lock);
 284	return false;
 285}
 286
 287/*
 288 * Check head of free list for an available worker. If one isn't available,
 289 * caller must create one.
 290 */
 291static bool io_acct_activate_free_worker(struct io_wq_acct *acct)
 292	__must_hold(RCU)
 293{
 294	struct hlist_nulls_node *n;
 295	struct io_worker *worker;
 296
 297	/*
 298	 * Iterate free_list and see if we can find an idle worker to
 299	 * activate. If a given worker is on the free_list but in the process
 300	 * of exiting, keep trying.
 301	 */
 302	hlist_nulls_for_each_entry_rcu(worker, n, &acct->free_list, nulls_node) {
 303		if (!io_worker_get(worker))
 304			continue;
 305		/*
 306		 * If the worker is already running, it's either already
 307		 * starting work or finishing work. In either case, if it does
 308		 * to go sleep, we'll kick off a new task for this work anyway.
 309		 */
 310		wake_up_process(worker->task);
 311		io_worker_release(worker);
 312		return true;
 313	}
 314
 315	return false;
 316}
 317
 318/*
 319 * We need a worker. If we find a free one, we're good. If not, and we're
 320 * below the max number of workers, create one.
 321 */
 322static bool io_wq_create_worker(struct io_wq *wq, struct io_wq_acct *acct)
 323{
 324	/*
 325	 * Most likely an attempt to queue unbounded work on an io_wq that
 326	 * wasn't setup with any unbounded workers.
 327	 */
 328	if (unlikely(!acct->max_workers))
 329		pr_warn_once("io-wq is not configured for unbound workers");
 330
 331	raw_spin_lock(&acct->workers_lock);
 332	if (acct->nr_workers >= acct->max_workers) {
 333		raw_spin_unlock(&acct->workers_lock);
 334		return true;
 335	}
 336	acct->nr_workers++;
 337	raw_spin_unlock(&acct->workers_lock);
 338	atomic_inc(&acct->nr_running);
 339	atomic_inc(&wq->worker_refs);
 340	return create_io_worker(wq, acct);
 341}
 342
 343static void io_wq_inc_running(struct io_worker *worker)
 344{
 345	struct io_wq_acct *acct = io_wq_get_acct(worker);
 346
 347	atomic_inc(&acct->nr_running);
 348}
 349
 350static void create_worker_cb(struct callback_head *cb)
 351{
 352	struct io_worker *worker;
 353	struct io_wq *wq;
 354
 355	struct io_wq_acct *acct;
 356	bool activated_free_worker, do_create = false;
 357
 358	worker = container_of(cb, struct io_worker, create_work);
 359	wq = worker->wq;
 360	acct = worker->acct;
 361
 362	rcu_read_lock();
 363	activated_free_worker = io_acct_activate_free_worker(acct);
 364	rcu_read_unlock();
 365	if (activated_free_worker)
 366		goto no_need_create;
 367
 368	raw_spin_lock(&acct->workers_lock);
 369
 370	if (acct->nr_workers < acct->max_workers) {
 371		acct->nr_workers++;
 372		do_create = true;
 373	}
 374	raw_spin_unlock(&acct->workers_lock);
 375	if (do_create) {
 376		create_io_worker(wq, acct);
 377	} else {
 378no_need_create:
 379		atomic_dec(&acct->nr_running);
 380		io_worker_ref_put(wq);
 381	}
 382	clear_bit_unlock(0, &worker->create_state);
 383	io_worker_release(worker);
 384}
 385
 386static bool io_queue_worker_create(struct io_worker *worker,
 387				   struct io_wq_acct *acct,
 388				   task_work_func_t func)
 389{
 390	struct io_wq *wq = worker->wq;
 391
 392	/* raced with exit, just ignore create call */
 393	if (test_bit(IO_WQ_BIT_EXIT, &wq->state))
 394		goto fail;
 395	if (!io_worker_get(worker))
 396		goto fail;
 397	/*
 398	 * create_state manages ownership of create_work/index. We should
 399	 * only need one entry per worker, as the worker going to sleep
 400	 * will trigger the condition, and waking will clear it once it
 401	 * runs the task_work.
 402	 */
 403	if (test_bit(0, &worker->create_state) ||
 404	    test_and_set_bit_lock(0, &worker->create_state))
 405		goto fail_release;
 406
 407	atomic_inc(&wq->worker_refs);
 408	init_task_work(&worker->create_work, func);
 409	if (!task_work_add(wq->task, &worker->create_work, TWA_SIGNAL)) {
 410		/*
 411		 * EXIT may have been set after checking it above, check after
 412		 * adding the task_work and remove any creation item if it is
 413		 * now set. wq exit does that too, but we can have added this
 414		 * work item after we canceled in io_wq_exit_workers().
 415		 */
 416		if (test_bit(IO_WQ_BIT_EXIT, &wq->state))
 417			io_wq_cancel_tw_create(wq);
 418		io_worker_ref_put(wq);
 419		return true;
 420	}
 421	io_worker_ref_put(wq);
 422	clear_bit_unlock(0, &worker->create_state);
 423fail_release:
 424	io_worker_release(worker);
 425fail:
 426	atomic_dec(&acct->nr_running);
 427	io_worker_ref_put(wq);
 428	return false;
 429}
 430
 431/* Defer if current and next work are both hashed to the same chain */
 432static bool io_wq_hash_defer(struct io_wq_work *work, struct io_wq_acct *acct)
 433{
 434	unsigned int hash, work_flags;
 435	struct io_wq_work *next;
 436
 437	lockdep_assert_held(&acct->lock);
 438
 439	work_flags = atomic_read(&work->flags);
 440	if (!__io_wq_is_hashed(work_flags))
 441		return false;
 442
 443	/* should not happen, io_acct_run_queue() said we had work */
 444	if (wq_list_empty(&acct->work_list))
 445		return true;
 446
 447	hash = __io_get_work_hash(work_flags);
 448	next = container_of(acct->work_list.first, struct io_wq_work, list);
 449	work_flags = atomic_read(&next->flags);
 450	if (!__io_wq_is_hashed(work_flags))
 451		return false;
 452	return hash == __io_get_work_hash(work_flags);
 453}
 454
 455static void io_wq_dec_running(struct io_worker *worker)
 456{
 457	struct io_wq_acct *acct = io_wq_get_acct(worker);
 458	struct io_wq *wq = worker->wq;
 459
 460	if (!test_bit(IO_WORKER_F_UP, &worker->flags))
 461		return;
 462
 463	if (!atomic_dec_and_test(&acct->nr_running))
 464		return;
 465	if (!worker->cur_work)
 466		return;
 467	if (!io_acct_run_queue(acct))
 468		return;
 469	if (io_wq_hash_defer(worker->cur_work, acct)) {
 470		raw_spin_unlock(&acct->lock);
 471		return;
 472	}
 473
 474	raw_spin_unlock(&acct->lock);
 475	atomic_inc(&acct->nr_running);
 476	atomic_inc(&wq->worker_refs);
 477	io_queue_worker_create(worker, acct, create_worker_cb);
 478}
 479
 480/*
 481 * Worker will start processing some work. Move it to the busy list, if
 482 * it's currently on the freelist
 483 */
 484static void __io_worker_busy(struct io_wq_acct *acct, struct io_worker *worker)
 485{
 486	if (test_bit(IO_WORKER_F_FREE, &worker->flags)) {
 487		clear_bit(IO_WORKER_F_FREE, &worker->flags);
 488		raw_spin_lock(&acct->workers_lock);
 489		hlist_nulls_del_init_rcu(&worker->nulls_node);
 490		raw_spin_unlock(&acct->workers_lock);
 491	}
 492}
 493
 494/*
 495 * No work, worker going to sleep. Move to freelist.
 496 */
 497static void __io_worker_idle(struct io_wq_acct *acct, struct io_worker *worker)
 498	__must_hold(acct->workers_lock)
 499{
 500	if (!test_bit(IO_WORKER_F_FREE, &worker->flags)) {
 501		set_bit(IO_WORKER_F_FREE, &worker->flags);
 502		hlist_nulls_add_head_rcu(&worker->nulls_node, &acct->free_list);
 503	}
 504}
 505
 506static bool io_wait_on_hash(struct io_wq *wq, unsigned int hash)
 507{
 508	bool ret = false;
 509
 510	spin_lock_irq(&wq->hash->wait.lock);
 511	if (list_empty(&wq->wait.entry)) {
 512		__add_wait_queue(&wq->hash->wait, &wq->wait);
 513		if (!test_bit(hash, &wq->hash->map)) {
 514			__set_current_state(TASK_RUNNING);
 515			list_del_init(&wq->wait.entry);
 516			ret = true;
 517		}
 518	}
 519	spin_unlock_irq(&wq->hash->wait.lock);
 520	return ret;
 521}
 522
 523static struct io_wq_work *io_get_next_work(struct io_wq_acct *acct,
 524					   struct io_wq *wq)
 525	__must_hold(acct->lock)
 526{
 527	struct io_wq_work_node *node, *prev;
 528	struct io_wq_work *work, *tail;
 529	unsigned int stall_hash = -1U;
 530
 531	wq_list_for_each(node, prev, &acct->work_list) {
 532		unsigned int work_flags;
 533		unsigned int hash;
 534
 535		work = container_of(node, struct io_wq_work, list);
 536
 537		/* not hashed, can run anytime */
 538		work_flags = atomic_read(&work->flags);
 539		if (!__io_wq_is_hashed(work_flags)) {
 540			wq_list_del(&acct->work_list, node, prev);
 541			return work;
 542		}
 543
 544		hash = __io_get_work_hash(work_flags);
 545		/* all items with this hash lie in [work, tail] */
 546		tail = wq->hash_tail[hash];
 547
 548		/* hashed, can run if not already running */
 549		if (!test_and_set_bit(hash, &wq->hash->map)) {
 550			wq->hash_tail[hash] = NULL;
 551			wq_list_cut(&acct->work_list, &tail->list, prev);
 552			return work;
 553		}
 554		if (stall_hash == -1U)
 555			stall_hash = hash;
 556		/* fast forward to a next hash, for-each will fix up @prev */
 557		node = &tail->list;
 558	}
 559
 560	if (stall_hash != -1U) {
 561		bool unstalled;
 562
 563		/*
 564		 * Set this before dropping the lock to avoid racing with new
 565		 * work being added and clearing the stalled bit.
 566		 */
 567		set_bit(IO_ACCT_STALLED_BIT, &acct->flags);
 568		raw_spin_unlock(&acct->lock);
 569		unstalled = io_wait_on_hash(wq, stall_hash);
 570		raw_spin_lock(&acct->lock);
 571		if (unstalled) {
 572			clear_bit(IO_ACCT_STALLED_BIT, &acct->flags);
 573			if (wq_has_sleeper(&wq->hash->wait))
 574				wake_up(&wq->hash->wait);
 575		}
 576	}
 577
 578	return NULL;
 579}
 580
 581static void io_assign_current_work(struct io_worker *worker,
 582				   struct io_wq_work *work)
 583{
 584	if (work) {
 585		io_run_task_work();
 586		cond_resched();
 587	}
 588
 589	raw_spin_lock(&worker->lock);
 590	worker->cur_work = work;
 591	raw_spin_unlock(&worker->lock);
 592}
 593
 594/*
 595 * Called with acct->lock held, drops it before returning
 596 */
 597static void io_worker_handle_work(struct io_wq_acct *acct,
 598				  struct io_worker *worker)
 599	__releases(&acct->lock)
 600{
 601	struct io_wq *wq = worker->wq;
 602
 603	do {
 604		bool do_kill = test_bit(IO_WQ_BIT_EXIT, &wq->state);
 605		struct io_wq_work *work;
 606
 607		/*
 608		 * If we got some work, mark us as busy. If we didn't, but
 609		 * the list isn't empty, it means we stalled on hashed work.
 610		 * Mark us stalled so we don't keep looking for work when we
 611		 * can't make progress, any work completion or insertion will
 612		 * clear the stalled flag.
 613		 */
 614		work = io_get_next_work(acct, wq);
 615		if (work) {
 616			/*
 617			 * Make sure cancelation can find this, even before
 618			 * it becomes the active work. That avoids a window
 619			 * where the work has been removed from our general
 620			 * work list, but isn't yet discoverable as the
 621			 * current work item for this worker.
 622			 */
 623			raw_spin_lock(&worker->lock);
 624			worker->cur_work = work;
 625			raw_spin_unlock(&worker->lock);
 626		}
 627
 628		raw_spin_unlock(&acct->lock);
 629
 630		if (!work)
 631			break;
 632
 633		__io_worker_busy(acct, worker);
 634
 635		io_assign_current_work(worker, work);
 636		__set_current_state(TASK_RUNNING);
 637
 638		/* handle a whole dependent link */
 639		do {
 640			struct io_wq_work *next_hashed, *linked;
 641			unsigned int work_flags = atomic_read(&work->flags);
 642			unsigned int hash = __io_wq_is_hashed(work_flags)
 643				? __io_get_work_hash(work_flags)
 644				: -1U;
 645
 646			next_hashed = wq_next_work(work);
 647
 648			if (do_kill &&
 649			    (work_flags & IO_WQ_WORK_UNBOUND))
 650				atomic_or(IO_WQ_WORK_CANCEL, &work->flags);
 651			io_wq_submit_work(work);
 652			io_assign_current_work(worker, NULL);
 653
 654			linked = io_wq_free_work(work);
 655			work = next_hashed;
 656			if (!work && linked && !io_wq_is_hashed(linked)) {
 657				work = linked;
 658				linked = NULL;
 659			}
 660			io_assign_current_work(worker, work);
 661			if (linked)
 662				io_wq_enqueue(wq, linked);
 663
 664			if (hash != -1U && !next_hashed) {
 665				/* serialize hash clear with wake_up() */
 666				spin_lock_irq(&wq->hash->wait.lock);
 667				clear_bit(hash, &wq->hash->map);
 668				clear_bit(IO_ACCT_STALLED_BIT, &acct->flags);
 669				spin_unlock_irq(&wq->hash->wait.lock);
 670				if (wq_has_sleeper(&wq->hash->wait))
 671					wake_up(&wq->hash->wait);
 672			}
 673		} while (work);
 674
 675		if (!__io_acct_run_queue(acct))
 676			break;
 677		raw_spin_lock(&acct->lock);
 678	} while (1);
 679}
 680
 681static int io_wq_worker(void *data)
 682{
 683	struct io_worker *worker = data;
 684	struct io_wq_acct *acct = io_wq_get_acct(worker);
 685	struct io_wq *wq = worker->wq;
 686	bool exit_mask = false, last_timeout = false;
 687	char buf[TASK_COMM_LEN] = {};
 688
 689	set_mask_bits(&worker->flags, 0,
 690		      BIT(IO_WORKER_F_UP) | BIT(IO_WORKER_F_RUNNING));
 691
 692	snprintf(buf, sizeof(buf), "iou-wrk-%d", wq->task->pid);
 693	set_task_comm(current, buf);
 694
 695	while (!test_bit(IO_WQ_BIT_EXIT, &wq->state)) {
 696		long ret;
 697
 698		set_current_state(TASK_INTERRUPTIBLE);
 699
 700		/*
 701		 * If we have work to do, io_acct_run_queue() returns with
 702		 * the acct->lock held. If not, it will drop it.
 703		 */
 704		while (io_acct_run_queue(acct))
 705			io_worker_handle_work(acct, worker);
 706
 707		raw_spin_lock(&acct->workers_lock);
 708		/*
 709		 * Last sleep timed out. Exit if we're not the last worker,
 710		 * or if someone modified our affinity. If wq is marked
 711		 * idle-exit, drop the worker as well. This is used to avoid
 712		 * keeping io-wq workers around for tasks that no longer have
 713		 * any active io_uring instances.
 714		 */
 715		if ((last_timeout && (exit_mask || acct->nr_workers > 1)) ||
 716		    test_bit(IO_WQ_BIT_EXIT_ON_IDLE, &wq->state)) {
 717			acct->nr_workers--;
 718			raw_spin_unlock(&acct->workers_lock);
 719			__set_current_state(TASK_RUNNING);
 720			break;
 721		}
 722		last_timeout = false;
 723		__io_worker_idle(acct, worker);
 724		raw_spin_unlock(&acct->workers_lock);
 725		if (io_run_task_work())
 726			continue;
 727		ret = schedule_timeout(WORKER_IDLE_TIMEOUT);
 728		if (signal_pending(current)) {
 729			struct ksignal ksig;
 730
 731			if (!get_signal(&ksig))
 732				continue;
 733			break;
 734		}
 735		if (!ret) {
 736			last_timeout = true;
 737			exit_mask = !cpumask_test_cpu(raw_smp_processor_id(),
 738							wq->cpu_mask);
 739		}
 740	}
 741
 742	if (test_bit(IO_WQ_BIT_EXIT, &wq->state) && io_acct_run_queue(acct))
 743		io_worker_handle_work(acct, worker);
 744
 745	io_worker_exit(worker);
 746	return 0;
 747}
 748
 749/*
 750 * Called when a worker is scheduled in. Mark us as currently running.
 751 */
 752void io_wq_worker_running(struct task_struct *tsk)
 753{
 754	struct io_worker *worker = tsk->worker_private;
 755
 756	if (!worker)
 757		return;
 758	if (!test_bit(IO_WORKER_F_UP, &worker->flags))
 759		return;
 760	if (test_bit(IO_WORKER_F_RUNNING, &worker->flags))
 761		return;
 762	set_bit(IO_WORKER_F_RUNNING, &worker->flags);
 763	io_wq_inc_running(worker);
 764}
 765
 766/*
 767 * Called when worker is going to sleep. If there are no workers currently
 768 * running and we have work pending, wake up a free one or create a new one.
 769 */
 770void io_wq_worker_sleeping(struct task_struct *tsk)
 771{
 772	struct io_worker *worker = tsk->worker_private;
 773
 774	if (!worker)
 775		return;
 776	if (!test_bit(IO_WORKER_F_UP, &worker->flags))
 777		return;
 778	if (!test_bit(IO_WORKER_F_RUNNING, &worker->flags))
 779		return;
 780
 781	clear_bit(IO_WORKER_F_RUNNING, &worker->flags);
 782	io_wq_dec_running(worker);
 783}
 784
 785static void io_init_new_worker(struct io_wq *wq, struct io_wq_acct *acct, struct io_worker *worker,
 786			       struct task_struct *tsk)
 787{
 788	tsk->worker_private = worker;
 789	worker->task = tsk;
 790	set_cpus_allowed_ptr(tsk, wq->cpu_mask);
 791
 792	raw_spin_lock(&acct->workers_lock);
 793	hlist_nulls_add_head_rcu(&worker->nulls_node, &acct->free_list);
 794	list_add_tail_rcu(&worker->all_list, &acct->all_list);
 795	set_bit(IO_WORKER_F_FREE, &worker->flags);
 796	raw_spin_unlock(&acct->workers_lock);
 797	wake_up_new_task(tsk);
 798}
 799
 800static bool io_wq_work_match_all(struct io_wq_work *work, void *data)
 801{
 802	return true;
 803}
 804
 805static inline bool io_should_retry_thread(struct io_worker *worker, long err)
 806{
 807	/*
 808	 * Prevent perpetual task_work retry, if the task (or its group) is
 809	 * exiting.
 810	 */
 811	if (fatal_signal_pending(current))
 812		return false;
 813
 814	worker->init_retries++;
 815	switch (err) {
 816	case -EAGAIN:
 817		return worker->init_retries <= WORKER_INIT_LIMIT;
 818	/* Analogous to a fork() syscall, always retry on a restartable error */
 819	case -ERESTARTSYS:
 820	case -ERESTARTNOINTR:
 821	case -ERESTARTNOHAND:
 822		return true;
 823	default:
 824		return false;
 825	}
 826}
 827
 828static void queue_create_worker_retry(struct io_worker *worker)
 829{
 830	/*
 831	 * We only bother retrying because there's a chance that the
 832	 * failure to create a worker is due to some temporary condition
 833	 * in the forking task (e.g. outstanding signal); give the task
 834	 * some time to clear that condition.
 835	 */
 836	schedule_delayed_work(&worker->work,
 837			      msecs_to_jiffies(worker->init_retries * 5));
 838}
 839
 840static void create_worker_cont(struct callback_head *cb)
 841{
 842	struct io_worker *worker;
 843	struct task_struct *tsk;
 844	struct io_wq *wq;
 845	struct io_wq_acct *acct;
 846
 847	worker = container_of(cb, struct io_worker, create_work);
 848	clear_bit_unlock(0, &worker->create_state);
 849	wq = worker->wq;
 850	acct = io_wq_get_acct(worker);
 851	tsk = create_io_thread(io_wq_worker, worker, NUMA_NO_NODE);
 852	if (!IS_ERR(tsk)) {
 853		io_init_new_worker(wq, acct, worker, tsk);
 854		io_worker_release(worker);
 855		return;
 856	} else if (!io_should_retry_thread(worker, PTR_ERR(tsk))) {
 857		atomic_dec(&acct->nr_running);
 858		raw_spin_lock(&acct->workers_lock);
 859		acct->nr_workers--;
 860		if (!acct->nr_workers) {
 861			struct io_cb_cancel_data match = {
 862				.fn		= io_wq_work_match_all,
 863				.cancel_all	= true,
 864			};
 865
 866			raw_spin_unlock(&acct->workers_lock);
 867			while (io_acct_cancel_pending_work(wq, acct, &match))
 868				;
 869		} else {
 870			raw_spin_unlock(&acct->workers_lock);
 871		}
 872		io_worker_ref_put(wq);
 873		kfree(worker);
 874		return;
 875	}
 876
 877	/* re-create attempts grab a new worker ref, drop the existing one */
 878	io_worker_release(worker);
 879	queue_create_worker_retry(worker);
 880}
 881
 882static void io_workqueue_create(struct work_struct *work)
 883{
 884	struct io_worker *worker = container_of(work, struct io_worker,
 885						work.work);
 886	struct io_wq_acct *acct = io_wq_get_acct(worker);
 887
 888	if (!io_queue_worker_create(worker, acct, create_worker_cont))
 889		kfree(worker);
 890}
 891
 892static bool create_io_worker(struct io_wq *wq, struct io_wq_acct *acct)
 893{
 894	struct io_worker *worker;
 895	struct task_struct *tsk;
 896
 897	__set_current_state(TASK_RUNNING);
 898
 899	worker = kzalloc(sizeof(*worker), GFP_KERNEL);
 900	if (!worker) {
 901fail:
 902		atomic_dec(&acct->nr_running);
 903		raw_spin_lock(&acct->workers_lock);
 904		acct->nr_workers--;
 905		raw_spin_unlock(&acct->workers_lock);
 906		io_worker_ref_put(wq);
 907		return false;
 908	}
 909
 910	refcount_set(&worker->ref, 1);
 911	worker->wq = wq;
 912	worker->acct = acct;
 913	raw_spin_lock_init(&worker->lock);
 914	init_completion(&worker->ref_done);
 915
 916	tsk = create_io_thread(io_wq_worker, worker, NUMA_NO_NODE);
 917	if (!IS_ERR(tsk)) {
 918		io_init_new_worker(wq, acct, worker, tsk);
 919	} else if (!io_should_retry_thread(worker, PTR_ERR(tsk))) {
 920		kfree(worker);
 921		goto fail;
 922	} else {
 923		INIT_DELAYED_WORK(&worker->work, io_workqueue_create);
 924		queue_create_worker_retry(worker);
 925	}
 926
 927	return true;
 928}
 929
 930/*
 931 * Iterate the passed in list and call the specific function for each
 932 * worker that isn't exiting
 933 */
 934static bool io_acct_for_each_worker(struct io_wq_acct *acct,
 935				    bool (*func)(struct io_worker *, void *),
 936				    void *data)
 937{
 938	struct io_worker *worker;
 939	bool ret = false;
 940
 941	list_for_each_entry_rcu(worker, &acct->all_list, all_list) {
 942		if (io_worker_get(worker)) {
 943			/* no task if node is/was offline */
 944			if (worker->task)
 945				ret = func(worker, data);
 946			io_worker_release(worker);
 947			if (ret)
 948				break;
 949		}
 950	}
 951
 952	return ret;
 953}
 954
 955static void io_wq_for_each_worker(struct io_wq *wq,
 956				  bool (*func)(struct io_worker *, void *),
 957				  void *data)
 958{
 959	for (int i = 0; i < IO_WQ_ACCT_NR; i++)
 960		if (io_acct_for_each_worker(&wq->acct[i], func, data))
 961			break;
 962}
 963
 964static bool io_wq_worker_wake(struct io_worker *worker, void *data)
 965{
 966	__set_notify_signal(worker->task);
 967	wake_up_process(worker->task);
 968	return false;
 969}
 970
 971void io_wq_set_exit_on_idle(struct io_wq *wq, bool enable)
 972{
 973	if (!wq->task)
 974		return;
 975
 976	if (!enable) {
 977		clear_bit(IO_WQ_BIT_EXIT_ON_IDLE, &wq->state);
 978		return;
 979	}
 980
 981	if (test_and_set_bit(IO_WQ_BIT_EXIT_ON_IDLE, &wq->state))
 982		return;
 983
 984	rcu_read_lock();
 985	io_wq_for_each_worker(wq, io_wq_worker_wake, NULL);
 986	rcu_read_unlock();
 987}
 988
 989static void io_run_cancel(struct io_wq_work *work, struct io_wq *wq)
 990{
 991	do {
 992		atomic_or(IO_WQ_WORK_CANCEL, &work->flags);
 993		io_wq_submit_work(work);
 994		work = io_wq_free_work(work);
 995	} while (work);
 996}
 997
 998static void io_wq_insert_work(struct io_wq *wq, struct io_wq_acct *acct,
 999			      struct io_wq_work *work, unsigned int work_flags)
1000{
1001	unsigned int hash;
1002	struct io_wq_work *tail;
1003
1004	if (!__io_wq_is_hashed(work_flags)) {
1005append:
1006		wq_list_add_tail(&work->list, &acct->work_list);
1007		return;
1008	}
1009
1010	hash = __io_get_work_hash(work_flags);
1011	tail = wq->hash_tail[hash];
1012	wq->hash_tail[hash] = work;
1013	if (!tail)
1014		goto append;
1015
1016	wq_list_add_after(&work->list, &tail->list, &acct->work_list);
1017}
1018
1019static bool io_wq_work_match_item(struct io_wq_work *work, void *data)
1020{
1021	return work == data;
1022}
1023
1024void io_wq_enqueue(struct io_wq *wq, struct io_wq_work *work)
1025{
1026	unsigned int work_flags = atomic_read(&work->flags);
1027	struct io_wq_acct *acct = io_work_get_acct(wq, work_flags);
1028	struct io_cb_cancel_data match = {
1029		.fn		= io_wq_work_match_item,
1030		.data		= work,
1031		.cancel_all	= false,
1032	};
1033	bool do_create;
1034
1035	/*
1036	 * If io-wq is exiting for this task, or if the request has explicitly
1037	 * been marked as one that should not get executed, cancel it here.
1038	 */
1039	if (test_bit(IO_WQ_BIT_EXIT, &wq->state) ||
1040	    (work_flags & IO_WQ_WORK_CANCEL)) {
1041		io_run_cancel(work, wq);
1042		return;
1043	}
1044
1045	raw_spin_lock(&acct->lock);
1046	io_wq_insert_work(wq, acct, work, work_flags);
1047	clear_bit(IO_ACCT_STALLED_BIT, &acct->flags);
1048	raw_spin_unlock(&acct->lock);
1049
1050	rcu_read_lock();
1051	do_create = !io_acct_activate_free_worker(acct);
1052	rcu_read_unlock();
1053
1054	if (do_create && ((work_flags & IO_WQ_WORK_CONCURRENT) ||
1055	    !atomic_read(&acct->nr_running))) {
1056		bool did_create;
1057
1058		did_create = io_wq_create_worker(wq, acct);
1059		if (likely(did_create))
1060			return;
1061
1062		raw_spin_lock(&acct->workers_lock);
1063		if (acct->nr_workers) {
1064			raw_spin_unlock(&acct->workers_lock);
1065			return;
1066		}
1067		raw_spin_unlock(&acct->workers_lock);
1068
1069		/* fatal condition, failed to create the first worker */
1070		io_acct_cancel_pending_work(wq, acct, &match);
1071	}
1072}
1073
1074/*
1075 * Work items that hash to the same value will not be done in parallel.
1076 * Used to limit concurrent writes, generally hashed by inode.
1077 */
1078void io_wq_hash_work(struct io_wq_work *work, void *val)
1079{
1080	unsigned int bit;
1081
1082	bit = hash_ptr(val, IO_WQ_HASH_ORDER);
1083	atomic_or(IO_WQ_WORK_HASHED | (bit << IO_WQ_HASH_SHIFT), &work->flags);
1084}
1085
1086static bool __io_wq_worker_cancel(struct io_worker *worker,
1087				  struct io_cb_cancel_data *match,
1088				  struct io_wq_work *work)
1089{
1090	if (work && match->fn(work, match->data)) {
1091		atomic_or(IO_WQ_WORK_CANCEL, &work->flags);
1092		__set_notify_signal(worker->task);
1093		return true;
1094	}
1095
1096	return false;
1097}
1098
1099static bool io_wq_worker_cancel(struct io_worker *worker, void *data)
1100{
1101	struct io_cb_cancel_data *match = data;
1102
1103	/*
1104	 * Hold the lock to avoid ->cur_work going out of scope, caller
1105	 * may dereference the passed in work.
1106	 */
1107	raw_spin_lock(&worker->lock);
1108	if (__io_wq_worker_cancel(worker, match, worker->cur_work))
1109		match->nr_running++;
1110	raw_spin_unlock(&worker->lock);
1111
1112	return match->nr_running && !match->cancel_all;
1113}
1114
1115static inline void io_wq_remove_pending(struct io_wq *wq,
1116					struct io_wq_acct *acct,
1117					 struct io_wq_work *work,
1118					 struct io_wq_work_node *prev)
1119{
1120	unsigned int hash = io_get_work_hash(work);
1121	struct io_wq_work *prev_work = NULL;
1122
1123	if (io_wq_is_hashed(work) && work == wq->hash_tail[hash]) {
1124		if (prev)
1125			prev_work = container_of(prev, struct io_wq_work, list);
1126		if (prev_work && io_get_work_hash(prev_work) == hash)
1127			wq->hash_tail[hash] = prev_work;
1128		else
1129			wq->hash_tail[hash] = NULL;
1130	}
1131	wq_list_del(&acct->work_list, &work->list, prev);
1132}
1133
1134static bool io_acct_cancel_pending_work(struct io_wq *wq,
1135					struct io_wq_acct *acct,
1136					struct io_cb_cancel_data *match)
1137{
1138	struct io_wq_work_node *node, *prev;
1139	struct io_wq_work *work;
1140
1141	raw_spin_lock(&acct->lock);
1142	wq_list_for_each(node, prev, &acct->work_list) {
1143		work = container_of(node, struct io_wq_work, list);
1144		if (!match->fn(work, match->data))
1145			continue;
1146		io_wq_remove_pending(wq, acct, work, prev);
1147		raw_spin_unlock(&acct->lock);
1148		io_run_cancel(work, wq);
1149		match->nr_pending++;
1150		/* not safe to continue after unlock */
1151		return true;
1152	}
1153	raw_spin_unlock(&acct->lock);
1154
1155	return false;
1156}
1157
1158static void io_wq_cancel_pending_work(struct io_wq *wq,
1159				      struct io_cb_cancel_data *match)
1160{
1161	int i;
1162retry:
1163	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1164		struct io_wq_acct *acct = io_get_acct(wq, i == 0);
1165
1166		if (io_acct_cancel_pending_work(wq, acct, match)) {
1167			if (match->cancel_all)
1168				goto retry;
1169			break;
1170		}
1171	}
1172}
1173
1174static void io_acct_cancel_running_work(struct io_wq_acct *acct,
1175					struct io_cb_cancel_data *match)
1176{
1177	raw_spin_lock(&acct->workers_lock);
1178	io_acct_for_each_worker(acct, io_wq_worker_cancel, match);
1179	raw_spin_unlock(&acct->workers_lock);
1180}
1181
1182static void io_wq_cancel_running_work(struct io_wq *wq,
1183				       struct io_cb_cancel_data *match)
1184{
1185	rcu_read_lock();
1186
1187	for (int i = 0; i < IO_WQ_ACCT_NR; i++)
1188		io_acct_cancel_running_work(&wq->acct[i], match);
1189
1190	rcu_read_unlock();
1191}
1192
1193enum io_wq_cancel io_wq_cancel_cb(struct io_wq *wq, work_cancel_fn *cancel,
1194				  void *data, bool cancel_all)
1195{
1196	struct io_cb_cancel_data match = {
1197		.fn		= cancel,
1198		.data		= data,
1199		.cancel_all	= cancel_all,
1200	};
1201
1202	/*
1203	 * First check pending list, if we're lucky we can just remove it
1204	 * from there. CANCEL_OK means that the work is returned as-new,
1205	 * no completion will be posted for it.
1206	 *
1207	 * Then check if a free (going busy) or busy worker has the work
1208	 * currently running. If we find it there, we'll return CANCEL_RUNNING
1209	 * as an indication that we attempt to signal cancellation. The
1210	 * completion will run normally in this case.
1211	 *
1212	 * Do both of these while holding the acct->workers_lock, to ensure that
1213	 * we'll find a work item regardless of state.
1214	 */
1215	io_wq_cancel_pending_work(wq, &match);
1216	if (match.nr_pending && !match.cancel_all)
1217		return IO_WQ_CANCEL_OK;
1218
1219	io_wq_cancel_running_work(wq, &match);
1220	if (match.nr_running && !match.cancel_all)
1221		return IO_WQ_CANCEL_RUNNING;
1222
1223	if (match.nr_running)
1224		return IO_WQ_CANCEL_RUNNING;
1225	if (match.nr_pending)
1226		return IO_WQ_CANCEL_OK;
1227	return IO_WQ_CANCEL_NOTFOUND;
1228}
1229
1230static int io_wq_hash_wake(struct wait_queue_entry *wait, unsigned mode,
1231			    int sync, void *key)
1232{
1233	struct io_wq *wq = container_of(wait, struct io_wq, wait);
1234	int i;
1235
1236	list_del_init(&wait->entry);
1237
1238	rcu_read_lock();
1239	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1240		struct io_wq_acct *acct = &wq->acct[i];
1241
1242		if (test_and_clear_bit(IO_ACCT_STALLED_BIT, &acct->flags))
1243			io_acct_activate_free_worker(acct);
1244	}
1245	rcu_read_unlock();
1246	return 1;
1247}
1248
1249struct io_wq *io_wq_create(unsigned bounded, struct io_wq_data *data)
1250{
1251	int ret, i;
1252	struct io_wq *wq;
1253
1254	if (WARN_ON_ONCE(!bounded))
1255		return ERR_PTR(-EINVAL);
1256
1257	wq = kzalloc(sizeof(struct io_wq), GFP_KERNEL);
1258	if (!wq)
1259		return ERR_PTR(-ENOMEM);
1260
1261	refcount_inc(&data->hash->refs);
1262	wq->hash = data->hash;
1263
1264	ret = -ENOMEM;
1265
1266	if (!alloc_cpumask_var(&wq->cpu_mask, GFP_KERNEL))
1267		goto err;
1268	cpuset_cpus_allowed(data->task, wq->cpu_mask);
1269	wq->acct[IO_WQ_ACCT_BOUND].max_workers = bounded;
1270	wq->acct[IO_WQ_ACCT_UNBOUND].max_workers =
1271				task_rlimit(current, RLIMIT_NPROC);
1272	INIT_LIST_HEAD(&wq->wait.entry);
1273	wq->wait.func = io_wq_hash_wake;
1274	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1275		struct io_wq_acct *acct = &wq->acct[i];
1276
1277		atomic_set(&acct->nr_running, 0);
1278
1279		raw_spin_lock_init(&acct->workers_lock);
1280		INIT_HLIST_NULLS_HEAD(&acct->free_list, 0);
1281		INIT_LIST_HEAD(&acct->all_list);
1282
1283		INIT_WQ_LIST(&acct->work_list);
1284		raw_spin_lock_init(&acct->lock);
1285	}
1286
1287	wq->task = get_task_struct(data->task);
1288	atomic_set(&wq->worker_refs, 1);
1289	init_completion(&wq->worker_done);
1290	ret = cpuhp_state_add_instance_nocalls(io_wq_online, &wq->cpuhp_node);
1291	if (ret) {
1292		put_task_struct(wq->task);
1293		goto err;
1294	}
1295
1296	return wq;
1297err:
1298	io_wq_put_hash(data->hash);
1299	free_cpumask_var(wq->cpu_mask);
1300	kfree(wq);
1301	return ERR_PTR(ret);
1302}
1303
1304static bool io_task_work_match(struct callback_head *cb, void *data)
1305{
1306	struct io_worker *worker;
1307
1308	if (cb->func != create_worker_cb && cb->func != create_worker_cont)
1309		return false;
1310	worker = container_of(cb, struct io_worker, create_work);
1311	return worker->wq == data;
1312}
1313
1314void io_wq_exit_start(struct io_wq *wq)
1315{
1316	set_bit(IO_WQ_BIT_EXIT, &wq->state);
1317}
1318
1319static void io_wq_cancel_tw_create(struct io_wq *wq)
1320{
1321	struct callback_head *cb;
1322
1323	while ((cb = task_work_cancel_match(wq->task, io_task_work_match, wq)) != NULL) {
1324		struct io_worker *worker;
1325
1326		worker = container_of(cb, struct io_worker, create_work);
1327		io_worker_cancel_cb(worker);
1328		/*
1329		 * Only the worker continuation helper has worker allocated and
1330		 * hence needs freeing.
1331		 */
1332		if (cb->func == create_worker_cont)
1333			kfree(worker);
1334	}
1335}
1336
1337static void io_wq_exit_workers(struct io_wq *wq)
1338{
1339	if (!wq->task)
1340		return;
1341
1342	io_wq_cancel_tw_create(wq);
1343
1344	rcu_read_lock();
1345	io_wq_for_each_worker(wq, io_wq_worker_wake, NULL);
1346	rcu_read_unlock();
1347	io_worker_ref_put(wq);
1348	wait_for_completion(&wq->worker_done);
1349
1350	spin_lock_irq(&wq->hash->wait.lock);
1351	list_del_init(&wq->wait.entry);
1352	spin_unlock_irq(&wq->hash->wait.lock);
1353
1354	put_task_struct(wq->task);
1355	wq->task = NULL;
1356}
1357
1358static void io_wq_destroy(struct io_wq *wq)
1359{
1360	struct io_cb_cancel_data match = {
1361		.fn		= io_wq_work_match_all,
1362		.cancel_all	= true,
1363	};
1364
1365	cpuhp_state_remove_instance_nocalls(io_wq_online, &wq->cpuhp_node);
1366	io_wq_cancel_pending_work(wq, &match);
1367	free_cpumask_var(wq->cpu_mask);
1368	io_wq_put_hash(wq->hash);
1369	kfree(wq);
1370}
1371
1372void io_wq_put_and_exit(struct io_wq *wq)
1373{
1374	WARN_ON_ONCE(!test_bit(IO_WQ_BIT_EXIT, &wq->state));
1375
1376	io_wq_exit_workers(wq);
1377	io_wq_destroy(wq);
1378}
1379
1380struct online_data {
1381	unsigned int cpu;
1382	bool online;
1383};
1384
1385static bool io_wq_worker_affinity(struct io_worker *worker, void *data)
1386{
1387	struct online_data *od = data;
1388
1389	if (od->online)
1390		cpumask_set_cpu(od->cpu, worker->wq->cpu_mask);
1391	else
1392		cpumask_clear_cpu(od->cpu, worker->wq->cpu_mask);
1393	return false;
1394}
1395
1396static int __io_wq_cpu_online(struct io_wq *wq, unsigned int cpu, bool online)
1397{
1398	struct online_data od = {
1399		.cpu = cpu,
1400		.online = online
1401	};
1402
1403	rcu_read_lock();
1404	io_wq_for_each_worker(wq, io_wq_worker_affinity, &od);
1405	rcu_read_unlock();
1406	return 0;
1407}
1408
1409static int io_wq_cpu_online(unsigned int cpu, struct hlist_node *node)
1410{
1411	struct io_wq *wq = hlist_entry_safe(node, struct io_wq, cpuhp_node);
1412
1413	return __io_wq_cpu_online(wq, cpu, true);
1414}
1415
1416static int io_wq_cpu_offline(unsigned int cpu, struct hlist_node *node)
1417{
1418	struct io_wq *wq = hlist_entry_safe(node, struct io_wq, cpuhp_node);
1419
1420	return __io_wq_cpu_online(wq, cpu, false);
1421}
1422
1423int io_wq_cpu_affinity(struct io_uring_task *tctx, cpumask_var_t mask)
1424{
1425	cpumask_var_t allowed_mask;
1426	int ret = 0;
1427
1428	if (!tctx || !tctx->io_wq)
1429		return -EINVAL;
1430
1431	if (!alloc_cpumask_var(&allowed_mask, GFP_KERNEL))
1432		return -ENOMEM;
1433
1434	rcu_read_lock();
1435	cpuset_cpus_allowed(tctx->io_wq->task, allowed_mask);
1436	if (mask) {
1437		if (cpumask_subset(mask, allowed_mask))
1438			cpumask_copy(tctx->io_wq->cpu_mask, mask);
1439		else
1440			ret = -EINVAL;
1441	} else {
1442		cpumask_copy(tctx->io_wq->cpu_mask, allowed_mask);
1443	}
1444	rcu_read_unlock();
1445
1446	free_cpumask_var(allowed_mask);
1447	return ret;
1448}
1449
1450/*
1451 * Set max number of unbounded workers, returns old value. If new_count is 0,
1452 * then just return the old value.
1453 */
1454int io_wq_max_workers(struct io_wq *wq, int *new_count)
1455{
1456	struct io_wq_acct *acct;
1457	int prev[IO_WQ_ACCT_NR];
1458	int i;
1459
1460	BUILD_BUG_ON((int) IO_WQ_ACCT_BOUND   != (int) IO_WQ_BOUND);
1461	BUILD_BUG_ON((int) IO_WQ_ACCT_UNBOUND != (int) IO_WQ_UNBOUND);
1462	BUILD_BUG_ON((int) IO_WQ_ACCT_NR      != 2);
1463
1464	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1465		if (new_count[i] > task_rlimit(current, RLIMIT_NPROC))
1466			new_count[i] = task_rlimit(current, RLIMIT_NPROC);
1467	}
1468
1469	for (i = 0; i < IO_WQ_ACCT_NR; i++)
1470		prev[i] = 0;
1471
1472	rcu_read_lock();
1473
1474	for (i = 0; i < IO_WQ_ACCT_NR; i++) {
1475		acct = &wq->acct[i];
1476		raw_spin_lock(&acct->workers_lock);
1477		prev[i] = max_t(int, acct->max_workers, prev[i]);
1478		if (new_count[i])
1479			acct->max_workers = new_count[i];
1480		raw_spin_unlock(&acct->workers_lock);
1481	}
1482	rcu_read_unlock();
1483
1484	for (i = 0; i < IO_WQ_ACCT_NR; i++)
1485		new_count[i] = prev[i];
1486
1487	return 0;
1488}
1489
1490static __init int io_wq_init(void)
1491{
1492	int ret;
1493
1494	ret = cpuhp_setup_state_multi(CPUHP_AP_ONLINE_DYN, "io-wq/online",
1495					io_wq_cpu_online, io_wq_cpu_offline);
1496	if (ret < 0)
1497		return ret;
1498	io_wq_online = ret;
1499	return 0;
1500}
1501subsys_initcall(io_wq_init);