1 files changed, 10 insertions, 13 deletions

diff --git a/block/stream.c b/block/stream.c
index 999d9e5..c0616b6 100644
--- a/block/stream.c
+++ b/block/stream.c
@@ -51,7 +51,7 @@ static int coroutine_fn stream_populate(BlockBackend *blk,
     return blk_co_preadv(blk, offset, bytes, NULL, BDRV_REQ_PREFETCH);
 }
 
-static int stream_prepare(Job *job)
+static int GRAPH_UNLOCKED stream_prepare(Job *job)
 {
     StreamBlockJob *s = container_of(job, StreamBlockJob, common.job);
     BlockDriverState *unfiltered_bs;
@@ -73,18 +73,16 @@ static int stream_prepare(Job *job)
     s->cor_filter_bs = NULL;
 
     /*
-     * bdrv_set_backing_hd() requires that the unfiltered_bs and the COW child
-     * of unfiltered_bs is drained. Drain already here and use
-     * bdrv_set_backing_hd_drained() instead because the polling during
-     * drained_begin() might change the graph, and if we do this only later, we
-     * may end up working with the wrong base node (or it might even have gone
-     * away by the time we want to use it).
+     * bdrv_set_backing_hd() requires that all block nodes are drained. Drain
+     * already here, because the polling during drained_begin() might change the
+     * graph, and if we do this only later, we may end up working with the wrong
+     * base node (or it might even have gone away by the time we want to use
+     * it).
      */
-    bdrv_drained_begin(unfiltered_bs);
     if (unfiltered_bs_cow) {
         bdrv_ref(unfiltered_bs_cow);
-        bdrv_drained_begin(unfiltered_bs_cow);
     }
+    bdrv_drain_all_begin();
 
     bdrv_graph_rdlock_main_loop();
     base = bdrv_filter_or_cow_bs(s->above_base);
@@ -106,7 +104,7 @@ static int stream_prepare(Job *job)
         }
 
         bdrv_graph_wrlock();
-        bdrv_set_backing_hd_drained(unfiltered_bs, base, &local_err);
+        bdrv_set_backing_hd(unfiltered_bs, base, &local_err);
         bdrv_graph_wrunlock();
 
         /*
@@ -123,11 +121,10 @@ static int stream_prepare(Job *job)
     }
 
 out:
+    bdrv_drain_all_end();
     if (unfiltered_bs_cow) {
-        bdrv_drained_end(unfiltered_bs_cow);
         bdrv_unref(unfiltered_bs_cow);
     }
-    bdrv_drained_end(unfiltered_bs);
     return ret;
 }
 
@@ -373,7 +370,7 @@ void stream_start(const char *job_id, BlockDriverState *bs,
      * already have our own plans. Also don't allow resize as the image size is
      * queried only at the job start and then cached.
      */
-    bdrv_graph_wrlock();
+    bdrv_graph_wrlock_drained();
     if (block_job_add_bdrv(&s->common, "active node", bs, 0,
                            basic_flags | BLK_PERM_WRITE, errp)) {
         bdrv_graph_wrunlock();