[RFC PATCH 25/38] blktrace: update blk_add_driver_data()

Signed-off-by: Chaitanya Kulkarni <[email protected]>
---
 kernel/trace/blktrace.c | 15 ++++++++++++---
 1 file changed, 12 insertions(+), 3 deletions(-)

diff --git a/kernel/trace/blktrace.c b/kernel/trace/blktrace.c
index 05f2d8acc789..c8b5955163ad 100644
--- a/kernel/trace/blktrace.c
+++ b/kernel/trace/blktrace.c
@@ -1604,17 +1604,26 @@ static void blk_add_trace_rq_remap(void *ignore, struct request *rq, dev_t dev,
 void blk_add_driver_data(struct request *rq, void *data, size_t len)
 {
 	struct blk_trace *bt;
+	struct blk_trace_ext *bte;
 
 	rcu_read_lock();
 	bt = rcu_dereference(rq->q->blk_trace);
-	if (likely(!bt)) {
+	bte = rcu_dereference(rq->q->blk_trace_ext);
+	if (likely(!bt) && likely(!bte)) {
 		rcu_read_unlock();
 		return;
 	}
 
-	__blk_add_trace(bt, blk_rq_trace_sector(rq), blk_rq_bytes(rq), 0, 0,
-				BLK_TA_DRV_DATA, 0, len, data,
+	if (bt) {
+		__blk_add_trace(bt, blk_rq_trace_sector(rq), blk_rq_bytes(rq), 0,
+				0, BLK_TA_DRV_DATA, 0, len, data,
 				blk_trace_request_get_cgid(rq));
+	} else if (bte) {
+		__blk_add_trace_ext(bte, blk_rq_trace_sector(rq),
+				blk_rq_bytes(rq), 0, 0, BLK_TA_DRV_DATA_EXT, 0,
+				len, data, blk_trace_request_get_cgid(rq),
+				req_get_ioprio(rq));
+	}
 	rcu_read_unlock();
 }
 EXPORT_SYMBOL_GPL(blk_add_driver_data);
-- 
2.22.1