librbd/migration: add nbd stream

Co-authored-by: Ilya Dryomov <idryomov@gmail.com>
Signed-off-by: Effi Ofer <effio@il.ibm.com>
Signed-off-by: Ilya Dryomov <idryomov@gmail.com>

CMakeLists.txt

@@ -259,6 +259,12 @@ if(WITH_RBD AND LINUX)
   set(HAVE_LIBCRYPTSETUP ${LIBCRYPTSETUP_FOUND})
 endif()
 
+# libnbd
+if(WITH_RBD AND NOT WIN32)
+  find_package(libnbd 1.0 REQUIRED)
+  set(HAVE_LIBNBD ${LIBNBD_FOUND})
+endif()
+
 include(CMakeDependentOption)
 
 CMAKE_DEPENDENT_OPTION(WITH_LIBURING "Enable io_uring bluestore backend" ON

ceph.spec.in

@@ -286,6 +286,7 @@ BuildRequires:	gperftools-devel >= 2.4
 BuildRequires:	libaio-devel
 BuildRequires:	libblkid-devel >= 2.17
 BuildRequires:	cryptsetup-devel
+BuildRequires:	libnbd-devel
 BuildRequires:	libcurl-devel
 BuildRequires:	libcap-devel
 BuildRequires:	libcap-ng-devel

cmake/modules/Findlibnbd.cmake

@@ -0,0 +1,33 @@
+# - Find libnbd
+# Sets the following:
+#
+# LIBNBD_INCLUDE_DIR
+# LIBNBD_LIBRARIES
+# LIBNBD_VERSION
+# LIBNBD_FOUND
+
+find_package(PkgConfig QUIET REQUIRED)
+pkg_search_module(PC_libnbd libnbd)
+
+find_path(LIBNBD_INCLUDE_DIR
+        NAMES libnbd.h
+        PATHS ${PC_libnbd_INCLUDE_DIRS})
+
+find_library(LIBNBD_LIBRARIES
+        NAMES libnbd.so
+        PATHS ${PC_libnbd_LIBRARY_DIRS})
+
+set(LIBNBD_VERSION ${PC_libnbd_VERSION})
+
+include(FindPackageHandleStandardArgs)
+
+find_package_handle_standard_args(libnbd
+        REQUIRED_VARS
+        LIBNBD_INCLUDE_DIR
+        LIBNBD_LIBRARIES
+        VERSION_VAR LIBNBD_VERSION)
+
+mark_as_advanced(
+  LIBNBD_LIBRARIES
+  LIBNBD_INCLUDE_DIR
+  LIBNBD_VERSION)

debian/control

@@ -54,6 +54,7 @@ Build-Depends: automake,
                liblttng-ust-dev,
                liblua5.3-dev,
                liblz4-dev (>= 0.0~r131),
+               libnbd-dev,
                libncurses-dev,
                libnss3-dev,
                liboath-dev,

doc/rbd/rbd-live-migration.rst

@@ -20,7 +20,7 @@ parent.
 The live-migration process can also be used in an import-only mode where the
 source image remains unmodified and the target image can be linked to an image
 in another Ceph cluster or to an external data source such as a backing file,
-HTTP(s) file, or S3 object.
+HTTP(s) file, S3 object, or NBD export.
 
 The live-migration copy process can safely run in the background while the new
 target image is in use. There is currently a requirement to temporarily stop
@@ -145,8 +145,8 @@ The general format for the ``source-spec`` JSON is as follows::
         }
 
 The following formats are currently supported: ``native``, ``qcow``, and
-``raw``. The following streams are currently supported: ``file``, ``http``, and
-``s3``.
+``raw``. The following streams are currently supported: ``file``, ``http``,
+``s3``, and ``nbd``.
 
 Formats
 ~~~~~~~
@@ -306,6 +306,19 @@ as follows::
   stored in the config-key store via ``ceph config-key set <key-path> <value>``
   (e.g. ``ceph config-key set rbd/s3/access_key NX5QOQKC6BH2IDN8HC7A``).
 
+The ``nbd`` stream can be used to import from a remote NBD export. Its
+``source-spec`` JSON is encoded as follows::
+
+        {
+            <format unique parameters>
+            "stream": {
+                "type": "nbd",
+                "server": "<server>",
+                "port": "<port>"
+            }
+        }
+
+
 Execute Migration
 =================
 

qa/workunits/rbd/cli_migration.sh

@@ -8,6 +8,7 @@ IMAGE3=image3
 IMAGES="${IMAGE1} ${IMAGE2} ${IMAGE3}"
 
 cleanup() {
+    kill_nbd_server
     cleanup_tempdir
     remove_images
 }
@@ -65,6 +66,10 @@ remove_images() {
     done
 }
 
+kill_nbd_server() {
+    pkill -9 qemu-nbd || true
+}
+
 show_diff()
 {
     local file1=$1
@@ -390,6 +395,42 @@ EOF
     remove_image "${dest_image}"
 }
 
+test_import_nbd_stream() {
+    local base_image=$1
+    local dest_image=$2
+
+    qemu-nbd -f qcow2 --read-only --shared 10 --persistent --fork \
+        ${TEMPDIR}/${base_image}.qcow2
+
+    cat > ${TEMPDIR}/spec.json <<EOF
+{
+  "type": "raw",
+  "stream": {
+    "type": "nbd",
+    "server": "localhost",
+    "port": "10809"
+  }
+}
+EOF
+    cat ${TEMPDIR}/spec.json
+
+    cat ${TEMPDIR}/spec.json | rbd migration prepare --import-only \
+        --source-spec-path - ${dest_image}
+    compare_images ${base_image} ${dest_image}
+    rbd migration abort ${dest_image}
+
+    rbd migration prepare --import-only \
+        --source-spec-path ${TEMPDIR}/spec.json ${dest_image}
+    compare_images ${base_image} ${dest_image}
+    rbd migration execute ${dest_image}
+    compare_images ${base_image} ${dest_image}
+    rbd migration commit ${dest_image}
+    compare_images ${base_image} ${dest_image}
+    remove_image "${dest_image}"
+
+    kill_nbd_server
+}
+
 # make sure rbd pool is EMPTY.. this is a test script!!
 rbd ls 2>&1 | wc -l | grep -v '^0$' && echo "nonempty rbd pool, aborting!  run this script on an empty test cluster only." && exit 1
 
@@ -401,7 +442,10 @@ export_base_image ${IMAGE1}
 
 test_import_native_format ${IMAGE1} ${IMAGE2}
 test_import_qcow_format ${IMAGE1} ${IMAGE2}
+
 test_import_qcow2_format ${IMAGE2} ${IMAGE3}
+test_import_nbd_stream ${IMAGE2} ${IMAGE3}
+
 test_import_raw_format ${IMAGE1} ${IMAGE2}
 
 echo OK

src/include/config-h.in.cmake

@@ -393,6 +393,9 @@
 /* Define if libcryptsetup can be used (linux only) */
 #cmakedefine HAVE_LIBCRYPTSETUP
 
+/* Define if libnbd can be used */
+#cmakedefine HAVE_LIBNBD
+
 /* Shared library extension, such as .so, .dll or .dylib */
 #cmakedefine CMAKE_SHARED_LIBRARY_SUFFIX "@CMAKE_SHARED_LIBRARY_SUFFIX@"
 

src/librbd/CMakeLists.txt

@@ -219,6 +219,11 @@ if(LINUX AND HAVE_LIBCRYPTSETUP)
           crypto/luks/Magic.cc)
 endif()
 
+if(HAVE_LIBNBD)
+  list(APPEND librbd_internal_srcs
+          migration/NBDStream.cc)
+endif()
+
 add_library(rbd_api STATIC librbd.cc)
 add_library(rbd_internal STATIC
   ${librbd_internal_srcs}
@@ -240,6 +245,10 @@ if(LINUX AND HAVE_LIBCRYPTSETUP)
   target_include_directories(rbd_internal PRIVATE ${LIBCRYPTSETUP_INCLUDE_DIR})
   target_link_libraries(rbd_internal PRIVATE ${LIBCRYPTSETUP_LIBRARIES})
 endif()
+if(HAVE_LIBNBD)
+  target_include_directories(rbd_internal PRIVATE ${LIBNBD_INCLUDE_DIR})
+  target_link_libraries(rbd_internal PRIVATE ${LIBNBD_LIBRARIES})
+endif()
 
 add_custom_target(librbd_plugins)
 set(librbd_plugins_dir ${CEPH_INSTALL_PKGLIBDIR}/librbd)

src/librbd/migration/FileStream.cc

@@ -226,6 +226,18 @@ void FileStream<I>::read(io::Extents&& byte_extents, bufferlist* data,
 
 #endif // BOOST_ASIO_HAS_POSIX_STREAM_DESCRIPTOR
 
+template <typename I>
+void FileStream<I>::list_sparse_extents(io::Extents&& byte_extents,
+                                        io::SparseExtents* sparse_extents,
+                                        Context* on_finish) {
+  // TODO: list sparse extents based on SEEK_HOLE/SEEK_DATA
+  for (auto [byte_offset, byte_length] : byte_extents) {
+    sparse_extents->insert(byte_offset, byte_length,
+                           {io::SPARSE_EXTENT_STATE_DATA, byte_length});
+  }
+  on_finish->complete(0);
+}
+
 } // namespace migration
 } // namespace librbd
 

src/librbd/migration/FileStream.h

@@ -44,6 +44,10 @@ public:
   void read(io::Extents&& byte_extents, bufferlist* data,
             Context* on_finish) override;
 
+  void list_sparse_extents(io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents,
+                           Context* on_finish) override;
+
 private:
   CephContext* m_cct;
   std::shared_ptr<AsioEngine> m_asio_engine;

src/librbd/migration/HttpStream.cc

@@ -77,6 +77,18 @@ void HttpStream<I>::read(io::Extents&& byte_extents, bufferlist* data,
   m_http_client->read(std::move(byte_extents), data, on_finish);
 }
 
+template <typename I>
+void HttpStream<I>::list_sparse_extents(io::Extents&& byte_extents,
+                                        io::SparseExtents* sparse_extents,
+                                        Context* on_finish) {
+  // no sparseness information -- list the full range as DATA
+  for (auto [byte_offset, byte_length] : byte_extents) {
+    sparse_extents->insert(byte_offset, byte_length,
+                           {io::SPARSE_EXTENT_STATE_DATA, byte_length});
+  }
+  on_finish->complete(0);
+}
+
 } // namespace migration
 } // namespace librbd
 

src/librbd/migration/HttpStream.h

@@ -45,6 +45,10 @@ public:
   void read(io::Extents&& byte_extents, bufferlist* data,
             Context* on_finish) override;
 
+  void list_sparse_extents(io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents,
+                           Context* on_finish) override;
+
 private:
   using HttpResponse = boost::beast::http::response<
     boost::beast::http::string_body>;

src/librbd/migration/NBDStream.cc

@@ -0,0 +1,275 @@
+// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:t -*-
+// vim: ts=8 sw=2 smarttab
+
+#include "librbd/migration/NBDStream.h"
+#include "common/dout.h"
+#include "common/errno.h"
+#include "librbd/AsioEngine.h"
+#include "librbd/ImageCtx.h"
+
+#include <libnbd.h>
+
+namespace librbd {
+namespace migration {
+
+namespace {
+
+const std::string SERVER_KEY {"server"};
+const std::string PORT_KEY {"port"};
+
+int extent_cb(void* data, const char* metacontext, uint64_t offset,
+              uint32_t* entries, size_t nr_entries, int* error) {
+  auto sparse_extents = reinterpret_cast<io::SparseExtents*>(data);
+
+  uint64_t length = 0;
+  for (size_t i=0; i<nr_entries; i+=2) {
+    length += entries[i];
+  }
+  auto state = io::SPARSE_EXTENT_STATE_DATA;
+  if (nr_entries == 2) {
+    if (entries[1] & (LIBNBD_STATE_HOLE | LIBNBD_STATE_ZERO)) {
+      state = io::SPARSE_EXTENT_STATE_ZEROED;
+    }
+  }
+  sparse_extents->insert(offset, length, {state, length});
+  return 1;
+}
+
+} // anonymous namespace
+
+#define dout_subsys ceph_subsys_rbd
+#undef dout_prefix
+#define dout_prefix *_dout << "librbd::migration::NBDStream::ReadRequest: " \
+                           << this << " " << __func__ << ": "
+
+template <typename I>
+struct NBDStream<I>::ReadRequest {
+  NBDStream* nbd_stream;
+  io::Extents byte_extents;
+  bufferlist* data;
+  Context* on_finish;
+  size_t index = 0;
+
+  ReadRequest(NBDStream* nbd_stream, io::Extents&& byte_extents,
+              bufferlist* data, Context* on_finish)
+    : nbd_stream(nbd_stream), byte_extents(std::move(byte_extents)),
+      data(data), on_finish(on_finish) {
+    auto cct = nbd_stream->m_cct;
+    ldout(cct, 20) << dendl;
+  }
+
+  void send() {
+    data->clear();
+    read();
+  }
+
+  void read() {
+    if (index >= byte_extents.size()) {
+      finish(0);
+      return;
+    }
+
+    auto cct = nbd_stream->m_cct;
+    auto [byte_offset, byte_length] = byte_extents[index++];
+    ldout(cct, 20) << "byte_offset=" << byte_offset << " byte_length="
+                   << byte_length << dendl;
+
+    auto ptr = buffer::ptr_node::create(buffer::create_small_page_aligned(
+      byte_length));
+    int rc = nbd_pread(nbd_stream->m_nbd, ptr->c_str(), byte_length,
+                       byte_offset, 0);
+    if (rc == -1) {
+      rc = nbd_get_errno();
+      lderr(cct) << "pread " << byte_offset << "~" << byte_length << ": "
+                 << nbd_get_error() << " (errno = " << rc << ")"
+                 << dendl;
+      finish(rc);
+      return;
+    }
+
+    data->push_back(std::move(ptr));
+    boost::asio::post(nbd_stream->m_strand, [this] { read(); });
+  }
+
+  void finish(int r) {
+    auto cct = nbd_stream->m_cct;
+    ldout(cct, 20) << "r=" << r << dendl;
+
+    if (r < 0) {
+      data->clear();
+    }
+
+    on_finish->complete(r);
+    delete this;
+  }
+};
+
+#undef dout_prefix
+#define dout_prefix *_dout << "librbd::migration::NBDStream::ListSparseExtentsRequest: " \
+                           << this << " " << __func__ << ": "
+
+template <typename I>
+struct NBDStream<I>::ListSparseExtentsRequest {
+  NBDStream* nbd_stream;
+  io::Extents byte_extents;
+  io::SparseExtents* sparse_extents;
+  Context* on_finish;
+  size_t index = 0;
+
+  ListSparseExtentsRequest(NBDStream* nbd_stream, io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents, Context* on_finish)
+    : nbd_stream(nbd_stream), byte_extents(std::move(byte_extents)),
+      sparse_extents(sparse_extents), on_finish(on_finish) {
+    auto cct = nbd_stream->m_cct;
+    ldout(cct, 20) << dendl;
+  }
+
+  void send() {
+    list_sparse_extents();
+  }
+
+  void list_sparse_extents() {
+    if (index >= byte_extents.size()) {
+      finish(0);
+      return;
+    }
+
+    auto cct = nbd_stream->m_cct;
+    auto [byte_offset, byte_length] = byte_extents[index++];
+    ldout(cct, 20) << "byte_offset=" << byte_offset << " byte_length="
+                   << byte_length << dendl;
+
+    int rc = nbd_block_status(nbd_stream->m_nbd, byte_length, byte_offset,
+                              {extent_cb, sparse_extents}, 0);
+    if (rc == -1) {
+      rc = nbd_get_errno();
+      lderr(cct) << "block_status " << byte_offset << "~" << byte_length << ": "
+                 << nbd_get_error() << " (errno = " << rc << ")"
+                 << dendl;
+      finish(rc);
+      return;
+    }
+
+    boost::asio::post(nbd_stream->m_strand, [this] { list_sparse_extents(); });
+  }
+
+  void finish(int r) {
+    auto cct = nbd_stream->m_cct;
+    ldout(cct, 20) << "r=" << r << dendl;
+
+    on_finish->complete(r);
+    delete this;
+  }
+};
+
+#undef dout_prefix
+#define dout_prefix *_dout << "librbd::migration::NBDStream: " \
+                           << this << " " << __func__ << ": "
+
+template <typename I>
+NBDStream<I>::NBDStream(I* image_ctx, const json_spirit::mObject& json_object)
+  : m_cct(image_ctx->cct), m_asio_engine(image_ctx->asio_engine),
+    m_json_object(json_object),
+    m_strand(boost::asio::make_strand(*m_asio_engine)) {
+}
+
+template <typename I>
+NBDStream<I>::~NBDStream() {
+  if (m_nbd != nullptr) {
+    nbd_close(m_nbd);
+  }
+}
+
+template <typename I>
+void NBDStream<I>::open(Context* on_finish) {
+  int rc;
+
+  auto& server_value = m_json_object[SERVER_KEY];
+  if (server_value.type() != json_spirit::str_type) {
+    lderr(m_cct) << "failed to locate '" << SERVER_KEY << "' key" << dendl;
+    on_finish->complete(-EINVAL);
+    return;
+  }
+
+  auto& port_value = m_json_object[PORT_KEY];
+  if (port_value.type() != json_spirit::str_type) {
+    lderr(m_cct) << "failed to locate '" << PORT_KEY << "' key" << dendl;
+    on_finish->complete(-EINVAL);
+    return;
+  }
+
+  const char *m_server = &(server_value.get_str())[0];
+  const char *m_port = &(port_value.get_str())[0];
+
+  m_nbd = nbd_create();
+  if (m_nbd == nullptr) {
+    lderr(m_cct) << "failed to create nbd object '" << dendl;
+    on_finish->complete(-EINVAL);
+    return;
+  }
+
+  rc = nbd_add_meta_context(m_nbd, LIBNBD_CONTEXT_BASE_ALLOCATION);
+  if (rc == -1) {
+    lderr(m_cct) << "failed to add nbd meta context '" << dendl;
+    on_finish->complete(-EINVAL);
+    return;
+  }
+
+  rc = nbd_connect_tcp(m_nbd, m_server, m_port);
+  if (rc == -1) {
+    rc = nbd_get_errno();
+    lderr(m_cct) << "failed to connect to nbd server: " << nbd_get_error()
+                 << " (errno=" << rc << ")" << dendl;
+    on_finish->complete(rc);
+    return;
+  }
+
+  ldout(m_cct, 20) << "server=" << m_server << ", "
+                   << "port=" << m_port << dendl;
+
+  on_finish->complete(0);
+}
+
+template <typename I>
+void NBDStream<I>::close(Context* on_finish) {
+  ldout(m_cct, 20) << dendl;
+
+  if (m_nbd != nullptr) {
+    nbd_close(m_nbd);
+    m_nbd = nullptr;
+  }
+
+  on_finish->complete(0);
+}
+
+template <typename I>
+void NBDStream<I>::get_size(uint64_t* size, Context* on_finish) {
+  ldout(m_cct, 20) << dendl;
+
+  *size = nbd_get_size(m_nbd);
+  on_finish->complete(0);
+}
+
+template <typename I>
+void NBDStream<I>::read(io::Extents&& byte_extents,
+                        bufferlist* data,
+                        Context* on_finish) {
+  ldout(m_cct, 20) << byte_extents << dendl;
+  auto ctx = new ReadRequest(this, std::move(byte_extents), data, on_finish);
+  boost::asio::post(m_strand, [ctx] { ctx->send(); });
+}
+
+template <typename I>
+void NBDStream<I>::list_sparse_extents(io::Extents&& byte_extents,
+                                       io::SparseExtents* sparse_extents,
+                                       Context* on_finish) {
+  ldout(m_cct, 20) << byte_extents << dendl;
+  auto ctx = new ListSparseExtentsRequest(this, std::move(byte_extents),
+                                          sparse_extents, on_finish);
+  boost::asio::post(m_strand, [ctx] { ctx->send(); });
+}
+
+} // namespace migration
+} // namespace librbd
+
+template class librbd::migration::NBDStream<librbd::ImageCtx>;

src/librbd/migration/NBDStream.h

@@ -0,0 +1,67 @@
+// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:t -*-
+// vim: ts=8 sw=2 smarttab
+
+#ifndef CEPH_LIBRBD_MIGRATION_NBD_STREAM_H
+#define CEPH_LIBRBD_MIGRATION_NBD_STREAM_H
+
+#include "include/int_types.h"
+#include "librbd/migration/StreamInterface.h"
+#include <json_spirit/json_spirit.h>
+#include <boost/asio/io_context.hpp>
+#include <boost/asio/strand.hpp>
+
+struct Context;
+
+struct nbd_handle;
+
+namespace librbd {
+
+struct AsioEngine;
+struct ImageCtx;
+
+namespace migration {
+
+template <typename ImageCtxT>
+class NBDStream : public StreamInterface {
+public:
+  static NBDStream* create(ImageCtxT* image_ctx,
+                           const json_spirit::mObject& json_object) {
+    return new NBDStream(image_ctx, json_object);
+  }
+
+  NBDStream(ImageCtxT* image_ctx, const json_spirit::mObject& json_object);
+  ~NBDStream() override;
+
+  NBDStream(const NBDStream&) = delete;
+  NBDStream& operator=(const NBDStream&) = delete;
+
+  void open(Context* on_finish) override;
+  void close(Context* on_finish) override;
+
+  void get_size(uint64_t* size, Context* on_finish) override;
+
+  void read(io::Extents&& byte_extents, bufferlist* data,
+            Context* on_finish) override;
+
+  void list_sparse_extents(io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents,
+                           Context* on_finish) override;
+
+private:
+  CephContext* m_cct;
+  std::shared_ptr<AsioEngine> m_asio_engine;
+  json_spirit::mObject m_json_object;
+  boost::asio::strand<boost::asio::io_context::executor_type> m_strand;
+
+  struct nbd_handle* m_nbd = nullptr;
+
+  struct ReadRequest;
+  struct ListSparseExtentsRequest;
+};
+
+} // namespace migration
+} // namespace librbd
+
+extern template class librbd::migration::NBDStream<librbd::ImageCtx>;
+
+#endif // CEPH_LIBRBD_MIGRATION_NBD_STREAM_H

src/librbd/migration/RawSnapshot.cc

@@ -205,13 +205,9 @@ void RawSnapshot<I>::list_snap(io::Extents&& image_extents,
   auto cct = m_image_ctx->cct;
   ldout(cct, 20) << "image_extents=" << image_extents << dendl;
 
-  // raw does support sparse extents so list the full IO extent as a delta
-  for (auto& [image_offset, image_length] : image_extents) {
-    sparse_extents->insert(image_offset, image_length,
-                           {io::SPARSE_EXTENT_STATE_DATA, image_length});
-  }
-
-  on_finish->complete(0);
+  // raw directly maps the image-extent IO down to a byte IO extent
+  m_stream->list_sparse_extents(std::move(image_extents), sparse_extents,
+                                on_finish);
 }
 
 } // namespace migration

src/librbd/migration/S3Stream.cc

@@ -194,6 +194,18 @@ void S3Stream<I>::process_request(HttpRequest& http_request) {
                    << "authorization=" << authorization << dendl;
 }
 
+template <typename I>
+void S3Stream<I>::list_sparse_extents(io::Extents&& byte_extents,
+                                      io::SparseExtents* sparse_extents,
+                                      Context* on_finish) {
+  // no sparseness information -- list the full range as DATA
+  for (auto [byte_offset, byte_length] : byte_extents) {
+    sparse_extents->insert(byte_offset, byte_length,
+                           {io::SPARSE_EXTENT_STATE_DATA, byte_length});
+  }
+  on_finish->complete(0);
+}
+
 } // namespace migration
 } // namespace librbd
 

src/librbd/migration/S3Stream.h

@@ -46,6 +46,10 @@ public:
   void read(io::Extents&& byte_extents, bufferlist* data,
             Context* on_finish) override;
 
+  void list_sparse_extents(io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents,
+                           Context* on_finish) override;
+
 private:
   using HttpRequest = boost::beast::http::request<
     boost::beast::http::empty_body>;

src/librbd/migration/SourceSpecBuilder.cc

@@ -7,6 +7,9 @@
 #include "librbd/migration/FileStream.h"
 #include "librbd/migration/HttpStream.h"
 #include "librbd/migration/S3Stream.h"
+#if defined(HAVE_LIBNBD)
+#include "librbd/migration/NBDStream.h"
+#endif
 #include "librbd/migration/NativeFormat.h"
 #include "librbd/migration/QCOWFormat.h"
 #include "librbd/migration/RawFormat.h"
@@ -125,6 +128,10 @@ int SourceSpecBuilder<I>::build_stream(
     stream->reset(HttpStream<I>::create(m_image_ctx, stream_obj));
   } else if (type == "s3") {
     stream->reset(S3Stream<I>::create(m_image_ctx, stream_obj));
+#if defined(HAVE_LIBNBD)
+  } else if (type == "nbd") {
+    stream->reset(NBDStream<I>::create(m_image_ctx, stream_obj));
+#endif
   } else {
     lderr(cct) << "unknown or unsupported stream type '" << type << "'"
                << dendl;

src/librbd/migration/StreamInterface.h

@@ -24,6 +24,10 @@ struct StreamInterface {
 
   virtual void read(io::Extents&& byte_extents, bufferlist* data,
                     Context* on_finish) = 0;
+
+  virtual void list_sparse_extents(io::Extents&& byte_extents,
+                                   io::SparseExtents* sparse_extents,
+                                   Context* on_finish) = 0;
 };
 
 } // namespace migration

src/test/librbd/CMakeLists.txt

@@ -153,6 +153,11 @@ if(LINUX AND HAVE_LIBCRYPTSETUP)
           crypto/luks/test_mock_LoadRequest.cc)
 endif()
 
+if(HAVE_LIBNBD)
+  list(APPEND unittest_librbd_srcs
+          migration/test_mock_NBDStream.cc)
+endif()
+
 # On Windows, we'll skip librbd unit tests for the time being, running just the
 # functional tests. The reason is that the unit tests require libcls*, which in
 # turn requires libos and libosd, however those libraries haven't been ported to

src/test/librbd/migration/test_mock_FileStream.cc

@@ -209,5 +209,33 @@ TEST_F(TestMockMigrationFileStream, ShortReadError) {
   ASSERT_EQ(0, ctx3.wait());
 }
 
+TEST_F(TestMockMigrationFileStream, ListSparseExtents) {
+  MockTestImageCtx mock_image_ctx(*m_image_ctx);
+
+  bufferlist bl;
+  ASSERT_EQ(0, bl.write_file(file_name.c_str()));
+
+  MockFileStream mock_file_stream(&mock_image_ctx, json_object);
+
+  C_SaferCond ctx1;
+  mock_file_stream.open(&ctx1);
+  ASSERT_EQ(0, ctx1.wait());
+
+  C_SaferCond ctx2;
+  io::SparseExtents sparse_extents;
+  mock_file_stream.list_sparse_extents({{0, 128}, {256, 64}}, &sparse_extents,
+                                       &ctx2);
+  ASSERT_EQ(0, ctx2.wait());
+
+  io::SparseExtents expected_sparse_extents;
+  expected_sparse_extents.insert(0, 128, {io::SPARSE_EXTENT_STATE_DATA, 128});
+  expected_sparse_extents.insert(256, 64, {io::SPARSE_EXTENT_STATE_DATA, 64});
+  ASSERT_EQ(expected_sparse_extents, sparse_extents);
+
+  C_SaferCond ctx3;
+  mock_file_stream.close(&ctx3);
+  ASSERT_EQ(0, ctx3.wait());
+}
+
 } // namespace migration
 } // namespace librbd

src/test/librbd/migration/test_mock_HttpStream.cc

@@ -190,5 +190,36 @@ TEST_F(TestMockMigrationHttpStream, Read) {
   ASSERT_EQ(0, ctx3.wait());
 }
 
+TEST_F(TestMockMigrationHttpStream, ListSparseExtents) {
+  MockTestImageCtx mock_image_ctx(*m_image_ctx);
+
+  InSequence seq;
+
+  auto mock_http_client = new MockHttpClient();
+  expect_open(*mock_http_client, 0);
+  expect_close(*mock_http_client, 0);
+
+  MockHttpStream mock_http_stream(&mock_image_ctx, json_object);
+
+  C_SaferCond ctx1;
+  mock_http_stream.open(&ctx1);
+  ASSERT_EQ(0, ctx1.wait());
+
+  C_SaferCond ctx2;
+  io::SparseExtents sparse_extents;
+  mock_http_stream.list_sparse_extents({{0, 128}, {256, 64}}, &sparse_extents,
+                                       &ctx2);
+  ASSERT_EQ(0, ctx2.wait());
+
+  io::SparseExtents expected_sparse_extents;
+  expected_sparse_extents.insert(0, 128, {io::SPARSE_EXTENT_STATE_DATA, 128});
+  expected_sparse_extents.insert(256, 64, {io::SPARSE_EXTENT_STATE_DATA, 64});
+  ASSERT_EQ(expected_sparse_extents, sparse_extents);
+
+  C_SaferCond ctx3;
+  mock_http_stream.close(&ctx3);
+  ASSERT_EQ(0, ctx3.wait());
+}
+
 } // namespace migration
 } // namespace librbd

src/test/librbd/migration/test_mock_NBDStream.cc

@@ -0,0 +1,62 @@
+// -*- mode:C++; tab-width:8; c-basic-offset:2; indent-tabs-mode:t -*-
+// vim: ts=8 sw=2 smarttab
+
+#include "test/librbd/test_mock_fixture.h"
+#include "test/librbd/test_support.h"
+#include "include/rbd_types.h"
+#include "librbd/migration/NBDStream.h"
+#include "gtest/gtest.h"
+#include "gmock/gmock.h"
+#include "json_spirit/json_spirit.h"
+
+namespace librbd {
+namespace {
+
+struct MockTestImageCtx : public MockImageCtx {
+  MockTestImageCtx(ImageCtx &image_ctx) : MockImageCtx(image_ctx) {
+  }
+};
+
+} // anonymous namespace
+} // namespace librbd
+
+#include "librbd/migration/NBDStream.cc"
+
+namespace librbd {
+namespace migration {
+
+using ::testing::Invoke;
+
+class TestMockMigrationNBDStream : public TestMockFixture {
+public:
+  typedef NBDStream<MockTestImageCtx> MockNBDStream;
+
+  void SetUp() override {
+    TestMockFixture::SetUp();
+
+    ASSERT_EQ(0, open_image(m_image_name, &m_image_ctx));
+    json_object["url"] = "localhost";
+    json_object["port"] = "10809";
+  }
+
+  librbd::ImageCtx *m_image_ctx;
+  json_spirit::mObject json_object;
+};
+
+TEST_F(TestMockMigrationNBDStream, OpenClose) {
+  MockTestImageCtx mock_image_ctx(*m_image_ctx);
+
+  MockNBDStream mock_nbd_stream(&mock_image_ctx, json_object);
+
+  C_SaferCond ctx1;
+  mock_nbd_stream.open(&ctx1);
+  // Since we don't have an nbd server running, we actually expect a failure.
+  ASSERT_EQ(-22, ctx1.wait());
+
+  C_SaferCond ctx2;
+  mock_nbd_stream.close(&ctx2);
+  ASSERT_EQ(0, ctx2.wait());
+}
+
+} // namespace migration
+} // namespace librbd

src/test/librbd/migration/test_mock_RawSnapshot.cc

@@ -105,6 +105,19 @@ public:
         })));
   }
 
+  void expect_stream_list_sparse_extents(MockStreamInterface& mock_stream_interface,
+                                         const io::Extents& byte_extents,
+                                         const io::SparseExtents& sparse_extents,
+                                         int r) {
+    EXPECT_CALL(mock_stream_interface, list_sparse_extents(byte_extents, _, _))
+      .WillOnce(WithArgs<1, 2>(Invoke(
+        [sparse_extents, r](io::SparseExtents* out_sparse_extents,
+                            Context* ctx) {
+          out_sparse_extents->insert(sparse_extents);
+          ctx->complete(r);
+        })));
+  }
+
   json_spirit::mObject json_object;
 };
 
@@ -232,6 +245,11 @@ TEST_F(TestMockMigrationRawSnapshot, ListSnap) {
   expect_stream_open(*mock_stream_interface, 0);
   expect_stream_get_size(*mock_stream_interface, 0, 0);
 
+  io::SparseExtents expected_sparse_extents;
+  expected_sparse_extents.insert(0, 123, {io::SPARSE_EXTENT_STATE_DATA, 123});
+  expect_stream_list_sparse_extents(*mock_stream_interface, {{0, 123}},
+                                    expected_sparse_extents, 0);
+
   expect_stream_close(*mock_stream_interface, 0);
 
   MockRawSnapshot mock_raw_snapshot(&mock_image_ctx, json_object,
@@ -245,6 +263,7 @@ TEST_F(TestMockMigrationRawSnapshot, ListSnap) {
   io::SparseExtents sparse_extents;
   mock_raw_snapshot.list_snap({{0, 123}}, 0, &sparse_extents, {}, &ctx2);
   ASSERT_EQ(0, ctx2.wait());
+  ASSERT_EQ(expected_sparse_extents, sparse_extents);
 
   C_SaferCond ctx3;
   mock_raw_snapshot.close(&ctx3);

src/test/librbd/migration/test_mock_S3Stream.cc

@@ -234,5 +234,36 @@ TEST_F(TestMockMigrationS3Stream, ProcessRequest) {
   ASSERT_EQ(0, ctx2.wait());
 }
 
+TEST_F(TestMockMigrationS3Stream, ListSparseExtents) {
+  MockTestImageCtx mock_image_ctx(*m_image_ctx);
+
+  InSequence seq;
+
+  auto mock_http_client = new MockHttpClient();
+  expect_open(*mock_http_client, 0);
+  expect_close(*mock_http_client, 0);
+
+  MockS3Stream mock_s3_stream(&mock_image_ctx, json_object);
+
+  C_SaferCond ctx1;
+  mock_s3_stream.open(&ctx1);
+  ASSERT_EQ(0, ctx1.wait());
+
+  C_SaferCond ctx2;
+  io::SparseExtents sparse_extents;
+  mock_s3_stream.list_sparse_extents({{0, 128}, {256, 64}}, &sparse_extents,
+                                     &ctx2);
+  ASSERT_EQ(0, ctx2.wait());
+
+  io::SparseExtents expected_sparse_extents;
+  expected_sparse_extents.insert(0, 128, {io::SPARSE_EXTENT_STATE_DATA, 128});
+  expected_sparse_extents.insert(256, 64, {io::SPARSE_EXTENT_STATE_DATA, 64});
+  ASSERT_EQ(expected_sparse_extents, sparse_extents);
+
+  C_SaferCond ctx3;
+  mock_s3_stream.close(&ctx3);
+  ASSERT_EQ(0, ctx3.wait());
+}
+
 } // namespace migration
 } // namespace librbd

src/test/librbd/mock/migration/MockStreamInterface.h

@@ -21,6 +21,14 @@ struct MockStreamInterface : public StreamInterface {
   void read(io::Extents&& byte_extents, bufferlist* bl, Context* on_finish) {
     read(byte_extents, bl, on_finish);
   }
+
+  MOCK_METHOD3(list_sparse_extents, void(const io::Extents&,
+                                         io::SparseExtents*, Context*));
+  void list_sparse_extents(io::Extents&& byte_extents,
+                           io::SparseExtents* sparse_extents,
+                           Context* on_finish) {
+    list_sparse_extents(byte_extents, sparse_extents, on_finish);
+  }
 };
 
 } // namespace migration