fs/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/mm.h>
  13 #include <linux/mmu_context.h>
  14 #include <linux/sched/mm.h>
  15 #include <linux/percpu.h>
  16 #include <linux/slab.h>
  17 #include <linux/kthread.h>
  18 #include <linux/rculist_nulls.h>
  19
  20 #include "io-wq.h"
  21
  22 #define WORKER_IDLE_TIMEOUT     (5 * HZ)
  23
  24 enum {
  25         IO_WORKER_F_UP          = 1,    /* up and active */
  26         IO_WORKER_F_RUNNING     = 2,    /* account as running */
  27         IO_WORKER_F_FREE        = 4,    /* worker on free list */
  28         IO_WORKER_F_EXITING     = 8,    /* worker exiting */
  29         IO_WORKER_F_FIXED       = 16,   /* static idle worker */
  30         IO_WORKER_F_BOUND       = 32,   /* is doing bounded work */
  31 };
  32
  33 enum {
  34         IO_WQ_BIT_EXIT          = 0,    /* wq exiting */
  35         IO_WQ_BIT_CANCEL        = 1,    /* cancel work on list */
  36 };
  37
  38 enum {
  39         IO_WQE_FLAG_STALLED     = 1,    /* stalled on hash */
  40 };
  41
  42 /*
  43  * One for each thread in a wqe pool
  44  */
  45 struct io_worker {
  46         refcount_t ref;
  47         unsigned flags;
  48         struct hlist_nulls_node nulls_node;
  49         struct list_head all_list;
  50         struct task_struct *task;
  51         wait_queue_head_t wait;
  52         struct io_wqe *wqe;
  53
  54         struct io_wq_work *cur_work;
  55         spinlock_t lock;
  56
  57         struct rcu_head rcu;
  58         struct mm_struct *mm;
  59         struct files_struct *restore_files;
  60 };
  61
  62 #if BITS_PER_LONG == 64
  63 #define IO_WQ_HASH_ORDER        6
  64 #else
  65 #define IO_WQ_HASH_ORDER        5
  66 #endif
  67
  68 struct io_wqe_acct {
  69         unsigned nr_workers;
  70         unsigned max_workers;
  71         atomic_t nr_running;
  72 };
  73
  74 enum {
  75         IO_WQ_ACCT_BOUND,
  76         IO_WQ_ACCT_UNBOUND,
  77 };
  78
  79 /*
  80  * Per-node worker thread pool
  81  */
  82 struct io_wqe {
  83         struct {
  84                 spinlock_t lock;
  85                 struct list_head work_list;
  86                 unsigned long hash_map;
  87                 unsigned flags;
  88         } ____cacheline_aligned_in_smp;
  89
  90         int node;
  91         struct io_wqe_acct acct[2];
  92
  93         struct hlist_nulls_head free_list;
  94         struct hlist_nulls_head busy_list;
  95         struct list_head all_list;
  96
  97         struct io_wq *wq;
  98 };
  99
 100 /*
 101  * Per io_wq state
 102   */
 103 struct io_wq {
 104         struct io_wqe **wqes;
 105         unsigned long state;
 106         unsigned nr_wqes;
 107
 108         get_work_fn *get_work;
 109         put_work_fn *put_work;
 110
 111         struct task_struct *manager;
 112         struct user_struct *user;
 113         struct mm_struct *mm;
 114         refcount_t refs;
 115         struct completion done;
 116 };
 117
 118 static bool io_worker_get(struct io_worker *worker)
 119 {
 120         return refcount_inc_not_zero(&worker->ref);
 121 }
 122
 123 static void io_worker_release(struct io_worker *worker)
 124 {
 125         if (refcount_dec_and_test(&worker->ref))
 126                 wake_up_process(worker->task);
 127 }
 128
 129 /*
 130  * Note: drops the wqe->lock if returning true! The caller must re-acquire
 131  * the lock in that case. Some callers need to restart handling if this
 132  * happens, so we can't just re-acquire the lock on behalf of the caller.
 133  */
 134 static bool __io_worker_unuse(struct io_wqe *wqe, struct io_worker *worker)
 135 {
 136         bool dropped_lock = false;
 137
 138         if (current->files != worker->restore_files) {
 139                 __acquire(&wqe->lock);
 140                 spin_unlock_irq(&wqe->lock);
 141                 dropped_lock = true;
 142
 143                 task_lock(current);
 144                 current->files = worker->restore_files;
 145                 task_unlock(current);
 146         }
 147
 148         /*
 149          * If we have an active mm, we need to drop the wq lock before unusing
 150          * it. If we do, return true and let the caller retry the idle loop.
 151          */
 152         if (worker->mm) {
 153                 if (!dropped_lock) {
 154                         __acquire(&wqe->lock);
 155                         spin_unlock_irq(&wqe->lock);
 156                         dropped_lock = true;
 157                 }
 158                 __set_current_state(TASK_RUNNING);
 159                 set_fs(KERNEL_DS);
 160                 unuse_mm(worker->mm);
 161                 mmput(worker->mm);
 162                 worker->mm = NULL;
 163         }
 164
 165         return dropped_lock;
 166 }
 167
 168 static inline struct io_wqe_acct *io_work_get_acct(struct io_wqe *wqe,
 169                                                    struct io_wq_work *work)
 170 {
 171         if (work->flags & IO_WQ_WORK_UNBOUND)
 172                 return &wqe->acct[IO_WQ_ACCT_UNBOUND];
 173
 174         return &wqe->acct[IO_WQ_ACCT_BOUND];
 175 }
 176
 177 static inline struct io_wqe_acct *io_wqe_get_acct(struct io_wqe *wqe,
 178                                                   struct io_worker *worker)
 179 {
 180         if (worker->flags & IO_WORKER_F_BOUND)
 181                 return &wqe->acct[IO_WQ_ACCT_BOUND];
 182
 183         return &wqe->acct[IO_WQ_ACCT_UNBOUND];
 184 }
 185
 186 static void io_worker_exit(struct io_worker *worker)
 187 {
 188         struct io_wqe *wqe = worker->wqe;
 189         struct io_wqe_acct *acct = io_wqe_get_acct(wqe, worker);
 190         unsigned nr_workers;
 191
 192         /*
 193          * If we're not at zero, someone else is holding a brief reference
 194          * to the worker. Wait for that to go away.
 195          */
 196         set_current_state(TASK_INTERRUPTIBLE);
 197         if (!refcount_dec_and_test(&worker->ref))
 198                 schedule();
 199         __set_current_state(TASK_RUNNING);
 200
 201         preempt_disable();
 202         current->flags &= ~PF_IO_WORKER;
 203         if (worker->flags & IO_WORKER_F_RUNNING)
 204                 atomic_dec(&acct->nr_running);
 205         if (!(worker->flags & IO_WORKER_F_BOUND))
 206                 atomic_dec(&wqe->wq->user->processes);
 207         worker->flags = 0;
 208         preempt_enable();
 209
 210         spin_lock_irq(&wqe->lock);
 211         hlist_nulls_del_rcu(&worker->nulls_node);
 212         list_del_rcu(&worker->all_list);
 213         if (__io_worker_unuse(wqe, worker)) {
 214                 __release(&wqe->lock);
 215                 spin_lock_irq(&wqe->lock);
 216         }
 217         acct->nr_workers--;
 218         nr_workers = wqe->acct[IO_WQ_ACCT_BOUND].nr_workers +
 219                         wqe->acct[IO_WQ_ACCT_UNBOUND].nr_workers;
 220         spin_unlock_irq(&wqe->lock);
 221
 222         /* all workers gone, wq exit can proceed */
 223         if (!nr_workers && refcount_dec_and_test(&wqe->wq->refs))
 224                 complete(&wqe->wq->done);
 225
 226         kfree_rcu(worker, rcu);
 227 }
 228
 229 static inline bool io_wqe_run_queue(struct io_wqe *wqe)
 230         __must_hold(wqe->lock)
 231 {
 232         if (!list_empty(&wqe->work_list) && !(wqe->flags & IO_WQE_FLAG_STALLED))
 233                 return true;
 234         return false;
 235 }
 236
 237 /*
 238  * Check head of free list for an available worker. If one isn't available,
 239  * caller must wake up the wq manager to create one.
 240  */
 241 static bool io_wqe_activate_free_worker(struct io_wqe *wqe)
 242         __must_hold(RCU)
 243 {
 244         struct hlist_nulls_node *n;
 245         struct io_worker *worker;
 246
 247         n = rcu_dereference(hlist_nulls_first_rcu(&wqe->free_list));
 248         if (is_a_nulls(n))
 249                 return false;
 250
 251         worker = hlist_nulls_entry(n, struct io_worker, nulls_node);
 252         if (io_worker_get(worker)) {
 253                 wake_up(&worker->wait);
 254                 io_worker_release(worker);
 255                 return true;
 256         }
 257
 258         return false;
 259 }
 260
 261 /*
 262  * We need a worker. If we find a free one, we're good. If not, and we're
 263  * below the max number of workers, wake up the manager to create one.
 264  */
 265 static void io_wqe_wake_worker(struct io_wqe *wqe, struct io_wqe_acct *acct)
 266 {
 267         bool ret;
 268
 269         /*
 270          * Most likely an attempt to queue unbounded work on an io_wq that
 271          * wasn't setup with any unbounded workers.
 272          */
 273         WARN_ON_ONCE(!acct->max_workers);
 274
 275         rcu_read_lock();
 276         ret = io_wqe_activate_free_worker(wqe);
 277         rcu_read_unlock();
 278
 279         if (!ret && acct->nr_workers < acct->max_workers)
 280                 wake_up_process(wqe->wq->manager);
 281 }
 282
 283 static void io_wqe_inc_running(struct io_wqe *wqe, struct io_worker *worker)
 284 {
 285         struct io_wqe_acct *acct = io_wqe_get_acct(wqe, worker);
 286
 287         atomic_inc(&acct->nr_running);
 288 }
 289
 290 static void io_wqe_dec_running(struct io_wqe *wqe, struct io_worker *worker)
 291         __must_hold(wqe->lock)
 292 {
 293         struct io_wqe_acct *acct = io_wqe_get_acct(wqe, worker);
 294
 295         if (atomic_dec_and_test(&acct->nr_running) && io_wqe_run_queue(wqe))
 296                 io_wqe_wake_worker(wqe, acct);
 297 }
 298
 299 static void io_worker_start(struct io_wqe *wqe, struct io_worker *worker)
 300 {
 301         allow_kernel_signal(SIGINT);
 302
 303         current->flags |= PF_IO_WORKER;
 304
 305         worker->flags |= (IO_WORKER_F_UP | IO_WORKER_F_RUNNING);
 306         worker->restore_files = current->files;
 307         io_wqe_inc_running(wqe, worker);
 308 }
 309
 310 /*
 311  * Worker will start processing some work. Move it to the busy list, if
 312  * it's currently on the freelist
 313  */
 314 static void __io_worker_busy(struct io_wqe *wqe, struct io_worker *worker,
 315                              struct io_wq_work *work)
 316         __must_hold(wqe->lock)
 317 {
 318         bool worker_bound, work_bound;
 319
 320         if (worker->flags & IO_WORKER_F_FREE) {
 321                 worker->flags &= ~IO_WORKER_F_FREE;
 322                 hlist_nulls_del_init_rcu(&worker->nulls_node);
 323                 hlist_nulls_add_head_rcu(&worker->nulls_node, &wqe->busy_list);
 324         }
 325
 326         /*
 327          * If worker is moving from bound to unbound (or vice versa), then
 328          * ensure we update the running accounting.
 329          */
 330          worker_bound = (worker->flags & IO_WORKER_F_BOUND) != 0;
 331          work_bound = (work->flags & IO_WQ_WORK_UNBOUND) == 0;
 332          if (worker_bound != work_bound) {
 333                 io_wqe_dec_running(wqe, worker);
 334                 if (work_bound) {
 335                         worker->flags |= IO_WORKER_F_BOUND;
 336                         wqe->acct[IO_WQ_ACCT_UNBOUND].nr_workers--;
 337                         wqe->acct[IO_WQ_ACCT_BOUND].nr_workers++;
 338                         atomic_dec(&wqe->wq->user->processes);
 339                 } else {
 340                         worker->flags &= ~IO_WORKER_F_BOUND;
 341                         wqe->acct[IO_WQ_ACCT_UNBOUND].nr_workers++;
 342                         wqe->acct[IO_WQ_ACCT_BOUND].nr_workers--;
 343                         atomic_inc(&wqe->wq->user->processes);
 344                 }
 345                 io_wqe_inc_running(wqe, worker);
 346          }
 347 }
 348
 349 /*
 350  * No work, worker going to sleep. Move to freelist, and unuse mm if we
 351  * have one attached. Dropping the mm may potentially sleep, so we drop
 352  * the lock in that case and return success. Since the caller has to
 353  * retry the loop in that case (we changed task state), we don't regrab
 354  * the lock if we return success.
 355  */
 356 static bool __io_worker_idle(struct io_wqe *wqe, struct io_worker *worker)
 357         __must_hold(wqe->lock)
 358 {
 359         if (!(worker->flags & IO_WORKER_F_FREE)) {
 360                 worker->flags |= IO_WORKER_F_FREE;
 361                 hlist_nulls_del_init_rcu(&worker->nulls_node);
 362                 hlist_nulls_add_head_rcu(&worker->nulls_node, &wqe->free_list);
 363         }
 364
 365         return __io_worker_unuse(wqe, worker);
 366 }
 367
 368 static struct io_wq_work *io_get_next_work(struct io_wqe *wqe, unsigned *hash)
 369         __must_hold(wqe->lock)
 370 {
 371         struct io_wq_work *work;
 372
 373         list_for_each_entry(work, &wqe->work_list, list) {
 374                 /* not hashed, can run anytime */
 375                 if (!(work->flags & IO_WQ_WORK_HASHED)) {
 376                         list_del(&work->list);
 377                         return work;
 378                 }
 379
 380                 /* hashed, can run if not already running */
 381                 *hash = work->flags >> IO_WQ_HASH_SHIFT;
 382                 if (!(wqe->hash_map & BIT_ULL(*hash))) {
 383                         wqe->hash_map |= BIT_ULL(*hash);
 384                         list_del(&work->list);
 385                         return work;
 386                 }
 387         }
 388
 389         return NULL;
 390 }
 391
 392 static void io_worker_handle_work(struct io_worker *worker)
 393         __releases(wqe->lock)
 394 {
 395         struct io_wq_work *work, *old_work = NULL, *put_work = NULL;
 396         struct io_wqe *wqe = worker->wqe;
 397         struct io_wq *wq = wqe->wq;
 398
 399         do {
 400                 unsigned hash = -1U;
 401
 402                 /*
 403                  * If we got some work, mark us as busy. If we didn't, but
 404                  * the list isn't empty, it means we stalled on hashed work.
 405                  * Mark us stalled so we don't keep looking for work when we
 406                  * can't make progress, any work completion or insertion will
 407                  * clear the stalled flag.
 408                  */
 409                 work = io_get_next_work(wqe, &hash);
 410                 if (work)
 411                         __io_worker_busy(wqe, worker, work);
 412                 else if (!list_empty(&wqe->work_list))
 413                         wqe->flags |= IO_WQE_FLAG_STALLED;
 414
 415                 spin_unlock_irq(&wqe->lock);
 416                 if (put_work && wq->put_work)
 417                         wq->put_work(old_work);
 418                 if (!work)
 419                         break;
 420 next:
 421                 /* flush any pending signals before assigning new work */
 422                 if (signal_pending(current))
 423                         flush_signals(current);
 424
 425                 spin_lock_irq(&worker->lock);
 426                 worker->cur_work = work;
 427                 spin_unlock_irq(&worker->lock);
 428
 429                 if ((work->flags & IO_WQ_WORK_NEEDS_FILES) &&
 430                     current->files != work->files) {
 431                         task_lock(current);
 432                         current->files = work->files;
 433                         task_unlock(current);
 434                 }
 435                 if ((work->flags & IO_WQ_WORK_NEEDS_USER) && !worker->mm &&
 436                     wq->mm && mmget_not_zero(wq->mm)) {
 437                         use_mm(wq->mm);
 438                         set_fs(USER_DS);
 439                         worker->mm = wq->mm;
 440                 }
 441                 if (test_bit(IO_WQ_BIT_CANCEL, &wq->state))
 442                         work->flags |= IO_WQ_WORK_CANCEL;
 443                 if (worker->mm)
 444                         work->flags |= IO_WQ_WORK_HAS_MM;
 445
 446                 if (wq->get_work && !(work->flags & IO_WQ_WORK_INTERNAL)) {
 447                         put_work = work;
 448                         wq->get_work(work);
 449                 }
 450
 451                 old_work = work;
 452                 work->func(&work);
 453
 454                 spin_lock_irq(&worker->lock);
 455                 worker->cur_work = NULL;
 456                 spin_unlock_irq(&worker->lock);
 457
 458                 spin_lock_irq(&wqe->lock);
 459
 460                 if (hash != -1U) {
 461                         wqe->hash_map &= ~BIT_ULL(hash);
 462                         wqe->flags &= ~IO_WQE_FLAG_STALLED;
 463                 }
 464                 if (work && work != old_work) {
 465                         spin_unlock_irq(&wqe->lock);
 466
 467                         if (put_work && wq->put_work) {
 468                                 wq->put_work(put_work);
 469                                 put_work = NULL;
 470                         }
 471
 472                         /* dependent work not hashed */
 473                         hash = -1U;
 474                         goto next;
 475                 }
 476         } while (1);
 477 }
 478
 479 static int io_wqe_worker(void *data)
 480 {
 481         struct io_worker *worker = data;
 482         struct io_wqe *wqe = worker->wqe;
 483         struct io_wq *wq = wqe->wq;
 484         DEFINE_WAIT(wait);
 485
 486         io_worker_start(wqe, worker);
 487
 488         while (!test_bit(IO_WQ_BIT_EXIT, &wq->state)) {
 489                 prepare_to_wait(&worker->wait, &wait, TASK_INTERRUPTIBLE);
 490
 491                 spin_lock_irq(&wqe->lock);
 492                 if (io_wqe_run_queue(wqe)) {
 493                         __set_current_state(TASK_RUNNING);
 494                         io_worker_handle_work(worker);
 495                         continue;
 496                 }
 497                 /* drops the lock on success, retry */
 498                 if (__io_worker_idle(wqe, worker)) {
 499                         __release(&wqe->lock);
 500                         continue;
 501                 }
 502                 spin_unlock_irq(&wqe->lock);
 503                 if (signal_pending(current))
 504                         flush_signals(current);
 505                 if (schedule_timeout(WORKER_IDLE_TIMEOUT))
 506                         continue;
 507                 /* timed out, exit unless we're the fixed worker */
 508                 if (test_bit(IO_WQ_BIT_EXIT, &wq->state) ||
 509                     !(worker->flags & IO_WORKER_F_FIXED))
 510                         break;
 511         }
 512
 513         finish_wait(&worker->wait, &wait);
 514
 515         if (test_bit(IO_WQ_BIT_EXIT, &wq->state)) {
 516                 spin_lock_irq(&wqe->lock);
 517                 if (!list_empty(&wqe->work_list))
 518                         io_worker_handle_work(worker);
 519                 else
 520                         spin_unlock_irq(&wqe->lock);
 521         }
 522
 523         io_worker_exit(worker);
 524         return 0;
 525 }
 526
 527 /*
 528  * Called when a worker is scheduled in. Mark us as currently running.
 529  */
 530 void io_wq_worker_running(struct task_struct *tsk)
 531 {
 532         struct io_worker *worker = kthread_data(tsk);
 533         struct io_wqe *wqe = worker->wqe;
 534
 535         if (!(worker->flags & IO_WORKER_F_UP))
 536                 return;
 537         if (worker->flags & IO_WORKER_F_RUNNING)
 538                 return;
 539         worker->flags |= IO_WORKER_F_RUNNING;
 540         io_wqe_inc_running(wqe, worker);
 541 }
 542
 543 /*
 544  * Called when worker is going to sleep. If there are no workers currently
 545  * running and we have work pending, wake up a free one or have the manager
 546  * set one up.
 547  */
 548 void io_wq_worker_sleeping(struct task_struct *tsk)
 549 {
 550         struct io_worker *worker = kthread_data(tsk);
 551         struct io_wqe *wqe = worker->wqe;
 552
 553         if (!(worker->flags & IO_WORKER_F_UP))
 554                 return;
 555         if (!(worker->flags & IO_WORKER_F_RUNNING))
 556                 return;
 557
 558         worker->flags &= ~IO_WORKER_F_RUNNING;
 559
 560         spin_lock_irq(&wqe->lock);
 561         io_wqe_dec_running(wqe, worker);
 562         spin_unlock_irq(&wqe->lock);
 563 }
 564
 565 static void create_io_worker(struct io_wq *wq, struct io_wqe *wqe, int index)
 566 {
 567         struct io_wqe_acct *acct =&wqe->acct[index];
 568         struct io_worker *worker;
 569
 570         worker = kcalloc_node(1, sizeof(*worker), GFP_KERNEL, wqe->node);
 571         if (!worker)
 572                 return;
 573
 574         refcount_set(&worker->ref, 1);
 575         worker->nulls_node.pprev = NULL;
 576         init_waitqueue_head(&worker->wait);
 577         worker->wqe = wqe;
 578         spin_lock_init(&worker->lock);
 579
 580         worker->task = kthread_create_on_node(io_wqe_worker, worker, wqe->node,
 581                                 "io_wqe_worker-%d/%d", index, wqe->node);
 582         if (IS_ERR(worker->task)) {
 583                 kfree(worker);
 584                 return;
 585         }
 586
 587         spin_lock_irq(&wqe->lock);
 588         hlist_nulls_add_head_rcu(&worker->nulls_node, &wqe->free_list);
 589         list_add_tail_rcu(&worker->all_list, &wqe->all_list);
 590         worker->flags |= IO_WORKER_F_FREE;
 591         if (index == IO_WQ_ACCT_BOUND)
 592                 worker->flags |= IO_WORKER_F_BOUND;
 593         if (!acct->nr_workers && (worker->flags & IO_WORKER_F_BOUND))
 594                 worker->flags |= IO_WORKER_F_FIXED;
 595         acct->nr_workers++;
 596         spin_unlock_irq(&wqe->lock);
 597
 598         if (index == IO_WQ_ACCT_UNBOUND)
 599                 atomic_inc(&wq->user->processes);
 600
 601         wake_up_process(worker->task);
 602 }
 603
 604 static inline bool io_wqe_need_worker(struct io_wqe *wqe, int index)
 605         __must_hold(wqe->lock)
 606 {
 607         struct io_wqe_acct *acct = &wqe->acct[index];
 608
 609         /* always ensure we have one bounded worker */
 610         if (index == IO_WQ_ACCT_BOUND && !acct->nr_workers)
 611                 return true;
 612         /* if we have available workers or no work, no need */
 613         if (!hlist_nulls_empty(&wqe->free_list) || !io_wqe_run_queue(wqe))
 614                 return false;
 615         return acct->nr_workers < acct->max_workers;
 616 }
 617
 618 /*
 619  * Manager thread. Tasked with creating new workers, if we need them.
 620  */
 621 static int io_wq_manager(void *data)
 622 {
 623         struct io_wq *wq = data;
 624
 625         while (!kthread_should_stop()) {
 626                 int i;
 627
 628                 for (i = 0; i < wq->nr_wqes; i++) {
 629                         struct io_wqe *wqe = wq->wqes[i];
 630                         bool fork_worker[2] = { false, false };
 631
 632                         spin_lock_irq(&wqe->lock);
 633                         if (io_wqe_need_worker(wqe, IO_WQ_ACCT_BOUND))
 634                                 fork_worker[IO_WQ_ACCT_BOUND] = true;
 635                         if (io_wqe_need_worker(wqe, IO_WQ_ACCT_UNBOUND))
 636                                 fork_worker[IO_WQ_ACCT_UNBOUND] = true;
 637                         spin_unlock_irq(&wqe->lock);
 638                         if (fork_worker[IO_WQ_ACCT_BOUND])
 639                                 create_io_worker(wq, wqe, IO_WQ_ACCT_BOUND);
 640                         if (fork_worker[IO_WQ_ACCT_UNBOUND])
 641                                 create_io_worker(wq, wqe, IO_WQ_ACCT_UNBOUND);
 642                 }
 643                 set_current_state(TASK_INTERRUPTIBLE);
 644                 schedule_timeout(HZ);
 645         }
 646
 647         return 0;
 648 }
 649
 650 static bool io_wq_can_queue(struct io_wqe *wqe, struct io_wqe_acct *acct,
 651                             struct io_wq_work *work)
 652 {
 653         bool free_worker;
 654
 655         if (!(work->flags & IO_WQ_WORK_UNBOUND))
 656                 return true;
 657         if (atomic_read(&acct->nr_running))
 658                 return true;
 659
 660         rcu_read_lock();
 661         free_worker = !hlist_nulls_empty(&wqe->free_list);
 662         rcu_read_unlock();
 663         if (free_worker)
 664                 return true;
 665
 666         if (atomic_read(&wqe->wq->user->processes) >= acct->max_workers &&
 667             !(capable(CAP_SYS_RESOURCE) || capable(CAP_SYS_ADMIN)))
 668                 return false;
 669
 670         return true;
 671 }
 672
 673 static void io_wqe_enqueue(struct io_wqe *wqe, struct io_wq_work *work)
 674 {
 675         struct io_wqe_acct *acct = io_work_get_acct(wqe, work);
 676         unsigned long flags;
 677
 678         /*
 679          * Do early check to see if we need a new unbound worker, and if we do,
 680          * if we're allowed to do so. This isn't 100% accurate as there's a
 681          * gap between this check and incrementing the value, but that's OK.
 682          * It's close enough to not be an issue, fork() has the same delay.
 683          */
 684         if (unlikely(!io_wq_can_queue(wqe, acct, work))) {
 685                 work->flags |= IO_WQ_WORK_CANCEL;
 686                 work->func(&work);
 687                 return;
 688         }
 689
 690         spin_lock_irqsave(&wqe->lock, flags);
 691         list_add_tail(&work->list, &wqe->work_list);
 692         wqe->flags &= ~IO_WQE_FLAG_STALLED;
 693         spin_unlock_irqrestore(&wqe->lock, flags);
 694
 695         if (!atomic_read(&acct->nr_running))
 696                 io_wqe_wake_worker(wqe, acct);
 697 }
 698
 699 void io_wq_enqueue(struct io_wq *wq, struct io_wq_work *work)
 700 {
 701         struct io_wqe *wqe = wq->wqes[numa_node_id()];
 702
 703         io_wqe_enqueue(wqe, work);
 704 }
 705
 706 /*
 707  * Enqueue work, hashed by some key. Work items that hash to the same value
 708  * will not be done in parallel. Used to limit concurrent writes, generally
 709  * hashed by inode.
 710  */
 711 void io_wq_enqueue_hashed(struct io_wq *wq, struct io_wq_work *work, void *val)
 712 {
 713         struct io_wqe *wqe = wq->wqes[numa_node_id()];
 714         unsigned bit;
 715
 716
 717         bit = hash_ptr(val, IO_WQ_HASH_ORDER);
 718         work->flags |= (IO_WQ_WORK_HASHED | (bit << IO_WQ_HASH_SHIFT));
 719         io_wqe_enqueue(wqe, work);
 720 }
 721
 722 static bool io_wqe_worker_send_sig(struct io_worker *worker, void *data)
 723 {
 724         send_sig(SIGINT, worker->task, 1);
 725         return false;
 726 }
 727
 728 /*
 729  * Iterate the passed in list and call the specific function for each
 730  * worker that isn't exiting
 731  */
 732 static bool io_wq_for_each_worker(struct io_wqe *wqe,
 733                                   bool (*func)(struct io_worker *, void *),
 734                                   void *data)
 735 {
 736         struct io_worker *worker;
 737         bool ret = false;
 738
 739         list_for_each_entry_rcu(worker, &wqe->all_list, all_list) {
 740                 if (io_worker_get(worker)) {
 741                         ret = func(worker, data);
 742                         io_worker_release(worker);
 743                         if (ret)
 744                                 break;
 745                 }
 746         }
 747
 748         return ret;
 749 }
 750
 751 void io_wq_cancel_all(struct io_wq *wq)
 752 {
 753         int i;
 754
 755         set_bit(IO_WQ_BIT_CANCEL, &wq->state);
 756
 757         /*
 758          * Browse both lists, as there's a gap between handing work off
 759          * to a worker and the worker putting itself on the busy_list
 760          */
 761         rcu_read_lock();
 762         for (i = 0; i < wq->nr_wqes; i++) {
 763                 struct io_wqe *wqe = wq->wqes[i];
 764
 765                 io_wq_for_each_worker(wqe, io_wqe_worker_send_sig, NULL);
 766         }
 767         rcu_read_unlock();
 768 }
 769
 770 struct io_cb_cancel_data {
 771         struct io_wqe *wqe;
 772         work_cancel_fn *cancel;
 773         void *caller_data;
 774 };
 775
 776 static bool io_work_cancel(struct io_worker *worker, void *cancel_data)
 777 {
 778         struct io_cb_cancel_data *data = cancel_data;
 779         unsigned long flags;
 780         bool ret = false;
 781
 782         /*
 783          * Hold the lock to avoid ->cur_work going out of scope, caller
 784          * may dereference the passed in work.
 785          */
 786         spin_lock_irqsave(&worker->lock, flags);
 787         if (worker->cur_work &&
 788             data->cancel(worker->cur_work, data->caller_data)) {
 789                 send_sig(SIGINT, worker->task, 1);
 790                 ret = true;
 791         }
 792         spin_unlock_irqrestore(&worker->lock, flags);
 793
 794         return ret;
 795 }
 796
 797 static enum io_wq_cancel io_wqe_cancel_cb_work(struct io_wqe *wqe,
 798                                                work_cancel_fn *cancel,
 799                                                void *cancel_data)
 800 {
 801         struct io_cb_cancel_data data = {
 802                 .wqe = wqe,
 803                 .cancel = cancel,
 804                 .caller_data = cancel_data,
 805         };
 806         struct io_wq_work *work;
 807         unsigned long flags;
 808         bool found = false;
 809
 810         spin_lock_irqsave(&wqe->lock, flags);
 811         list_for_each_entry(work, &wqe->work_list, list) {
 812                 if (cancel(work, cancel_data)) {
 813                         list_del(&work->list);
 814                         found = true;
 815                         break;
 816                 }
 817         }
 818         spin_unlock_irqrestore(&wqe->lock, flags);
 819
 820         if (found) {
 821                 work->flags |= IO_WQ_WORK_CANCEL;
 822                 work->func(&work);
 823                 return IO_WQ_CANCEL_OK;
 824         }
 825
 826         rcu_read_lock();
 827         found = io_wq_for_each_worker(wqe, io_work_cancel, &data);
 828         rcu_read_unlock();
 829         return found ? IO_WQ_CANCEL_RUNNING : IO_WQ_CANCEL_NOTFOUND;
 830 }
 831
 832 enum io_wq_cancel io_wq_cancel_cb(struct io_wq *wq, work_cancel_fn *cancel,
 833                                   void *data)
 834 {
 835         enum io_wq_cancel ret = IO_WQ_CANCEL_NOTFOUND;
 836         int i;
 837
 838         for (i = 0; i < wq->nr_wqes; i++) {
 839                 struct io_wqe *wqe = wq->wqes[i];
 840
 841                 ret = io_wqe_cancel_cb_work(wqe, cancel, data);
 842                 if (ret != IO_WQ_CANCEL_NOTFOUND)
 843                         break;
 844         }
 845
 846         return ret;
 847 }
 848
 849 static bool io_wq_worker_cancel(struct io_worker *worker, void *data)
 850 {
 851         struct io_wq_work *work = data;
 852         unsigned long flags;
 853         bool ret = false;
 854
 855         if (worker->cur_work != work)
 856                 return false;
 857
 858         spin_lock_irqsave(&worker->lock, flags);
 859         if (worker->cur_work == work) {
 860                 send_sig(SIGINT, worker->task, 1);
 861                 ret = true;
 862         }
 863         spin_unlock_irqrestore(&worker->lock, flags);
 864
 865         return ret;
 866 }
 867
 868 static enum io_wq_cancel io_wqe_cancel_work(struct io_wqe *wqe,
 869                                             struct io_wq_work *cwork)
 870 {
 871         struct io_wq_work *work;
 872         unsigned long flags;
 873         bool found = false;
 874
 875         cwork->flags |= IO_WQ_WORK_CANCEL;
 876
 877         /*
 878          * First check pending list, if we're lucky we can just remove it
 879          * from there. CANCEL_OK means that the work is returned as-new,
 880          * no completion will be posted for it.
 881          */
 882         spin_lock_irqsave(&wqe->lock, flags);
 883         list_for_each_entry(work, &wqe->work_list, list) {
 884                 if (work == cwork) {
 885                         list_del(&work->list);
 886                         found = true;
 887                         break;
 888                 }
 889         }
 890         spin_unlock_irqrestore(&wqe->lock, flags);
 891
 892         if (found) {
 893                 work->flags |= IO_WQ_WORK_CANCEL;
 894                 work->func(&work);
 895                 return IO_WQ_CANCEL_OK;
 896         }
 897
 898         /*
 899          * Now check if a free (going busy) or busy worker has the work
 900          * currently running. If we find it there, we'll return CANCEL_RUNNING
 901          * as an indication that we attempte to signal cancellation. The
 902          * completion will run normally in this case.
 903          */
 904         rcu_read_lock();
 905         found = io_wq_for_each_worker(wqe, io_wq_worker_cancel, cwork);
 906         rcu_read_unlock();
 907         return found ? IO_WQ_CANCEL_RUNNING : IO_WQ_CANCEL_NOTFOUND;
 908 }
 909
 910 enum io_wq_cancel io_wq_cancel_work(struct io_wq *wq, struct io_wq_work *cwork)
 911 {
 912         enum io_wq_cancel ret = IO_WQ_CANCEL_NOTFOUND;
 913         int i;
 914
 915         for (i = 0; i < wq->nr_wqes; i++) {
 916                 struct io_wqe *wqe = wq->wqes[i];
 917
 918                 ret = io_wqe_cancel_work(wqe, cwork);
 919                 if (ret != IO_WQ_CANCEL_NOTFOUND)
 920                         break;
 921         }
 922
 923         return ret;
 924 }
 925
 926 struct io_wq_flush_data {
 927         struct io_wq_work work;
 928         struct completion done;
 929 };
 930
 931 static void io_wq_flush_func(struct io_wq_work **workptr)
 932 {
 933         struct io_wq_work *work = *workptr;
 934         struct io_wq_flush_data *data;
 935
 936         data = container_of(work, struct io_wq_flush_data, work);
 937         complete(&data->done);
 938 }
 939
 940 /*
 941  * Doesn't wait for previously queued work to finish. When this completes,
 942  * it just means that previously queued work was started.
 943  */
 944 void io_wq_flush(struct io_wq *wq)
 945 {
 946         struct io_wq_flush_data data;
 947         int i;
 948
 949         for (i = 0; i < wq->nr_wqes; i++) {
 950                 struct io_wqe *wqe = wq->wqes[i];
 951
 952                 init_completion(&data.done);
 953                 INIT_IO_WORK(&data.work, io_wq_flush_func);
 954                 data.work.flags |= IO_WQ_WORK_INTERNAL;
 955                 io_wqe_enqueue(wqe, &data.work);
 956                 wait_for_completion(&data.done);
 957         }
 958 }
 959
 960 struct io_wq *io_wq_create(unsigned bounded, struct mm_struct *mm,
 961                            struct user_struct *user, get_work_fn *get_work,
 962                            put_work_fn *put_work)
 963 {
 964         int ret = -ENOMEM, i, node;
 965         struct io_wq *wq;
 966
 967         wq = kcalloc(1, sizeof(*wq), GFP_KERNEL);
 968         if (!wq)
 969                 return ERR_PTR(-ENOMEM);
 970
 971         wq->nr_wqes = num_online_nodes();
 972         wq->wqes = kcalloc(wq->nr_wqes, sizeof(struct io_wqe *), GFP_KERNEL);
 973         if (!wq->wqes) {
 974                 kfree(wq);
 975                 return ERR_PTR(-ENOMEM);
 976         }
 977
 978         wq->get_work = get_work;
 979         wq->put_work = put_work;
 980
 981         /* caller must already hold a reference to this */
 982         wq->user = user;
 983
 984         i = 0;
 985         refcount_set(&wq->refs, wq->nr_wqes);
 986         for_each_online_node(node) {
 987                 struct io_wqe *wqe;
 988
 989                 wqe = kcalloc_node(1, sizeof(struct io_wqe), GFP_KERNEL, node);
 990                 if (!wqe)
 991                         break;
 992                 wq->wqes[i] = wqe;
 993                 wqe->node = node;
 994                 wqe->acct[IO_WQ_ACCT_BOUND].max_workers = bounded;
 995                 atomic_set(&wqe->acct[IO_WQ_ACCT_BOUND].nr_running, 0);
 996                 if (user) {
 997                         wqe->acct[IO_WQ_ACCT_UNBOUND].max_workers =
 998                                         task_rlimit(current, RLIMIT_NPROC);
 999                 }
1000                 atomic_set(&wqe->acct[IO_WQ_ACCT_UNBOUND].nr_running, 0);
1001                 wqe->node = node;
1002                 wqe->wq = wq;
1003                 spin_lock_init(&wqe->lock);
1004                 INIT_LIST_HEAD(&wqe->work_list);
1005                 INIT_HLIST_NULLS_HEAD(&wqe->free_list, 0);
1006                 INIT_HLIST_NULLS_HEAD(&wqe->busy_list, 1);
1007                 INIT_LIST_HEAD(&wqe->all_list);
1008
1009                 i++;
1010         }
1011
1012         init_completion(&wq->done);
1013
1014         if (i != wq->nr_wqes)
1015                 goto err;
1016
1017         /* caller must have already done mmgrab() on this mm */
1018         wq->mm = mm;
1019
1020         wq->manager = kthread_create(io_wq_manager, wq, "io_wq_manager");
1021         if (!IS_ERR(wq->manager)) {
1022                 wake_up_process(wq->manager);
1023                 return wq;
1024         }
1025
1026         ret = PTR_ERR(wq->manager);
1027         wq->manager = NULL;
1028 err:
1029         complete(&wq->done);
1030         io_wq_destroy(wq);
1031         return ERR_PTR(ret);
1032 }
1033
1034 static bool io_wq_worker_wake(struct io_worker *worker, void *data)
1035 {
1036         wake_up_process(worker->task);
1037         return false;
1038 }
1039
1040 void io_wq_destroy(struct io_wq *wq)
1041 {
1042         int i;
1043
1044         if (wq->manager) {
1045                 set_bit(IO_WQ_BIT_EXIT, &wq->state);
1046                 kthread_stop(wq->manager);
1047         }
1048
1049         rcu_read_lock();
1050         for (i = 0; i < wq->nr_wqes; i++) {
1051                 struct io_wqe *wqe = wq->wqes[i];
1052
1053                 if (!wqe)
1054                         continue;
1055                 io_wq_for_each_worker(wqe, io_wq_worker_wake, NULL);
1056         }
1057         rcu_read_unlock();
1058
1059         wait_for_completion(&wq->done);
1060
1061         for (i = 0; i < wq->nr_wqes; i++)
1062                 kfree(wq->wqes[i]);
1063         kfree(wq->wqes);
1064         kfree(wq);
1065 }