lib/percpu_ida.c

   1 /*
   2  * Percpu IDA library
   3  *
   4  * Copyright (C) 2013 Datera, Inc. Kent Overstreet
   5  *
   6  * This program is free software; you can redistribute it and/or
   7  * modify it under the terms of the GNU General Public License as
   8  * published by the Free Software Foundation; either version 2, or (at
   9  * your option) any later version.
  10  *
  11  * This program is distributed in the hope that it will be useful, but
  12  * WITHOUT ANY WARRANTY; without even the implied warranty of
  13  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
  14  * General Public License for more details.
  15  */
  16
  17 #include <linux/mm.h>
  18 #include <linux/bitmap.h>
  19 #include <linux/bitops.h>
  20 #include <linux/bug.h>
  21 #include <linux/err.h>
  22 #include <linux/export.h>
  23 #include <linux/init.h>
  24 #include <linux/kernel.h>
  25 #include <linux/percpu.h>
  26 #include <linux/sched/signal.h>
  27 #include <linux/string.h>
  28 #include <linux/spinlock.h>
  29 #include <linux/percpu_ida.h>
  30
  31 struct percpu_ida_cpu {
  32         /*
  33          * Even though this is percpu, we need a lock for tag stealing by remote
  34          * CPUs:
  35          */
  36         spinlock_t                      lock;
  37
  38         /* nr_free/freelist form a stack of free IDs */
  39         unsigned                        nr_free;
  40         unsigned                        freelist[];
  41 };
  42
  43 static inline void move_tags(unsigned *dst, unsigned *dst_nr,
  44                              unsigned *src, unsigned *src_nr,
  45                              unsigned nr)
  46 {
  47         *src_nr -= nr;
  48         memcpy(dst + *dst_nr, src + *src_nr, sizeof(unsigned) * nr);
  49         *dst_nr += nr;
  50 }
  51
  52 /*
  53  * Try to steal tags from a remote cpu's percpu freelist.
  54  *
  55  * We first check how many percpu freelists have tags
  56  *
  57  * Then we iterate through the cpus until we find some tags - we don't attempt
  58  * to find the "best" cpu to steal from, to keep cacheline bouncing to a
  59  * minimum.
  60  */
  61 static inline void steal_tags(struct percpu_ida *pool,
  62                               struct percpu_ida_cpu *tags)
  63 {
  64         unsigned cpus_have_tags, cpu = pool->cpu_last_stolen;
  65         struct percpu_ida_cpu *remote;
  66
  67         for (cpus_have_tags = cpumask_weight(&pool->cpus_have_tags);
  68              cpus_have_tags; cpus_have_tags--) {
  69                 cpu = cpumask_next(cpu, &pool->cpus_have_tags);
  70
  71                 if (cpu >= nr_cpu_ids) {
  72                         cpu = cpumask_first(&pool->cpus_have_tags);
  73                         if (cpu >= nr_cpu_ids)
  74                                 BUG();
  75                 }
  76
  77                 pool->cpu_last_stolen = cpu;
  78                 remote = per_cpu_ptr(pool->tag_cpu, cpu);
  79
  80                 cpumask_clear_cpu(cpu, &pool->cpus_have_tags);
  81
  82                 if (remote == tags)
  83                         continue;
  84
  85                 spin_lock(&remote->lock);
  86
  87                 if (remote->nr_free) {
  88                         memcpy(tags->freelist,
  89                                remote->freelist,
  90                                sizeof(unsigned) * remote->nr_free);
  91
  92                         tags->nr_free = remote->nr_free;
  93                         remote->nr_free = 0;
  94                 }
  95
  96                 spin_unlock(&remote->lock);
  97
  98                 if (tags->nr_free)
  99                         break;
 100         }
 101 }
 102
 103 /*
 104  * Pop up to IDA_PCPU_BATCH_MOVE IDs off the global freelist, and push them onto
 105  * our percpu freelist:
 106  */
 107 static inline void alloc_global_tags(struct percpu_ida *pool,
 108                                      struct percpu_ida_cpu *tags)
 109 {
 110         move_tags(tags->freelist, &tags->nr_free,
 111                   pool->freelist, &pool->nr_free,
 112                   min(pool->nr_free, pool->percpu_batch_size));
 113 }
 114
 115 /**
 116  * percpu_ida_alloc - allocate a tag
 117  * @pool: pool to allocate from
 118  * @state: task state for prepare_to_wait
 119  *
 120  * Returns a tag - an integer in the range [0..nr_tags) (passed to
 121  * tag_pool_init()), or otherwise -ENOSPC on allocation failure.
 122  *
 123  * Safe to be called from interrupt context (assuming it isn't passed
 124  * TASK_UNINTERRUPTIBLE | TASK_INTERRUPTIBLE, of course).
 125  *
 126  * @gfp indicates whether or not to wait until a free id is available (it's not
 127  * used for internal memory allocations); thus if passed __GFP_RECLAIM we may sleep
 128  * however long it takes until another thread frees an id (same semantics as a
 129  * mempool).
 130  *
 131  * Will not fail if passed TASK_UNINTERRUPTIBLE | TASK_INTERRUPTIBLE.
 132  */
 133 int percpu_ida_alloc(struct percpu_ida *pool, int state)
 134 {
 135         DEFINE_WAIT(wait);
 136         struct percpu_ida_cpu *tags;
 137         unsigned long flags;
 138         int tag = -ENOSPC;
 139
 140         tags = raw_cpu_ptr(pool->tag_cpu);
 141         spin_lock_irqsave(&tags->lock, flags);
 142
 143         /* Fastpath */
 144         if (likely(tags->nr_free >= 0)) {
 145                 tag = tags->freelist[--tags->nr_free];
 146                 spin_unlock_irqrestore(&tags->lock, flags);
 147                 return tag;
 148         }
 149         spin_unlock_irqrestore(&tags->lock, flags);
 150
 151         while (1) {
 152                 spin_lock_irqsave(&pool->lock, flags);
 153                 tags = this_cpu_ptr(pool->tag_cpu);
 154
 155                 /*
 156                  * prepare_to_wait() must come before steal_tags(), in case
 157                  * percpu_ida_free() on another cpu flips a bit in
 158                  * cpus_have_tags
 159                  *
 160                  * global lock held and irqs disabled, don't need percpu lock
 161                  */
 162                 if (state != TASK_RUNNING)
 163                         prepare_to_wait(&pool->wait, &wait, state);
 164
 165                 if (!tags->nr_free)
 166                         alloc_global_tags(pool, tags);
 167                 if (!tags->nr_free)
 168                         steal_tags(pool, tags);
 169
 170                 if (tags->nr_free) {
 171                         tag = tags->freelist[--tags->nr_free];
 172                         if (tags->nr_free)
 173                                 cpumask_set_cpu(smp_processor_id(),
 174                                                 &pool->cpus_have_tags);
 175                 }
 176
 177                 spin_unlock_irqrestore(&pool->lock, flags);
 178
 179                 if (tag >= 0 || state == TASK_RUNNING)
 180                         break;
 181
 182                 if (signal_pending_state(state, current)) {
 183                         tag = -ERESTARTSYS;
 184                         break;
 185                 }
 186
 187                 schedule();
 188         }
 189         if (state != TASK_RUNNING)
 190                 finish_wait(&pool->wait, &wait);
 191
 192         return tag;
 193 }
 194 EXPORT_SYMBOL_GPL(percpu_ida_alloc);
 195
 196 /**
 197  * percpu_ida_free - free a tag
 198  * @pool: pool @tag was allocated from
 199  * @tag: a tag previously allocated with percpu_ida_alloc()
 200  *
 201  * Safe to be called from interrupt context.
 202  */
 203 void percpu_ida_free(struct percpu_ida *pool, unsigned tag)
 204 {
 205         struct percpu_ida_cpu *tags;
 206         unsigned long flags;
 207         unsigned nr_free;
 208
 209         BUG_ON(tag >= pool->nr_tags);
 210
 211         tags = raw_cpu_ptr(pool->tag_cpu);
 212
 213         spin_lock_irqsave(&tags->lock, flags);
 214         tags->freelist[tags->nr_free++] = tag;
 215
 216         nr_free = tags->nr_free;
 217
 218         if (nr_free == 1) {
 219                 cpumask_set_cpu(smp_processor_id(),
 220                                 &pool->cpus_have_tags);
 221                 wake_up(&pool->wait);
 222         }
 223         spin_unlock_irqrestore(&tags->lock, flags);
 224
 225         if (nr_free == pool->percpu_max_size) {
 226                 spin_lock_irqsave(&pool->lock, flags);
 227                 spin_lock(&tags->lock);
 228
 229                 if (tags->nr_free == pool->percpu_max_size) {
 230                         move_tags(pool->freelist, &pool->nr_free,
 231                                   tags->freelist, &tags->nr_free,
 232                                   pool->percpu_batch_size);
 233
 234                         wake_up(&pool->wait);
 235                 }
 236                 spin_unlock(&tags->lock);
 237                 spin_unlock_irqrestore(&pool->lock, flags);
 238         }
 239 }
 240 EXPORT_SYMBOL_GPL(percpu_ida_free);
 241
 242 /**
 243  * percpu_ida_destroy - release a tag pool's resources
 244  * @pool: pool to free
 245  *
 246  * Frees the resources allocated by percpu_ida_init().
 247  */
 248 void percpu_ida_destroy(struct percpu_ida *pool)
 249 {
 250         free_percpu(pool->tag_cpu);
 251         free_pages((unsigned long) pool->freelist,
 252                    get_order(pool->nr_tags * sizeof(unsigned)));
 253 }
 254 EXPORT_SYMBOL_GPL(percpu_ida_destroy);
 255
 256 /**
 257  * percpu_ida_init - initialize a percpu tag pool
 258  * @pool: pool to initialize
 259  * @nr_tags: number of tags that will be available for allocation
 260  *
 261  * Initializes @pool so that it can be used to allocate tags - integers in the
 262  * range [0, nr_tags). Typically, they'll be used by driver code to refer to a
 263  * preallocated array of tag structures.
 264  *
 265  * Allocation is percpu, but sharding is limited by nr_tags - for best
 266  * performance, the workload should not span more cpus than nr_tags / 128.
 267  */
 268 int __percpu_ida_init(struct percpu_ida *pool, unsigned long nr_tags,
 269         unsigned long max_size, unsigned long batch_size)
 270 {
 271         unsigned i, cpu, order;
 272
 273         memset(pool, 0, sizeof(*pool));
 274
 275         init_waitqueue_head(&pool->wait);
 276         spin_lock_init(&pool->lock);
 277         pool->nr_tags = nr_tags;
 278         pool->percpu_max_size = max_size;
 279         pool->percpu_batch_size = batch_size;
 280
 281         /* Guard against overflow */
 282         if (nr_tags > (unsigned) INT_MAX + 1) {
 283                 pr_err("percpu_ida_init(): nr_tags too large\n");
 284                 return -EINVAL;
 285         }
 286
 287         order = get_order(nr_tags * sizeof(unsigned));
 288         pool->freelist = (void *) __get_free_pages(GFP_KERNEL, order);
 289         if (!pool->freelist)
 290                 return -ENOMEM;
 291
 292         for (i = 0; i < nr_tags; i++)
 293                 pool->freelist[i] = i;
 294
 295         pool->nr_free = nr_tags;
 296
 297         pool->tag_cpu = __alloc_percpu(sizeof(struct percpu_ida_cpu) +
 298                                        pool->percpu_max_size * sizeof(unsigned),
 299                                        sizeof(unsigned));
 300         if (!pool->tag_cpu)
 301                 goto err;
 302
 303         for_each_possible_cpu(cpu)
 304                 spin_lock_init(&per_cpu_ptr(pool->tag_cpu, cpu)->lock);
 305
 306         return 0;
 307 err:
 308         percpu_ida_destroy(pool);
 309         return -ENOMEM;
 310 }
 311 EXPORT_SYMBOL_GPL(__percpu_ida_init);
 312
 313 /**
 314  * percpu_ida_for_each_free - iterate free ids of a pool
 315  * @pool: pool to iterate
 316  * @fn: interate callback function
 317  * @data: parameter for @fn
 318  *
 319  * Note, this doesn't guarantee to iterate all free ids restrictly. Some free
 320  * ids might be missed, some might be iterated duplicated, and some might
 321  * be iterated and not free soon.
 322  */
 323 int percpu_ida_for_each_free(struct percpu_ida *pool, percpu_ida_cb fn,
 324         void *data)
 325 {
 326         unsigned long flags;
 327         struct percpu_ida_cpu *remote;
 328         unsigned cpu, i, err = 0;
 329
 330         for_each_possible_cpu(cpu) {
 331                 remote = per_cpu_ptr(pool->tag_cpu, cpu);
 332                 spin_lock_irqsave(&remote->lock, flags);
 333                 for (i = 0; i < remote->nr_free; i++) {
 334                         err = fn(remote->freelist[i], data);
 335                         if (err)
 336                                 break;
 337                 }
 338                 spin_unlock_irqrestore(&remote->lock, flags);
 339                 if (err)
 340                         goto out;
 341         }
 342
 343         spin_lock_irqsave(&pool->lock, flags);
 344         for (i = 0; i < pool->nr_free; i++) {
 345                 err = fn(pool->freelist[i], data);
 346                 if (err)
 347                         break;
 348         }
 349         spin_unlock_irqrestore(&pool->lock, flags);
 350 out:
 351         return err;
 352 }
 353 EXPORT_SYMBOL_GPL(percpu_ida_for_each_free);
 354
 355 /**
 356  * percpu_ida_free_tags - return free tags number of a specific cpu or global pool
 357  * @pool: pool related
 358  * @cpu: specific cpu or global pool if @cpu == nr_cpu_ids
 359  *
 360  * Note: this just returns a snapshot of free tags number.
 361  */
 362 unsigned percpu_ida_free_tags(struct percpu_ida *pool, int cpu)
 363 {
 364         struct percpu_ida_cpu *remote;
 365         if (cpu == nr_cpu_ids)
 366                 return pool->nr_free;
 367         remote = per_cpu_ptr(pool->tag_cpu, cpu);
 368         return remote->nr_free;
 369 }
 370 EXPORT_SYMBOL_GPL(percpu_ida_free_tags);