librbd: integrate neorados into ImageCtx

Also create an up-to-date data_io_context that mimics the function
of ImageCtx::data_ctx. The data_io_context will eventually be passed
via the IO dispatch specs to replace the passing of the snapshot
id vectors.

Signed-off-by: Jason Dillaman <dillaman@redhat.com>

10 files changed, 96 insertions, 4 deletions

diff --git a/src/librbd/ImageCtx.cc b/src/librbd/ImageCtx.cc
index 6b9e934f9c4..f8730a40cc6 100644
--- a/src/librbd/ImageCtx.cc
+++ b/src/librbd/ImageCtx.cc
@@ -4,6 +4,8 @@
 #include <boost/assign/list_of.hpp>
 #include <stddef.h>
 
+#include "include/neorados/RADOS.hpp"
+
 #include "common/ceph_context.h"
 #include "common/dout.h"
 #include "common/errno.h"
@@ -36,7 +38,6 @@
 #include "librbd/operation/ResizeRequest.h"
 
 #include "osdc/Striper.h"
-#include <boost/bind.hpp>
 #include <boost/algorithm/string/predicate.hpp>
 
 #define dout_subsys ceph_subsys_rbd
@@ -71,6 +72,12 @@ public:
   }
 };
 
+librados::IoCtx duplicate_io_ctx(librados::IoCtx& io_ctx) {
+  librados::IoCtx dup_io_ctx;
+  dup_io_ctx.dup(io_ctx);
+  return dup_io_ctx;
+}
+
 } // anonymous namespace
 
   const string ImageCtx::METADATA_CONF_PREFIX = "conf_";
@@ -87,6 +94,9 @@ public:
       exclusive_locked(false),
       name(image_name),
       asio_engine(std::make_shared<AsioEngine>(p)),
+      rados_api(asio_engine->get_rados_api()),
+      data_ctx(duplicate_io_ctx(p)),
+      md_ctx(duplicate_io_ctx(p)),
       image_watcher(NULL),
       journal(NULL),
       owner_lock(ceph::make_shared_mutex(util::unique_lock_name("librbd::ImageCtx::owner_lock", this))),
@@ -112,11 +122,11 @@ public:
       asok_hook(nullptr),
       trace_endpoint("librbd")
   {
-    md_ctx.dup(p);
-    data_ctx.dup(p);
     if (snap)
       snap_name = snap;
 
+    rebuild_data_io_context();
+
     // FIPS zeroization audit 20191117: this memset is not security related.
     memset(&header, 0, sizeof(header));
 
@@ -316,6 +326,7 @@ public:
       snap_exists = true;
       if (data_ctx.is_valid()) {
         data_ctx.snap_set_read(snap_id);
+        rebuild_data_io_context();
       }
       return 0;
     }
@@ -331,6 +342,7 @@ public:
     snap_exists = true;
     if (data_ctx.is_valid()) {
       data_ctx.snap_set_read(snap_id);
+      rebuild_data_io_context();
     }
   }
 
@@ -893,6 +905,21 @@ public:
     journal_policy = policy;
   }
 
+  void ImageCtx::rebuild_data_io_context() {
+    auto ctx = std::make_shared<neorados::IOContext>(
+      data_ctx.get_id(), data_ctx.get_namespace());
+    ctx->read_snap(snap_id);
+    ctx->write_snap_context(
+      {{snapc.seq, {snapc.snaps.begin(), snapc.snaps.end()}}});
+
+    // atomically reset the data IOContext to new version
+    atomic_store(&data_io_context, ctx);
+  }
+
+  IOContext ImageCtx::get_data_io_context() const {
+    return atomic_load(&data_io_context);
+  }
+
   void ImageCtx::get_timer_instance(CephContext *cct, SafeTimer **timer,
                                     ceph::mutex **timer_lock) {
     auto safe_timer_singleton =
diff --git a/src/librbd/ImageCtx.h b/src/librbd/ImageCtx.h
index 43ed32cac35..af501ced216 100644
--- a/src/librbd/ImageCtx.h
+++ b/src/librbd/ImageCtx.h
@@ -39,6 +39,7 @@
 class SafeTimer;
 
 namespace neorados {
+class IOContext;
 class RADOS;
 } // namespace neorados
 
@@ -114,7 +115,13 @@ namespace librbd {
 
     std::shared_ptr<AsioEngine> asio_engine;
 
-    IoCtx data_ctx, md_ctx;
+    // New ASIO-style RADOS API
+    neorados::RADOS& rados_api;
+
+    // Legacy RADOS API
+    librados::IoCtx data_ctx;
+    librados::IoCtx md_ctx;
+
     ImageWatcher<ImageCtx> *image_watcher;
     Journal<ImageCtx> *journal;
 
@@ -347,8 +354,14 @@ namespace librbd {
     journal::Policy *get_journal_policy() const;
     void set_journal_policy(journal::Policy *policy);
 
+    void rebuild_data_io_context();
+    IOContext get_data_io_context() const;
+
     static void get_timer_instance(CephContext *cct, SafeTimer **timer,
                                    ceph::mutex **timer_lock);
+
+  private:
+    std::shared_ptr<neorados::IOContext> data_io_context;
   };
 }
 
diff --git a/src/librbd/Types.h b/src/librbd/Types.h
index 66e3a967f6f..d6d7abe9aca 100644
--- a/src/librbd/Types.h
+++ b/src/librbd/Types.h
@@ -8,8 +8,11 @@
 #include "cls/rbd/cls_rbd_types.h"
 #include "deep_copy/Types.h"
 #include <map>
+#include <memory>
 #include <string>
 
+namespace neorados { class IOContext; }
+
 namespace librbd {
 
 // Performance counters
@@ -55,6 +58,8 @@ enum {
   l_librbd_last,
 };
 
+typedef std::shared_ptr<neorados::IOContext> IOContext;
+
 typedef std::map<uint64_t, uint64_t> SnapSeqs;
 
 /// Full information about an image's parent.
diff --git a/src/librbd/image/OpenRequest.cc b/src/librbd/image/OpenRequest.cc
index ccca85ff476..227fadff2d3 100644
--- a/src/librbd/image/OpenRequest.cc
+++ b/src/librbd/image/OpenRequest.cc
@@ -483,6 +483,7 @@ Context *OpenRequest<I>::handle_v2_get_data_pool(int *result) {
       m_image_ctx->data_ctx.close();
     } else {
       m_image_ctx->data_ctx.set_namespace(m_image_ctx->md_ctx.get_namespace());
+      m_image_ctx->rebuild_data_io_context();
     }
   } else {
     data_pool_id = m_image_ctx->md_ctx.get_id();
diff --git a/src/librbd/image/RefreshRequest.cc b/src/librbd/image/RefreshRequest.cc
index d70e026e36a..c1ccde4000e 100644
--- a/src/librbd/image/RefreshRequest.cc
+++ b/src/librbd/image/RefreshRequest.cc
@@ -1385,6 +1385,7 @@ void RefreshRequest<I>::apply() {
   if (m_image_ctx.data_ctx.is_valid()) {
     m_image_ctx.data_ctx.selfmanaged_snap_set_write_ctx(m_image_ctx.snapc.seq,
                                                         m_image_ctx.snaps);
+    m_image_ctx.rebuild_data_io_context();
   }
 
   // handle dynamically enabled / disabled features
diff --git a/src/librbd/operation/SnapshotCreateRequest.cc b/src/librbd/operation/SnapshotCreateRequest.cc
index 13885cf3011..344dcf781b2 100644
--- a/src/librbd/operation/SnapshotCreateRequest.cc
+++ b/src/librbd/operation/SnapshotCreateRequest.cc
@@ -433,6 +433,7 @@ void SnapshotCreateRequest<I>::update_snap_context() {
   image_ctx.snapc.snaps.swap(snaps);
   image_ctx.data_ctx.selfmanaged_snap_set_write_ctx(
     image_ctx.snapc.seq, image_ctx.snaps);
+  image_ctx.rebuild_data_io_context();
 
   if (!image_ctx.migration_info.empty()) {
     auto it = image_ctx.migration_info.snap_map.find(CEPH_NOSNAP);
diff --git a/src/test/librbd/image/test_mock_RefreshRequest.cc b/src/test/librbd/image/test_mock_RefreshRequest.cc
index 68c17d42ae3..c032cb7f791 100644
--- a/src/test/librbd/image/test_mock_RefreshRequest.cc
+++ b/src/test/librbd/image/test_mock_RefreshRequest.cc
@@ -558,6 +558,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessV1) {
   expect_v1_get_snapshots(mock_image_ctx, 0);
   expect_v1_get_locks(mock_image_ctx, 0);
   expect_init_layout(mock_image_ctx);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -583,6 +584,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessSnapshotV1) {
   expect_v1_get_locks(mock_image_ctx, 0);
   expect_init_layout(mock_image_ctx);
   expect_add_snap(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -617,6 +619,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessV2) {
   if (ictx->test_features(RBD_FEATURE_EXCLUSIVE_LOCK)) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -654,6 +657,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessSnapshotV2) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
   expect_add_snap(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -693,6 +697,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessLegacySnapshotV2) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
   expect_add_snap(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -733,6 +738,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessLegacySnapshotNoTimestampV2) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
   expect_add_snap(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -775,6 +781,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessSetSnapshotV2) {
   }
   expect_add_snap(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
   expect_get_snap_id(mock_image_ctx, "snap", ictx->snap_ids.begin()->second);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -832,6 +839,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessChild) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
   expect_refresh_parent_apply(*mock_refresh_parent_request);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   expect_refresh_parent_finalize(mock_image_ctx, *mock_refresh_parent_request, 0);
 
   C_SaferCond ctx;
@@ -886,6 +894,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessChildDontOpenParent) {
   if (ictx->test_features(RBD_FEATURE_EXCLUSIVE_LOCK)) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, true, &ctx);
@@ -923,6 +932,7 @@ TEST_F(TestMockImageRefreshRequest, SuccessOpFeatures) {
   if (ictx->test_features(RBD_FEATURE_EXCLUSIVE_LOCK)) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -988,6 +998,7 @@ TEST_F(TestMockImageRefreshRequest, DisableExclusiveLock) {
   expect_apply_metadata(mock_image_ctx, 0);
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   expect_shut_down_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
 
   C_SaferCond ctx;
@@ -1042,6 +1053,7 @@ TEST_F(TestMockImageRefreshRequest, DisableExclusiveLockWhileAcquiringLock) {
   expect_apply_metadata(mock_image_ctx, 0);
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, true, false, &ctx);
@@ -1090,6 +1102,7 @@ TEST_F(TestMockImageRefreshRequest, JournalDisabledByPolicy) {
   MockJournalPolicy mock_journal_policy;
   expect_get_journal_policy(mock_image_ctx, mock_journal_policy);
   expect_journal_disabled(mock_journal_policy, true);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1140,6 +1153,7 @@ TEST_F(TestMockImageRefreshRequest, EnableJournalWithExclusiveLock) {
   expect_get_journal_policy(mock_image_ctx, mock_journal_policy);
   expect_journal_disabled(mock_journal_policy, false);
   expect_open_journal(mock_image_ctx, mock_journal, 0);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1184,6 +1198,7 @@ TEST_F(TestMockImageRefreshRequest, EnableJournalWithoutExclusiveLock) {
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
   expect_set_require_lock(mock_exclusive_lock, librbd::io::DIRECTION_BOTH);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1235,6 +1250,7 @@ TEST_F(TestMockImageRefreshRequest, DisableJournal) {
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
   expect_block_writes(mock_image_ctx, 0);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   if (!mock_image_ctx.clone_copy_on_read) {
     expect_unset_require_lock(mock_exclusive_lock, librbd::io::DIRECTION_READ);
   }
@@ -1287,6 +1303,7 @@ TEST_F(TestMockImageRefreshRequest, EnableObjectMapWithExclusiveLock) {
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
   expect_open_object_map(mock_image_ctx, &mock_object_map, 0);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1330,6 +1347,7 @@ TEST_F(TestMockImageRefreshRequest, EnableObjectMapWithoutExclusiveLock) {
   expect_apply_metadata(mock_image_ctx, 0);
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1380,6 +1398,7 @@ TEST_F(TestMockImageRefreshRequest, DisableObjectMap) {
   expect_apply_metadata(mock_image_ctx, 0);
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   expect_close_object_map(mock_image_ctx, mock_object_map, 0);
 
   C_SaferCond ctx;
@@ -1427,6 +1446,7 @@ TEST_F(TestMockImageRefreshRequest, OpenObjectMapError) {
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
   expect_open_object_map(mock_image_ctx, &mock_object_map, -EBLACKLISTED);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false,
@@ -1475,6 +1495,7 @@ TEST_F(TestMockImageRefreshRequest, OpenObjectMapTooLarge) {
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
   expect_open_object_map(mock_image_ctx, &mock_object_map, -EFBIG);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false,
@@ -1511,6 +1532,7 @@ TEST_F(TestMockImageRefreshRequest, ApplyMetadataError) {
   if (ictx->test_features(RBD_FEATURE_EXCLUSIVE_LOCK)) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx;
   MockRefreshRequest *req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx);
@@ -1545,6 +1567,7 @@ TEST_F(TestMockImageRefreshRequest, NonPrimaryFeature) {
   expect_apply_metadata(mock_image_ctx, 0);
   expect_get_group(mock_image_ctx, 0);
   expect_refresh_parent_is_required(mock_refresh_parent_request, false);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx1;
   auto req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx1);
@@ -1570,6 +1593,7 @@ TEST_F(TestMockImageRefreshRequest, NonPrimaryFeature) {
   if (ictx->test_features(RBD_FEATURE_EXCLUSIVE_LOCK)) {
     expect_init_exclusive_lock(mock_image_ctx, mock_exclusive_lock, 0);
   }
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
 
   C_SaferCond ctx2;
   req = new MockRefreshRequest(mock_image_ctx, false, false, &ctx2);
diff --git a/src/test/librbd/mock/MockImageCtx.cc b/src/test/librbd/mock/MockImageCtx.cc
index 9a1c5430daa..a2670f896dc 100644
--- a/src/test/librbd/mock/MockImageCtx.cc
+++ b/src/test/librbd/mock/MockImageCtx.cc
@@ -1,6 +1,7 @@
 // -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:t -*-
 // vim: ts=8 sw=2 smarttab
 
+#include "include/neorados/RADOS.hpp"
 #include "test/librbd/mock/MockImageCtx.h"
 #include "test/librbd/mock/MockSafeTimer.h"
 #include "librbd/io/AsyncOperation.h"
@@ -35,4 +36,13 @@ void MockImageCtx::wait_for_async_ops() {
   async_op.finish_op();
 }
 
+IOContext MockImageCtx::get_data_io_context() {
+  auto ctx = std::make_shared<neorados::IOContext>(
+    data_ctx.get_id(), data_ctx.get_namespace());
+  ctx->read_snap(snap_id);
+  ctx->write_snap_context(
+    {{snapc.seq, {snapc.snaps.begin(), snapc.snaps.end()}}});
+  return ctx;
+}
+
 } // namespace librbd
diff --git a/src/test/librbd/mock/MockImageCtx.h b/src/test/librbd/mock/MockImageCtx.h
index ab8b0696a1f..a06ce09e6a5 100644
--- a/src/test/librbd/mock/MockImageCtx.h
+++ b/src/test/librbd/mock/MockImageCtx.h
@@ -63,6 +63,7 @@ struct MockImageCtx {
       exclusive_locked(image_ctx.exclusive_locked),
       lock_tag(image_ctx.lock_tag),
       asio_engine(image_ctx.asio_engine),
+      rados_api(image_ctx.rados_api),
       owner_lock(image_ctx.owner_lock),
       image_lock(image_ctx.image_lock),
       timestamp_lock(image_ctx.timestamp_lock),
@@ -220,6 +221,9 @@ struct MockImageCtx {
   MOCK_CONST_METHOD0(get_stripe_count, uint64_t());
   MOCK_CONST_METHOD0(get_stripe_period, uint64_t());
 
+  MOCK_METHOD0(rebuild_data_io_context, void());
+  IOContext get_data_io_context();
+
   static void set_timer_instance(MockSafeTimer *timer, ceph::mutex *timer_lock);
   static void get_timer_instance(CephContext *cct, MockSafeTimer **timer,
                                  ceph::mutex **timer_lock);
@@ -251,6 +255,7 @@ struct MockImageCtx {
   std::string lock_tag;
 
   std::shared_ptr<AsioEngine> asio_engine;
+  neorados::RADOS& rados_api;
 
   librados::IoCtx md_ctx;
   librados::IoCtx data_ctx;
diff --git a/src/test/librbd/operation/test_mock_SnapshotCreateRequest.cc b/src/test/librbd/operation/test_mock_SnapshotCreateRequest.cc
index fb0b1f1dc48..2aed5403c34 100644
--- a/src/test/librbd/operation/test_mock_SnapshotCreateRequest.cc
+++ b/src/test/librbd/operation/test_mock_SnapshotCreateRequest.cc
@@ -179,6 +179,7 @@ TEST_F(TestMockOperationSnapshotCreateRequest, Success) {
   if (!mock_image_ctx.old_format) {
     expect_object_map_snap_create(mock_image_ctx);
     expect_update_snap_context(mock_image_ctx);
+    EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   }
   expect_unblock_writes(mock_image_ctx);
   expect_notify_unquiesce(mock_image_ctx, -EINVAL);
@@ -278,6 +279,7 @@ TEST_F(TestMockOperationSnapshotCreateRequest, CreateSnapStale) {
   if (!mock_image_ctx.old_format) {
     expect_object_map_snap_create(mock_image_ctx);
     expect_update_snap_context(mock_image_ctx);
+    EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   }
   expect_unblock_writes(mock_image_ctx);
   expect_notify_unquiesce(mock_image_ctx, 0);
@@ -387,6 +389,7 @@ TEST_F(TestMockOperationSnapshotCreateRequest, SkipObjectMap) {
   expect_allocate_snap_id(mock_image_ctx, 0);
   expect_snap_create(mock_image_ctx, 0);
   expect_update_snap_context(mock_image_ctx);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   expect_unblock_writes(mock_image_ctx);
   expect_notify_unquiesce(mock_image_ctx, 0);
 
@@ -430,6 +433,7 @@ TEST_F(TestMockOperationSnapshotCreateRequest, SkipNotifyQuiesce) {
   if (!mock_image_ctx.old_format) {
     expect_object_map_snap_create(mock_image_ctx);
     expect_update_snap_context(mock_image_ctx);
+    EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   }
   expect_unblock_writes(mock_image_ctx);
 
@@ -475,6 +479,7 @@ TEST_F(TestMockOperationSnapshotCreateRequest, SetImageState) {
   MockSetImageStateRequest mock_set_image_state_request;
   expect_set_image_state(mock_image_ctx, mock_set_image_state_request, 0);
   expect_update_snap_context(mock_image_ctx);
+  EXPECT_CALL(mock_image_ctx, rebuild_data_io_context());
   expect_unblock_writes(mock_image_ctx);
   expect_notify_unquiesce(mock_image_ctx, 0);