6 files changed, 23 insertions, 14 deletions

diff --git a/block/blk-mq-tag.c b/block/blk-mq-tag.c
index d64a02fb1f73..5d70edc9855f 100644
--- a/block/blk-mq-tag.c
+++ b/block/blk-mq-tag.c
@@ -36,7 +36,8 @@ static unsigned int __blk_mq_get_tag(struct blk_mq_tags *tags, gfp_t gfp)
 {
         int tag;
 
-        tag = percpu_ida_alloc(&tags->free_tags, gfp);
+        tag = percpu_ida_alloc(&tags->free_tags, (gfp & __GFP_WAIT) ?
+                               TASK_UNINTERRUPTIBLE : TASK_RUNNING);
         if (tag < 0)
                 return BLK_MQ_TAG_FAIL;
         return tag + tags->nr_reserved_tags;
@@ -52,7 +53,8 @@ static unsigned int __blk_mq_get_reserved_tag(struct blk_mq_tags *tags,
                 return BLK_MQ_TAG_FAIL;
         }
 
-        tag = percpu_ida_alloc(&tags->reserved_tags, gfp);
+        tag = percpu_ida_alloc(&tags->reserved_tags, (gfp & __GFP_WAIT) ?
+                               TASK_UNINTERRUPTIBLE : TASK_RUNNING);
         if (tag < 0)
                 return BLK_MQ_TAG_FAIL;
         return tag;
diff --git a/drivers/target/iscsi/iscsi_target_util.c b/drivers/target/iscsi/iscsi_target_util.c
index 0819e688a398..9b8e1db1257b 100644
--- a/drivers/target/iscsi/iscsi_target_util.c
+++ b/drivers/target/iscsi/iscsi_target_util.c
@@ -156,9 +156,13 @@ struct iscsi_cmd *iscsit_allocate_cmd(struct iscsi_conn *conn, gfp_t gfp_mask)
 {
         struct iscsi_cmd *cmd;
         struct se_session *se_sess = conn->sess->se_sess;
-        int size, tag;
+        int size, tag, state = (gfp_mask & __GFP_WAIT) ? TASK_UNINTERRUPTIBLE :
+                                TASK_RUNNING;
+
+        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, state);
+        if (tag < 0)
+                return NULL;
 
-        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, gfp_mask);
         size = sizeof(struct iscsi_cmd) + conn->conn_transport->priv_size;
         cmd = (struct iscsi_cmd *)(se_sess->sess_cmd_map + (tag * size));
         memset(cmd, 0, size);
diff --git a/drivers/target/tcm_fc/tfc_cmd.c b/drivers/target/tcm_fc/tfc_cmd.c
index 479ec5621a4e..8b2c1aaf81de 100644
--- a/drivers/target/tcm_fc/tfc_cmd.c
+++ b/drivers/target/tcm_fc/tfc_cmd.c
@@ -438,7 +438,7 @@ static void ft_recv_cmd(struct ft_sess *sess, struct fc_frame *fp)
         struct se_session *se_sess = sess->se_sess;
         int tag;
 
-        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, GFP_ATOMIC);
+        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, TASK_RUNNING);
         if (tag < 0)
                 goto busy;
 
diff --git a/drivers/vhost/scsi.c b/drivers/vhost/scsi.c
index 84488a890b69..2d084fb8d4d3 100644
--- a/drivers/vhost/scsi.c
+++ b/drivers/vhost/scsi.c
@@ -728,7 +728,7 @@ vhost_scsi_get_tag(struct vhost_virtqueue *vq,
         }
         se_sess = tv_nexus->tvn_se_sess;
 
-        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, GFP_ATOMIC);
+        tag = percpu_ida_alloc(&se_sess->sess_tag_pool, TASK_RUNNING);
         if (tag < 0) {
                 pr_err("Unable to obtain tag for tcm_vhost_cmd\n");
                 return ERR_PTR(-ENOMEM);
diff --git a/include/linux/percpu_ida.h b/include/linux/percpu_ida.h
index 1900bd0fa639..f5cfdd6a5539 100644
--- a/include/linux/percpu_ida.h
+++ b/include/linux/percpu_ida.h
@@ -4,6 +4,7 @@
 #include <linux/types.h>
 #include <linux/bitops.h>
 #include <linux/init.h>
+#include <linux/sched.h>
 #include <linux/spinlock_types.h>
 #include <linux/wait.h>
 #include <linux/cpumask.h>
@@ -61,7 +62,7 @@ struct percpu_ida {
 /* Max size of percpu freelist, */
 #define IDA_DEFAULT_PCPU_SIZE   ((IDA_DEFAULT_PCPU_BATCH_MOVE * 3) / 2)
 
-int percpu_ida_alloc(struct percpu_ida *pool, gfp_t gfp);
+int percpu_ida_alloc(struct percpu_ida *pool, int state);
 void percpu_ida_free(struct percpu_ida *pool, unsigned tag);
 
 void percpu_ida_destroy(struct percpu_ida *pool);
diff --git a/lib/percpu_ida.c b/lib/percpu_ida.c
index 9d054bf91d0f..58b671484ac2 100644
--- a/lib/percpu_ida.c
+++ b/lib/percpu_ida.c
@@ -132,22 +132,22 @@ static inline unsigned alloc_local_tag(struct percpu_ida_cpu *tags)
 /**
  * percpu_ida_alloc - allocate a tag
  * @pool: pool to allocate from
- * @gfp: gfp flags
+ * @state: task state for prepare_to_wait
  *
  * Returns a tag - an integer in the range [0..nr_tags) (passed to
  * tag_pool_init()), or otherwise -ENOSPC on allocation failure.
  *
  * Safe to be called from interrupt context (assuming it isn't passed
- * __GFP_WAIT, of course).
+ * TASK_UNINTERRUPTIBLE, of course).
  *
  * @gfp indicates whether or not to wait until a free id is available (it's not
  * used for internal memory allocations); thus if passed __GFP_WAIT we may sleep
  * however long it takes until another thread frees an id (same semantics as a
  * mempool).
  *
- * Will not fail if passed __GFP_WAIT.
+ * Will not fail if passed TASK_UNINTERRUPTIBLE.
  */
-int percpu_ida_alloc(struct percpu_ida *pool, gfp_t gfp)
+int percpu_ida_alloc(struct percpu_ida *pool, int state)
 {
         DEFINE_WAIT(wait);
         struct percpu_ida_cpu *tags;
@@ -174,7 +174,8 @@ int percpu_ida_alloc(struct percpu_ida *pool, gfp_t gfp)
                  *
                  * global lock held and irqs disabled, don't need percpu lock
                  */
-                prepare_to_wait(&pool->wait, &wait, TASK_UNINTERRUPTIBLE);
+                if (state != TASK_RUNNING)
+                        prepare_to_wait(&pool->wait, &wait, state);
 
                 if (!tags->nr_free)
                         alloc_global_tags(pool, tags);
@@ -191,7 +192,7 @@ int percpu_ida_alloc(struct percpu_ida *pool, gfp_t gfp)
                 spin_unlock(&pool->lock);
                 local_irq_restore(flags);
 
-                if (tag >= 0 || !(gfp & __GFP_WAIT))
+                if (tag >= 0 || state == TASK_RUNNING)
                         break;
 
                 schedule();
@@ -199,8 +200,9 @@ int percpu_ida_alloc(struct percpu_ida *pool, gfp_t gfp)
                 local_irq_save(flags);
                 tags = this_cpu_ptr(pool->tag_cpu);
         }
+        if (state != TASK_RUNNING)
+                finish_wait(&pool->wait, &wait);
 
-        finish_wait(&pool->wait, &wait);
         return tag;
 }
 EXPORT_SYMBOL_GPL(percpu_ida_alloc);