librbd: add rbd_diff_iterate3() API to take source snapshot by ID

Allow a diff to start from a non-user snapshot.  This would be used by
"rbd du" command to account for non-user snapshots which are currently
just skipped potentially resulting in underreported space usage and in
other places.

Fixes: https://tracker.ceph.com/issues/65720
Co-authored-by: Ilya Dryomov <idryomov@gmail.com>
Signed-off-by: Vinay Bhaskar Varada <vvarada@redhat.com>
Signed-off-by: Ilya Dryomov <idryomov@gmail.com>

src/include/rbd/librbd.h

@@ -52,6 +52,7 @@ extern "C" {
 #define LIBRBD_SUPPORTS_ENCRYPTION 1
 #define LIBRBD_SUPPORTS_ENCRYPTION_LOAD2 1
 #define LIBRBD_SUPPORTS_GROUP_SNAP_GET_INFO 1
+#define LIBRBD_SUPPORTS_DIFF_ITERATE3 1
 
 #if __GNUC__ >= 4
   #define CEPH_RBD_API          __attribute__ ((visibility ("default")))
@@ -403,6 +404,14 @@ enum {
   RBD_WRITE_ZEROES_FLAG_THICK_PROVISION = (1U<<0), /* fully allocated zeroed extent */
 };
 
+/* rbd_diff_iterate3 flags */
+enum {
+  /* full history diff should include parent */
+  RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT = (1U<<0),
+  /* diff extents should cover whole object */
+  RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT = (1U<<1),
+};
+
 typedef enum {
     RBD_ENCRYPTION_FORMAT_LUKS1 = 0,
     RBD_ENCRYPTION_FORMAT_LUKS2 = 1,
@@ -1235,6 +1244,11 @@ CEPH_RBD_API int rbd_diff_iterate2(rbd_image_t image,
                                    uint8_t include_parent, uint8_t whole_object,
 		                   int (*cb)(uint64_t, size_t, int, void *),
                                    void *arg);
+CEPH_RBD_API int rbd_diff_iterate3(rbd_image_t image, uint64_t from_snap_id,
+                                   uint64_t ofs, uint64_t len, uint32_t flags,
+                                   int (*cb)(uint64_t, size_t, int, void *),
+                                   void *arg);
+
 CEPH_RBD_API ssize_t rbd_write(rbd_image_t image, uint64_t ofs, size_t len,
                                const char *buf);
 /*

src/include/rbd/librbd.hpp

@@ -758,6 +758,9 @@ public:
 		    uint64_t ofs, uint64_t len,
                     bool include_parent, bool whole_object,
 		    int (*cb)(uint64_t, size_t, int, void *), void *arg);
+  int diff_iterate3(uint64_t from_snap_id,
+                    uint64_t ofs, uint64_t len, uint32_t flags,
+                    int (*cb)(uint64_t, size_t, int, void *), void *arg);
 
   ssize_t write(uint64_t ofs, size_t len, ceph::bufferlist& bl);
   /* @param op_flags see librados.h constants beginning with LIBRADOS_OP_FLAG */

src/librbd/api/DiffIterate.cc

@@ -196,15 +196,12 @@ int simple_diff_cb(uint64_t off, size_t len, int exists, void *arg) {
 } // anonymous namespace
 
 template <typename I>
-int DiffIterate<I>::diff_iterate(I *ictx,
-				 const cls::rbd::SnapshotNamespace& from_snap_namespace,
-				 const char *fromsnapname,
+int DiffIterate<I>::diff_iterate(I *ictx, uint64_t from_snap_id,
                                  uint64_t off, uint64_t len,
                                  bool include_parent, bool whole_object,
                                  int (*cb)(uint64_t, size_t, int, void *),
                                  void *arg) {
-  ldout(ictx->cct, 10) << "from_snap_namespace=" << from_snap_namespace
-                       << ", fromsnapname=" << (fromsnapname ?: "")
+  ldout(ictx->cct, 10) << "from_snap_id=" << from_snap_id
                        << ", off=" << off
                        << ", len=" << len
                        << ", include_parent=" << include_parent
@@ -250,7 +247,7 @@ int DiffIterate<I>::diff_iterate(I *ictx,
     // lock is acquired
     // acquire exclusive lock only if not busy (i.e. don't request),
     // throttle acquisition attempts and ignore errors
-    if (fromsnapname == nullptr && whole_object &&
+    if (from_snap_id == 0 && whole_object &&
         should_try_acquire_lock(ictx)) {
       C_SaferCond lock_ctx;
       ictx->exclusive_lock->try_acquire_lock(&lock_ctx);
@@ -260,7 +257,7 @@ int DiffIterate<I>::diff_iterate(I *ictx,
     }
   }
 
-  DiffIterate command(*ictx, from_snap_namespace, fromsnapname, off, len,
+  DiffIterate command(*ictx, from_snap_id, off, len,
 		      include_parent, whole_object, cb, arg);
   r = command.execute();
   return r;
@@ -296,24 +293,23 @@ int DiffIterate<I>::execute() {
 
   ceph_assert(m_image_ctx.data_ctx.is_valid());
 
-  librados::snap_t from_snap_id = 0;
+  librados::snap_t from_snap_id = m_from_snap_id;
   librados::snap_t end_snap_id;
   uint64_t from_size = 0;
   uint64_t end_size;
   {
     std::shared_lock image_locker{m_image_ctx.image_lock};
-    if (m_from_snap_name) {
-      from_snap_id = m_image_ctx.get_snap_id(m_from_snap_namespace,
-                                             m_from_snap_name);
-      from_size = m_image_ctx.get_image_size(from_snap_id);
+    if (from_snap_id != 0) {
+      auto info = m_image_ctx.get_snap_info(from_snap_id);
+      if (info == nullptr) {
+        return -ENOENT;
+      }
+      from_size = info->size;
     }
     end_snap_id = m_image_ctx.snap_id;
     end_size = m_image_ctx.get_image_size(end_snap_id);
   }
 
-  if (from_snap_id == CEPH_NOSNAP) {
-    return -ENOENT;
-  }
   if (from_snap_id > end_snap_id) {
     return -EINVAL;
   }
@@ -353,7 +349,7 @@ int DiffIterate<I>::execute() {
         if (m_image_ctx.prune_parent_extents(parent_extents, io::ImageArea::DATA,
                                              raw_overlap, false) > 0) {
           ldout(cct, 10) << " first getting parent diff" << dendl;
-          DiffIterate diff_parent(*m_image_ctx.parent, {}, nullptr,
+          DiffIterate diff_parent(*m_image_ctx.parent, 0,
                                   parent_extents[0].first,
                                   parent_extents[0].second, true, true,
                                   &simple_diff_cb, &parent_diff);

src/librbd/api/DiffIterate.h

@@ -20,9 +20,7 @@ class DiffIterate {
 public:
   typedef int (*Callback)(uint64_t, size_t, int, void *);
 
-  static int diff_iterate(ImageCtxT *ictx,
-			  const cls::rbd::SnapshotNamespace& from_snap_namespace,
-			  const char *fromsnapname,
+  static int diff_iterate(ImageCtxT *ictx, uint64_t from_snap_id,
                           uint64_t off, uint64_t len, bool include_parent,
                           bool whole_object,
 		          int (*cb)(uint64_t, size_t, int, void *),
@@ -30,8 +28,7 @@ public:
 
 private:
   ImageCtxT &m_image_ctx;
-  cls::rbd::SnapshotNamespace m_from_snap_namespace;
-  const char* m_from_snap_name;
+  uint64_t m_from_snap_id;
   uint64_t m_offset;
   uint64_t m_length;
   bool m_include_parent;
@@ -39,13 +36,12 @@ private:
   Callback m_callback;
   void *m_callback_arg;
 
-  DiffIterate(ImageCtxT &image_ctx, 
-	      const cls::rbd::SnapshotNamespace& from_snap_namespace,
-	      const char *from_snap_name, uint64_t off, uint64_t len,
+  DiffIterate(ImageCtxT &image_ctx, uint64_t from_snap_id,
+	      uint64_t off, uint64_t len,
 	      bool include_parent, bool whole_object, Callback callback,
 	      void *callback_arg)
-    : m_image_ctx(image_ctx), m_from_snap_namespace(from_snap_namespace),
-      m_from_snap_name(from_snap_name), m_offset(off),
+    : m_image_ctx(image_ctx),
+      m_from_snap_id(from_snap_id), m_offset(off),
       m_length(len), m_include_parent(include_parent),
       m_whole_object(whole_object), m_callback(callback),
       m_callback_arg(callback_arg)

src/librbd/api/Trash.cc

@@ -483,8 +483,7 @@ int Trash<I>::purge(IoCtx& io_ctx, time_t expire_ts,
           }
 
           r = librbd::api::DiffIterate<I>::diff_iterate(
-            ictx, cls::rbd::UserSnapshotNamespace(), nullptr, 0, ictx->size,
-            false, true,
+            ictx, 0, 0, ictx->size, false, true,
             [](uint64_t offset, size_t len, int exists, void *arg) {
                 auto *to_free = reinterpret_cast<uint64_t *>(arg);
                 if (exists)

src/librbd/librbd.cc

@@ -2688,9 +2688,13 @@ namespace librbd {
     tracepoint(librbd, diff_iterate_enter, ictx, ictx->name.c_str(),
                ictx->snap_name.c_str(), ictx->read_only, fromsnapname, ofs, len,
                true, false);
-    int r = librbd::api::DiffIterate<>::diff_iterate(ictx,
-						     cls::rbd::UserSnapshotNamespace(),
-						     fromsnapname, ofs,
+    uint64_t from_snap_id = 0;
+    if (fromsnapname != nullptr) {
+      std::shared_lock image_locker{ictx->image_lock};
+      from_snap_id = ictx->get_snap_id(cls::rbd::UserSnapshotNamespace(),
+                                       fromsnapname);
+    }
+    int r = librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id, ofs,
                                                      len, true, false, cb, arg);
     tracepoint(librbd, diff_iterate_exit, r);
     return r;
@@ -2704,15 +2708,36 @@ namespace librbd {
     tracepoint(librbd, diff_iterate_enter, ictx, ictx->name.c_str(),
               ictx->snap_name.c_str(), ictx->read_only, fromsnapname, ofs, len,
               include_parent, whole_object);
-    int r = librbd::api::DiffIterate<>::diff_iterate(ictx,
-						     cls::rbd::UserSnapshotNamespace(),
-						     fromsnapname, ofs,
-                                                     len, include_parent,
+    uint64_t from_snap_id = 0;
+    if (fromsnapname != nullptr) {
+      std::shared_lock image_locker{ictx->image_lock};
+      from_snap_id = ictx->get_snap_id(cls::rbd::UserSnapshotNamespace(),
+                                       fromsnapname);
+    }
+    int r = librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id,
+                                                     ofs, len, include_parent,
                                                      whole_object, cb, arg);
     tracepoint(librbd, diff_iterate_exit, r);
     return r;
   }
 
+  int Image::diff_iterate3(uint64_t from_snap_id,
+                           uint64_t ofs, uint64_t len, uint32_t flags,
+                           int (*cb)(uint64_t, size_t, int, void *), void *arg)
+  {
+    if ((flags & ~(RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT |
+                   RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT)) != 0) {
+      return -EINVAL;
+    }
+
+    ImageCtx *ictx = (ImageCtx *)ctx;
+    bool include_parent = flags & RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT;
+    bool whole_object = flags & RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT;
+    return librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id,
+                                                    ofs, len, include_parent,
+                                                    whole_object, cb, arg);
+  }
+
   ssize_t Image::write(uint64_t ofs, size_t len, bufferlist& bl)
   {
     ImageCtx *ictx = (ImageCtx *)ctx;
@@ -6182,10 +6207,14 @@ extern "C" int rbd_diff_iterate(rbd_image_t image,
   tracepoint(librbd, diff_iterate_enter, ictx, ictx->name.c_str(),
              ictx->snap_name.c_str(), ictx->read_only, fromsnapname, ofs, len,
              true, false);
-  int r = librbd::api::DiffIterate<>::diff_iterate(ictx,
-						   cls::rbd::UserSnapshotNamespace(),
-						   fromsnapname, ofs, len,
-                                                   true, false, cb, arg);
+  uint64_t from_snap_id = 0;
+  if (fromsnapname != nullptr) {
+    std::shared_lock image_locker{ictx->image_lock};
+    from_snap_id = ictx->get_snap_id(cls::rbd::UserSnapshotNamespace(),
+                                     fromsnapname);
+  }
+  int r = librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id, ofs,
+                                                   len, true, false, cb, arg);
   tracepoint(librbd, diff_iterate_exit, r);
   return r;
 }
@@ -6200,15 +6229,37 @@ extern "C" int rbd_diff_iterate2(rbd_image_t image, const char *fromsnapname,
   tracepoint(librbd, diff_iterate_enter, ictx, ictx->name.c_str(),
             ictx->snap_name.c_str(), ictx->read_only, fromsnapname, ofs, len,
             include_parent != 0, whole_object != 0);
-  int r = librbd::api::DiffIterate<>::diff_iterate(ictx,
-						   cls::rbd::UserSnapshotNamespace(),
-						   fromsnapname, ofs, len,
-                                                   include_parent, whole_object,
-                                                   cb, arg);
+  uint64_t from_snap_id = 0;
+  if (fromsnapname != nullptr) {
+    std::shared_lock image_locker{ictx->image_lock};
+    from_snap_id = ictx->get_snap_id(cls::rbd::UserSnapshotNamespace(),
+                                     fromsnapname);
+  }
+  int r = librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id,
+                                                   ofs, len, include_parent,
+                                                   whole_object, cb, arg);
   tracepoint(librbd, diff_iterate_exit, r);
   return r;
 }
 
+extern "C" int rbd_diff_iterate3(rbd_image_t image, uint64_t from_snap_id,
+                                 uint64_t ofs, uint64_t len, uint32_t flags,
+                                 int (*cb)(uint64_t, size_t, int, void *),
+                                 void *arg)
+{
+  if ((flags & ~(RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT |
+                 RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT)) != 0) {
+    return -EINVAL;
+  }
+
+  librbd::ImageCtx *ictx = (librbd::ImageCtx *)image;
+  bool include_parent = flags & RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT;
+  bool whole_object = flags & RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT;
+  return librbd::api::DiffIterate<>::diff_iterate(ictx, from_snap_id,
+                                                  ofs, len, include_parent,
+                                                  whole_object, cb, arg);
+}
+
 extern "C" ssize_t rbd_write(rbd_image_t image, uint64_t ofs, size_t len,
 			     const char *buf)
 {

src/pybind/rbd/c_rbd.pxd

@@ -60,6 +60,9 @@ cdef extern from "rbd/librbd.h" nogil:
 
         _RBD_WRITE_ZEROES_FLAG_THICK_PROVISION "RBD_WRITE_ZEROES_FLAG_THICK_PROVISION"
 
+        _RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT "RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT"
+        _RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT "RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT"
+
     ctypedef void* rados_t
     ctypedef void* rados_ioctx_t
     ctypedef void* rbd_image_t
@@ -628,6 +631,11 @@ cdef extern from "rbd/librbd.h" nogil:
                          int (*cb)(uint64_t, size_t, int, void *)
                              nogil except? -9000,
                          void *arg) except? -9000
+    int rbd_diff_iterate3(rbd_image_t image, uint64_t from_snap_id,
+                          uint64_t ofs, uint64_t len, uint32_t flags,
+                          int (*cb)(uint64_t, size_t, int, void *)
+                              nogil except? -9000,
+                          void *arg) except? -9000
 
     int rbd_flush(rbd_image_t image)
     int rbd_invalidate_cache(rbd_image_t image)

src/pybind/rbd/mock_rbd.pxi

@@ -64,6 +64,9 @@ cdef nogil:
 
         _RBD_WRITE_ZEROES_FLAG_THICK_PROVISION "RBD_WRITE_ZEROES_FLAG_THICK_PROVISION"
 
+        _RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT "RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT"
+        _RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT "RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT"
+
     ctypedef void* rados_t
     ctypedef void* rados_ioctx_t
     ctypedef void* rbd_image_t
@@ -779,6 +782,12 @@ cdef nogil:
                              nogil except? -9000,
                          void *arg) except? -9000:
         pass
+    int rbd_diff_iterate3(rbd_image_t image, uint64_t from_snap_id,
+                         uint64_t ofs, uint64_t len, uint32_t flags,
+                         int (*cb)(uint64_t, size_t, int, void *)
+                             nogil except? -9000,
+                         void *arg) except? -9000:
+        pass
 
     int rbd_flush(rbd_image_t image):
         pass

src/pybind/rbd/rbd.pyx

@@ -4102,15 +4102,16 @@ cdef class Image(object):
         Raises :class:`InvalidArgument` if from_snapshot is after
         the currently set snapshot.
 
-        Raises :class:`ImageNotFound` if from_snapshot is not the name
+        Raises :class:`ImageNotFound` if from_snapshot is not the name or id
         of a snapshot of the image.
 
         :param offset: start offset in bytes
         :type offset: int
         :param length: size of region to report on, in bytes
         :type length: int
-        :param from_snapshot: starting snapshot name, or None
-        :type from_snapshot: str or None
+        :param from_snapshot: starting snapshot name or id, or None to
+                              get all allocated extents
+        :type from_snapshot: str or int
         :param iterate_cb: function to call for each extent
         :type iterate_cb: function acception arguments for offset,
                            length, and exists
@@ -4121,18 +4122,45 @@ cdef class Image(object):
         :raises: :class:`InvalidArgument`, :class:`IOError`,
                  :class:`ImageNotFound`
         """
-        from_snapshot = cstr(from_snapshot, 'from_snapshot', opt=True)
         cdef:
-            char *_from_snapshot = opt_str(from_snapshot)
+            char *_from_snap_name = NULL
+            uint64_t _from_snap_id = 0
             uint64_t _offset = offset, _length = length
             uint8_t _include_parent = include_parent
             uint8_t _whole_object = whole_object
-        with nogil:
-            ret = rbd_diff_iterate2(self.image, _from_snapshot, _offset,
-                                    _length, _include_parent, _whole_object,
-                                    &diff_iterate_cb, <void *>iterate_cb)
+            uint32_t _flags = 0
+
+        if from_snapshot is not None:
+            if isinstance(from_snapshot, str):
+                from_snap_name = cstr(from_snapshot, 'from_snapshot')
+                _from_snap_name = from_snap_name
+            elif isinstance(from_snapshot, int):
+                from_snap_name = None
+                _from_snap_id = from_snapshot
+            else:
+                raise TypeError("from_snapshot must be a string or an integer")
+        else:
+            from_snap_name = None
+
+        if from_snap_name is not None:
+            with nogil:
+                ret = rbd_diff_iterate2(self.image, _from_snap_name, _offset,
+                                        _length, _include_parent, _whole_object,
+                                        &diff_iterate_cb, <void *>iterate_cb)
+        else:
+            if include_parent:
+                _flags |= _RBD_DIFF_ITERATE_FLAG_INCLUDE_PARENT
+            if whole_object:
+                _flags |= _RBD_DIFF_ITERATE_FLAG_WHOLE_OBJECT
+            with nogil:
+                ret = rbd_diff_iterate3(self.image, _from_snap_id, _offset,
+                                        _length, _flags, &diff_iterate_cb,
+                                        <void *>iterate_cb)
         if ret < 0:
-            msg = 'error generating diff from snapshot %s' % from_snapshot
+            if from_snap_name is not None:
+                msg = 'error generating diff from snapshot %s' % from_snapshot
+            else:
+                msg = 'error generating diff from snapshot id %d' % _from_snap_id
             raise make_ex(ret, msg)
 
     @requires_not_closed

src/test/librbd/test_internal.cc

@@ -1298,8 +1298,7 @@ TEST_F(TestInternal, DiffIterateCloneOverwrite) {
                                               cls::rbd::UserSnapshotNamespace(),
                                               "one"));
   ASSERT_EQ(0, librbd::api::DiffIterate<>::diff_iterate(
-    ictx, cls::rbd::UserSnapshotNamespace(), nullptr, 0, size, true, false,
-    iterate_cb, (void *)&diff));
+    ictx, 0, 0, size, true, false, iterate_cb, (void *)&diff));
   ASSERT_EQ(one, diff);
 }
 

src/test/pybind/test_rbd.py

@@ -695,14 +695,14 @@ class TestImage(object):
         self.image.write(data, 0)
         self.image.write_zeroes(0, 256)
         eq(self.image.read(256, 256), b'\0' * 256)
-        check_diff(self.image, 0, IMG_SIZE, None, [])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [])
 
     def test_write_zeroes_thick_provision(self):
         data = rand_data(256)
         self.image.write(data, 0)
         self.image.write_zeroes(0, 256, RBD_WRITE_ZEROES_FLAG_THICK_PROVISION)
         eq(self.image.read(256, 256), b'\0' * 256)
-        check_diff(self.image, 0, IMG_SIZE, None, [(0, 256, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 256, True)])
 
     def test_read(self):
         data = self.image.read(0, 20)
@@ -1394,22 +1394,146 @@ class TestImage(object):
         eq([], self.image.list_lockers())
 
     def test_diff_iterate(self):
-        check_diff(self.image, 0, IMG_SIZE, None, [])
+        def cb(offset, length, exists):
+            raise Exception()
+
+        assert_raises(TypeError, self.image.diff_iterate, 0, IMG_SIZE, 1.0, cb)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      LIBRADOS_SNAP_HEAD, cb)
+
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [])
         self.image.write(b'a' * 256, 0)
-        check_diff(self.image, 0, IMG_SIZE, None, [(0, 256, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 256, True)])
         self.image.write(b'b' * 256, 256)
-        check_diff(self.image, 0, IMG_SIZE, None, [(0, 512, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
         self.image.discard(128, 256)
-        check_diff(self.image, 0, IMG_SIZE, None, [(0, 512, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+
+        check_diff(self.image, 0, 0, None, 0, [])
+        check_diff(self.image, IMG_SIZE - 1, 0, None, 0, [])
+        check_diff(self.image, IMG_SIZE, 0, None, 0, [])
 
         self.image.create_snap('snap1')
+        snap_id1 = self.image.snap_get_id('snap1')
+        self.image.remove_snap('snap1')
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      'snap1', cb)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      snap_id1, cb)
+
+        self.image.create_snap('snap1')
+        snap_id1 = self.image.snap_get_id('snap1')
         self.image.discard(0, 1 << IMG_ORDER)
         self.image.create_snap('snap2')
+        snap_id2 = self.image.snap_get_id('snap2')
+
+        self.image.write(b'c' * 16, 16)
+        self.image.write(b'c', 1 << IMG_ORDER)
+
+        self.image.set_snap('snap1')
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0,
+                   [(0, 1 << IMG_ORDER, True)], whole_object=True)
+        check_diff(self.image, 0, IMG_SIZE, 'snap1', snap_id1, [])
+        assert_raises(InvalidArgument, self.image.diff_iterate, 0, IMG_SIZE,
+                      'snap2', cb)
+        assert_raises(InvalidArgument, self.image.diff_iterate, 0, IMG_SIZE,
+                      snap_id2, cb)
+
         self.image.set_snap('snap2')
-        check_diff(self.image, 0, IMG_SIZE, 'snap1', [(0, 512, False)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [])
+        check_diff(self.image, 0, IMG_SIZE, 'snap1', snap_id1,
+                   [(0, 512, False)])
+        check_diff(self.image, 0, IMG_SIZE, 'snap1', snap_id1,
+                   [(0, 1 << IMG_ORDER, False)], whole_object=True)
+        check_diff(self.image, 0, IMG_SIZE, 'snap2', snap_id2, [])
+
+        self.image.set_snap(None)
+        expected_whole_object = [(0, 1 << IMG_ORDER, True),
+                                 (1 << IMG_ORDER, 1 << IMG_ORDER, True)]
+        check_diff(self.image, 0, IMG_SIZE, None, 0,
+                   [(0, 32, True), (1 << IMG_ORDER, 1, True)])
+        check_diff(self.image, 0, IMG_SIZE, None, 0,
+                   expected_whole_object, whole_object=True)
+        check_diff(self.image, 0, IMG_SIZE, 'snap1', snap_id1,
+                   [(0, 32, True),
+                    (32, 480, False),
+                    (1 << IMG_ORDER, 1, True)])
+        check_diff(self.image, 0, IMG_SIZE, 'snap1', snap_id1,
+                   expected_whole_object, whole_object=True)
+        check_diff(self.image, 0, IMG_SIZE, 'snap2', snap_id2,
+                   [(0, 32, True), (1 << IMG_ORDER, 1, True)])
+        check_diff(self.image, 0, IMG_SIZE, 'snap2', snap_id2,
+                   expected_whole_object, whole_object=True)
+
         self.image.remove_snap('snap1')
         self.image.remove_snap('snap2')
 
+    @require_features([RBD_FEATURE_LAYERING])
+    def test_diff_iterate_from_trash_snap(self):
+        def cb(offset, length, exists):
+            raise Exception()
+
+        self.image.write(b'a' * 256, 0)
+        self.image.create_snap('snap')
+        snap_id = self.image.snap_get_id('snap')
+        clone_name = get_temp_image_name()
+        self.rbd.clone(ioctx, image_name, 'snap', ioctx, clone_name, features,
+                       clone_format=2)
+
+        self.image.write(b'b' * 256, 256)
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+        check_diff(self.image, 0, IMG_SIZE, 'snap', snap_id, [(256, 256, True)])
+
+        self.image.remove_snap('snap')
+        image_snaps = list(self.image.list_snaps())
+        assert [s['namespace'] for s in image_snaps] == [RBD_SNAP_NAMESPACE_TYPE_TRASH]
+        assert image_snaps[0]['id'] == snap_id
+
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      'snap', cb)
+        check_diff_one(self.image, 0, IMG_SIZE, snap_id, [(256, 256, True)])
+        self.image.write(b'c' * 256, 128)
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+        check_diff_one(self.image, 0, IMG_SIZE, snap_id, [(128, 384, True)])
+        check_diff_one(self.image, 0, IMG_SIZE, snap_id,
+                       [(0, 1 << IMG_ORDER, True)], whole_object=True)
+
+        self.rbd.remove(ioctx, clone_name)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      'snap', cb)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      snap_id, cb)
+
+    @require_features([RBD_FEATURE_LAYERING])
+    def test_diff_iterate_exclude_parent(self):
+        self.image.write(b'a' * 256, 0)
+        self.image.create_snap('snap')
+        clone_name = get_temp_image_name()
+        self.rbd.clone(ioctx, image_name, 'snap', ioctx, clone_name, features,
+                       clone_format=2)
+
+        self.image.write(b'b' * 256, 256)
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+
+        with Image(ioctx, clone_name) as clone:
+            check_diff(clone, 0, IMG_SIZE, None, 0, [(0, 256, True)])
+            clone.write(b'c' * 256, 1 << IMG_ORDER)
+            check_diff(clone, 0, IMG_SIZE, None, 0,
+                       [(0, 256, True), (1 << IMG_ORDER, 256, True)])
+            check_diff(clone, 0, IMG_SIZE, None, 0,
+                       [(0, 1 << IMG_ORDER, True),
+                        (1 << IMG_ORDER, 1 << IMG_ORDER, True)],
+                       whole_object=True)
+            check_diff(clone, 0, IMG_SIZE, None, 0,
+                       [(1 << IMG_ORDER, 256, True)], include_parent=False)
+            check_diff(clone, 0, IMG_SIZE, None, 0,
+                       [(1 << IMG_ORDER, 1 << IMG_ORDER, True)],
+                       include_parent=False, whole_object=True)
+
+        self.rbd.remove(ioctx, clone_name)
+        self.image.remove_snap('snap')
+
     def test_aio_read(self):
         # this is a list so that the local cb() can modify it
         retval = [None]
@@ -1635,13 +1759,20 @@ class TestImageId(object):
         info = self.image2.stat()
         check_stat(info, new_size, IMG_ORDER)
 
-def check_diff(image, offset, length, from_snapshot, expected):
+def check_diff_one(image, offset, length, from_snapshot, expected, **kwargs):
     extents = []
     def cb(offset, length, exists):
         extents.append((offset, length, exists))
-    image.diff_iterate(0, IMG_SIZE, from_snapshot, cb)
+    image.diff_iterate(offset, length, from_snapshot, cb, **kwargs)
     eq(extents, expected)
 
+def check_diff(image, offset, length, from_snap_name, from_snap_id, expected,
+               **kwargs):
+    assert from_snap_name is None or isinstance(from_snap_name, str)
+    assert isinstance(from_snap_id, int)
+    check_diff_one(image, offset, length, from_snap_name, expected, **kwargs)
+    check_diff_one(image, offset, length, from_snap_id, expected, **kwargs)
+
 class TestClone(object):
 
     @require_features([RBD_FEATURE_LAYERING])
@@ -3209,6 +3340,37 @@ class TestGroups(object):
 
         self.rbd.remove(ioctx, clone_name)
 
+    def test_group_snap_diff_iterate(self):
+        def cb(offset, length, exists):
+            raise Exception()
+
+        self.image.write(b'a' * 256, 0)
+        self.group.add_image(ioctx, image_name)
+        self.group.create_snap(snap_name)
+        image_snaps = list(self.image.list_snaps())
+        assert [s['namespace'] for s in image_snaps] == [RBD_SNAP_NAMESPACE_TYPE_GROUP]
+        image_snap_name = image_snaps[0]['name']
+        image_snap_id = image_snaps[0]['id']
+
+        self.image.write(b'b' * 256, 256)
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      image_snap_name, cb)
+        check_diff_one(self.image, 0, IMG_SIZE, image_snap_id,
+                       [(256, 256, True)])
+        self.image.write(b'c' * 256, 128)
+        check_diff(self.image, 0, IMG_SIZE, None, 0, [(0, 512, True)])
+        check_diff_one(self.image, 0, IMG_SIZE, image_snap_id,
+                       [(128, 384, True)])
+        check_diff_one(self.image, 0, IMG_SIZE, image_snap_id,
+                       [(0, 1 << IMG_ORDER, True)], whole_object=True)
+
+        self.group.remove_snap(snap_name)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      image_snap_name, cb)
+        assert_raises(ImageNotFound, self.image.diff_iterate, 0, IMG_SIZE,
+                      image_snap_id, cb)
+
     def test_group_snap_rollback(self):
         for _ in range(1, 3):
             create_image()