2 files changed, 6 insertions, 3 deletions

diff --git a/block/blk-merge.c b/block/blk-merge.c
index 39f2e077a014..bef1b4d0fc02 100644
--- a/block/blk-merge.c
+++ b/block/blk-merge.c
@@ -220,7 +220,7 @@ new_segment:
                 bvprv = bvec;
         } /* segments in rq */
 
-        if (q->dma_drain_size) {
+        if (q->dma_drain_size && q->dma_drain_needed(rq)) {
                 sg->page_link &= ~0x02;
                 sg = sg_next(sg);
                 sg_set_page(sg, virt_to_page(q->dma_drain_buffer),
diff --git a/block/blk-settings.c b/block/blk-settings.c
index 13536a388d27..9a8ffdd0ce3d 100644
--- a/block/blk-settings.c
+++ b/block/blk-settings.c
@@ -296,6 +296,7 @@ EXPORT_SYMBOL(blk_queue_stack_limits);
  * blk_queue_dma_drain - Set up a drain buffer for excess dma.
  *
  * @q:  the request queue for the device
+ * @dma_drain_needed: fn which returns non-zero if drain is necessary
  * @buf:        physically contiguous buffer
  * @size:       size of the buffer in bytes
  *
@@ -315,14 +316,16 @@ EXPORT_SYMBOL(blk_queue_stack_limits);
  * device can support otherwise there won't be room for the drain
  * buffer.
  */
-int blk_queue_dma_drain(struct request_queue *q, void *buf,
-                                unsigned int size)
+extern int blk_queue_dma_drain(struct request_queue *q,
+                               dma_drain_needed_fn *dma_drain_needed,
+                               void *buf, unsigned int size)
 {
         if (q->max_hw_segments < 2 || q->max_phys_segments < 2)
                 return -EINVAL;
         /* make room for appending the drain */
         --q->max_hw_segments;
         --q->max_phys_segments;
+        q->dma_drain_needed = dma_drain_needed;
         q->dma_drain_buffer = buf;
         q->dma_drain_size = size;