do {
                struct request *rq;
+               int budget_token;
 
                if (e->type->ops.has_work && !e->type->ops.has_work(hctx))
                        break;
                        break;
                }
 
-               if (!blk_mq_get_dispatch_budget(q))
+               budget_token = blk_mq_get_dispatch_budget(q);
+               if (budget_token < 0)
                        break;
 
                rq = e->type->ops.dispatch_request(hctx);
                if (!rq) {
-                       blk_mq_put_dispatch_budget(q);
+                       blk_mq_put_dispatch_budget(q, budget_token);
                        /*
                         * We're releasing without dispatching. Holding the
                         * budget could have blocked any "hctx"s with the
                        break;
                }
 
+               blk_mq_set_rq_budget_token(rq, budget_token);
+
                /*
                 * Now this rq owns the budget which has to be released
                 * if this rq won't be queued to driver via .queue_rq()
        struct request *rq;
 
        do {
+               int budget_token;
+
                if (!list_empty_careful(&hctx->dispatch)) {
                        ret = -EAGAIN;
                        break;
                if (!sbitmap_any_bit_set(&hctx->ctx_map))
                        break;
 
-               if (!blk_mq_get_dispatch_budget(q))
+               budget_token = blk_mq_get_dispatch_budget(q);
+               if (budget_token < 0)
                        break;
 
                rq = blk_mq_dequeue_from_ctx(hctx, ctx);
                if (!rq) {
-                       blk_mq_put_dispatch_budget(q);
+                       blk_mq_put_dispatch_budget(q, budget_token);
                        /*
                         * We're releasing without dispatching. Holding the
                         * budget could have blocked any "hctx"s with the
                        break;
                }
 
+               blk_mq_set_rq_budget_token(rq, budget_token);
+
                /*
                 * Now this rq owns the budget which has to be released
                 * if this rq won't be queued to driver via .queue_rq()
 
                                                  bool need_budget)
 {
        struct blk_mq_hw_ctx *hctx = rq->mq_hctx;
+       int budget_token = -1;
 
-       if (need_budget && !blk_mq_get_dispatch_budget(rq->q)) {
-               blk_mq_put_driver_tag(rq);
-               return PREP_DISPATCH_NO_BUDGET;
+       if (need_budget) {
+               budget_token = blk_mq_get_dispatch_budget(rq->q);
+               if (budget_token < 0) {
+                       blk_mq_put_driver_tag(rq);
+                       return PREP_DISPATCH_NO_BUDGET;
+               }
+               blk_mq_set_rq_budget_token(rq, budget_token);
        }
 
        if (!blk_mq_get_driver_tag(rq)) {
                         * together during handling partial dispatch
                         */
                        if (need_budget)
-                               blk_mq_put_dispatch_budget(rq->q);
+                               blk_mq_put_dispatch_budget(rq->q, budget_token);
                        return PREP_DISPATCH_NO_TAG;
                }
        }
 
 /* release all allocated budgets before calling to blk_mq_dispatch_rq_list */
 static void blk_mq_release_budgets(struct request_queue *q,
-               unsigned int nr_budgets)
+               struct list_head *list)
 {
-       int i;
+       struct request *rq;
 
-       for (i = 0; i < nr_budgets; i++)
-               blk_mq_put_dispatch_budget(q);
+       list_for_each_entry(rq, list, queuelist) {
+               int budget_token = blk_mq_get_rq_budget_token(rq);
+
+               if (budget_token >= 0)
+                       blk_mq_put_dispatch_budget(q, budget_token);
+       }
 }
 
 /*
                        (hctx->flags & BLK_MQ_F_TAG_QUEUE_SHARED);
                bool no_budget_avail = prep == PREP_DISPATCH_NO_BUDGET;
 
-               blk_mq_release_budgets(q, nr_budgets);
+               if (nr_budgets)
+                       blk_mq_release_budgets(q, list);
 
                spin_lock(&hctx->lock);
                list_splice_tail_init(list, &hctx->dispatch);
 {
        struct request_queue *q = rq->q;
        bool run_queue = true;
+       int budget_token;
 
        /*
         * RCU or SRCU read lock is needed before checking quiesced flag.
        if (q->elevator && !bypass_insert)
                goto insert;
 
-       if (!blk_mq_get_dispatch_budget(q))
+       budget_token = blk_mq_get_dispatch_budget(q);
+       if (budget_token < 0)
                goto insert;
 
+       blk_mq_set_rq_budget_token(rq, budget_token);
+
        if (!blk_mq_get_driver_tag(rq)) {
-               blk_mq_put_dispatch_budget(q);
+               blk_mq_put_dispatch_budget(q, budget_token);
                goto insert;
        }
 
 
 void blk_mq_in_flight_rw(struct request_queue *q, struct block_device *part,
                unsigned int inflight[2]);
 
-static inline void blk_mq_put_dispatch_budget(struct request_queue *q)
+static inline void blk_mq_put_dispatch_budget(struct request_queue *q,
+                                             int budget_token)
 {
        if (q->mq_ops->put_budget)
-               q->mq_ops->put_budget(q);
+               q->mq_ops->put_budget(q, budget_token);
 }
 
-static inline bool blk_mq_get_dispatch_budget(struct request_queue *q)
+static inline int blk_mq_get_dispatch_budget(struct request_queue *q)
 {
        if (q->mq_ops->get_budget)
                return q->mq_ops->get_budget(q);
-       return true;
+       return 0;
+}
+
+static inline void blk_mq_set_rq_budget_token(struct request *rq, int token)
+{
+       if (token < 0)
+               return;
+
+       if (rq->q->mq_ops->set_rq_budget_token)
+               rq->q->mq_ops->set_rq_budget_token(rq, token);
+}
+
+static inline int blk_mq_get_rq_budget_token(struct request *rq)
+{
+       if (rq->q->mq_ops->get_rq_budget_token)
+               return rq->q->mq_ops->get_rq_budget_token(rq);
+       return -1;
 }
 
 static inline void __blk_mq_inc_active_requests(struct blk_mq_hw_ctx *hctx)
 
                atomic_dec(&starget->target_busy);
 
        atomic_dec(&sdev->device_busy);
+       cmd->budget_token = -1;
 }
 
 static void scsi_kick_queue(struct request_queue *q)
        unsigned long jiffies_at_alloc;
        int retries, to_clear;
        bool in_flight;
+       int budget_token = cmd->budget_token;
 
        if (!blk_rq_is_scsi(rq) && !(flags & SCMD_INITIALIZED)) {
                flags |= SCMD_INITIALIZED;
        cmd->retries = retries;
        if (in_flight)
                __set_bit(SCMD_STATE_INFLIGHT, &cmd->state);
+       cmd->budget_token = budget_token;
 
 }
 
        blk_mq_complete_request(cmd->request);
 }
 
-static void scsi_mq_put_budget(struct request_queue *q)
+static void scsi_mq_put_budget(struct request_queue *q, int budget_token)
 {
        struct scsi_device *sdev = q->queuedata;
 
        atomic_dec(&sdev->device_busy);
 }
 
-static bool scsi_mq_get_budget(struct request_queue *q)
+static int scsi_mq_get_budget(struct request_queue *q)
 {
        struct scsi_device *sdev = q->queuedata;
 
        if (scsi_dev_queue_ready(q, sdev))
-               return true;
+               return 0;
 
        atomic_inc(&sdev->restarts);
 
        if (unlikely(atomic_read(&sdev->device_busy) == 0 &&
                                !scsi_device_blocked(sdev)))
                blk_mq_delay_run_hw_queues(sdev->request_queue, SCSI_QUEUE_DELAY);
-       return false;
+       return -1;
 }
 
 static void scsi_mq_set_rq_budget_token(struct request *req, int token)
        blk_status_t ret;
        int reason;
 
+       WARN_ON_ONCE(cmd->budget_token < 0);
+
        /*
         * If the device is not in running state we will reject some or all
         * commands.
        if (scsi_target(sdev)->can_queue > 0)
                atomic_dec(&scsi_target(sdev)->target_busy);
 out_put_budget:
-       scsi_mq_put_budget(q);
+       scsi_mq_put_budget(q, cmd->budget_token);
+       cmd->budget_token = -1;
        switch (ret) {
        case BLK_STS_OK:
                break;
 
         * reserved budget. Also we have to handle failure case
         * of .get_budget for avoiding I/O deadlock.
         */
-       bool (*get_budget)(struct request_queue *);
+       int (*get_budget)(struct request_queue *);
 
        /**
         * @put_budget: Release the reserved budget.
         */
-       void (*put_budget)(struct request_queue *);
+       void (*put_budget)(struct request_queue *, int);
 
        /*
         * @set_rq_budget_toekn: store rq's budget token