3 files changed, 24 insertions, 6 deletions

diff --git a/block/blk-core.c b/block/blk-core.c
index 3689319a5974..500c080a6a6b 100644
--- a/block/blk-core.c
+++ b/block/blk-core.c
@@ -64,13 +64,27 @@ static void drive_stat_acct(struct request *rq, int new_io)
                 return;
 
         cpu = part_stat_lock();
-        part = disk_map_sector_rcu(rq->rq_disk, blk_rq_pos(rq));
 
-        if (!new_io)
+        if (!new_io) {
+                part = rq->part;
                 part_stat_inc(cpu, part, merges[rw]);
-        else {
+        } else {
+                part = disk_map_sector_rcu(rq->rq_disk, blk_rq_pos(rq));
+                if (!kref_test_and_get(&part->ref)) {
+                        /*
+                         * The partition is already being removed,
+                         * the request will be accounted on the disk only
+                         *
+                         * We take a reference on disk->part0 although that
+                         * partition will never be deleted, so we can treat
+                         * it as any other partition.
+                         */
+                        part = &rq->rq_disk->part0;
+                        kref_get(&part->ref);
+                }
                 part_round_stats(cpu, part);
                 part_inc_in_flight(part, rw);
+                rq->part = part;
         }
 
         part_stat_unlock();
@@ -128,6 +142,7 @@ void blk_rq_init(struct request_queue *q, struct request *rq)
         rq->ref_count = 1;
         rq->start_time = jiffies;
         set_start_time_ns(rq);
+        rq->part = NULL;
 }
 EXPORT_SYMBOL(blk_rq_init);
 
@@ -1776,7 +1791,7 @@ static void blk_account_io_completion(struct request *req, unsigned int bytes)
                 int cpu;
 
                 cpu = part_stat_lock();
-                part = disk_map_sector_rcu(req->rq_disk, blk_rq_pos(req));
+                part = req->part;
                 part_stat_add(cpu, part, sectors[rw], bytes >> 9);
                 part_stat_unlock();
         }
@@ -1796,13 +1811,14 @@ static void blk_account_io_done(struct request *req)
                 int cpu;
 
                 cpu = part_stat_lock();
-                part = disk_map_sector_rcu(req->rq_disk, blk_rq_pos(req));
+                part = req->part;
 
                 part_stat_inc(cpu, part, ios[rw]);
                 part_stat_add(cpu, part, ticks[rw], duration);
                 part_round_stats(cpu, part);
                 part_dec_in_flight(part, rw);
 
+                kref_put(&part->ref, __delete_partition);
                 part_stat_unlock();
         }
 }
diff --git a/block/blk-merge.c b/block/blk-merge.c
index 77b7c26df6b5..b06b83b89d89 100644
--- a/block/blk-merge.c
+++ b/block/blk-merge.c
@@ -351,11 +351,12 @@ static void blk_account_io_merge(struct request *req)
                 int cpu;
 
                 cpu = part_stat_lock();
-                part = disk_map_sector_rcu(req->rq_disk, blk_rq_pos(req));
+                part = req->part;
 
                 part_round_stats(cpu, part);
                 part_dec_in_flight(part, rq_data_dir(req));
 
+                kref_put(&part->ref, __delete_partition);
                 part_stat_unlock();
         }
 }
diff --git a/block/genhd.c b/block/genhd.c
index 16ccc0d2d5d3..85c150598830 100644
--- a/block/genhd.c
+++ b/block/genhd.c
@@ -1192,6 +1192,7 @@ struct gendisk *alloc_disk_node(int minors, int node_id)
                         return NULL;
                 }
                 disk->part_tbl->part[0] = &disk->part0;
+                kref_init(&disk->part0.ref);
 
                 disk->minors = minors;
                 rand_initialize_disk(disk);