1 files changed, 48 insertions, 8 deletions

diff --git a/block/blk-iolatency.c b/block/blk-iolatency.c
index fc714ef402a6..2620baa1f699 100644
--- a/block/blk-iolatency.c
+++ b/block/blk-iolatency.c
@@ -72,6 +72,7 @@
 #include <linux/sched/loadavg.h>
 #include <linux/sched/signal.h>
 #include <trace/events/block.h>
+#include <linux/blk-mq.h>
 #include "blk-rq-qos.h"
 #include "blk-stat.h"
 
@@ -591,6 +592,7 @@ static void blkcg_iolatency_done_bio(struct rq_qos *rqos, struct bio *bio)
         u64 now = ktime_to_ns(ktime_get());
         bool issue_as_root = bio_issue_as_root_blkg(bio);
         bool enabled = false;
+        int inflight = 0;
 
         blkg = bio->bi_blkg;
         if (!blkg || !bio_flagged(bio, BIO_TRACKED))
@@ -601,6 +603,9 @@ static void blkcg_iolatency_done_bio(struct rq_qos *rqos, struct bio *bio)
                 return;
 
         enabled = blk_iolatency_enabled(iolat->blkiolat);
+        if (!enabled)
+                return;
+
         while (blkg && blkg->parent) {
                 iolat = blkg_to_lat(blkg);
                 if (!iolat) {
@@ -609,8 +614,9 @@ static void blkcg_iolatency_done_bio(struct rq_qos *rqos, struct bio *bio)
                 }
                 rqw = &iolat->rq_wait;
 
-                atomic_dec(&rqw->inflight);
-                if (!enabled || iolat->min_lat_nsec == 0)
+                inflight = atomic_dec_return(&rqw->inflight);
+                WARN_ON_ONCE(inflight < 0);
+                if (iolat->min_lat_nsec == 0)
                         goto next;
                 iolatency_record_time(iolat, &bio->bi_issue, now,
                                       issue_as_root);
@@ -754,10 +760,13 @@ int blk_iolatency_init(struct request_queue *q)
         return 0;
 }
 
-static void iolatency_set_min_lat_nsec(struct blkcg_gq *blkg, u64 val)
+/*
+ * return 1 for enabling iolatency, return -1 for disabling iolatency, otherwise
+ * return 0.
+ */
+static int iolatency_set_min_lat_nsec(struct blkcg_gq *blkg, u64 val)
 {
         struct iolatency_grp *iolat = blkg_to_lat(blkg);
-        struct blk_iolatency *blkiolat = iolat->blkiolat;
         u64 oldval = iolat->min_lat_nsec;
 
         iolat->min_lat_nsec = val;
@@ -766,9 +775,10 @@ static void iolatency_set_min_lat_nsec(struct blkcg_gq *blkg, u64 val)
                                     BLKIOLATENCY_MAX_WIN_SIZE);
 
         if (!oldval && val)
-                atomic_inc(&blkiolat->enabled);
+                return 1;
         if (oldval && !val)
-                atomic_dec(&blkiolat->enabled);
+                return -1;
+        return 0;
 }
 
 static void iolatency_clear_scaling(struct blkcg_gq *blkg)
@@ -800,6 +810,7 @@ static ssize_t iolatency_set_limit(struct kernfs_open_file *of, char *buf,
         u64 lat_val = 0;
         u64 oldval;
         int ret;
+        int enable = 0;
 
         ret = blkg_conf_prep(blkcg, &blkcg_policy_iolatency, buf, &ctx);
         if (ret)
@@ -834,7 +845,12 @@ static ssize_t iolatency_set_limit(struct kernfs_open_file *of, char *buf,
         blkg = ctx.blkg;
         oldval = iolat->min_lat_nsec;
 
-        iolatency_set_min_lat_nsec(blkg, lat_val);
+        enable = iolatency_set_min_lat_nsec(blkg, lat_val);
+        if (enable) {
+                WARN_ON_ONCE(!blk_get_queue(blkg->q));
+                blkg_get(blkg);
+        }
+
         if (oldval != iolat->min_lat_nsec) {
                 iolatency_clear_scaling(blkg);
         }
@@ -842,6 +858,24 @@ static ssize_t iolatency_set_limit(struct kernfs_open_file *of, char *buf,
         ret = 0;
 out:
         blkg_conf_finish(&ctx);
+        if (ret == 0 && enable) {
+                struct iolatency_grp *tmp = blkg_to_lat(blkg);
+                struct blk_iolatency *blkiolat = tmp->blkiolat;
+
+                blk_mq_freeze_queue(blkg->q);
+
+                if (enable == 1)
+                        atomic_inc(&blkiolat->enabled);
+                else if (enable == -1)
+                        atomic_dec(&blkiolat->enabled);
+                else
+                        WARN_ON_ONCE(1);
+
+                blk_mq_unfreeze_queue(blkg->q);
+
+                blkg_put(blkg);
+                blk_put_queue(blkg->q);
+        }
         return ret ?: nbytes;
 }
 
@@ -977,8 +1011,14 @@ static void iolatency_pd_offline(struct blkg_policy_data *pd)
 {
         struct iolatency_grp *iolat = pd_to_lat(pd);
         struct blkcg_gq *blkg = lat_to_blkg(iolat);
+        struct blk_iolatency *blkiolat = iolat->blkiolat;
+        int ret;
 
-        iolatency_set_min_lat_nsec(blkg, 0);
+        ret = iolatency_set_min_lat_nsec(blkg, 0);
+        if (ret == 1)
+                atomic_inc(&blkiolat->enabled);
+        if (ret == -1)
+                atomic_dec(&blkiolat->enabled);
         iolatency_clear_scaling(blkg);
 }