re-work SG dirty tracking (#672)

Move SG dirtying to vfu_unmap_sg(): as we don't want to track SGs
ourselves, doing this in vfu_map_sg() is no longer the right place.

Note that the lack of tracking implies that any SGs must be unmapped
before the final stop and copy phase. To avoid the need for this, add
vfu_mark_sg_dirty(): this allows a consumer to mark a region as dirty
explicitly without needing to unmap it. Currently it's the same as
vfu_unmap_sg(), but that's an implementation detail.

Note this still marks current maps after a get operation; that will
change subsequently.

Signed-off-by: John Levon <john.levon@nutanix.com>
Reviewed-by: Thanos Makatos <thanos.makatos@nutanix.com>

4 files changed, 91 insertions, 37 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index 629727f..3fe4b9f 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -803,9 +803,22 @@ vfu_map_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, struct iovec *iov, int cnt,
            int flags);
 
 /**
+ * Mark scatter/gather entries (previously mapped by vfu_map_sg()) as dirty
+ * (written to). This is only necessary if vfu_unmap_sg() is not called.
+ *
+ * @vfu_ctx: the libvfio-user context
+ * @sg: array of scatter/gather entries to mark as dirty
+ * @cnt: number of scatter/gather entries to mark as dirty
+ */
+void
+vfu_mark_sg_dirty(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, int cnt);
+
+/**
  * Unmaps scatter/gather entries (previously mapped by vfu_map_sg()) from
  * the process's virtual memory.
  *
+ * This will automatically mark the sg as dirty if needed.
+ *
  * @vfu_ctx: the libvfio-user context
  * @sg: array of scatter/gather entries to unmap
  * @iov: array of iovec structures for each scatter/gather entry
diff --git a/lib/dma.h b/lib/dma.h
index 190a016..afd7345 100644
--- a/lib/dma.h
+++ b/lib/dma.h
@@ -156,6 +156,7 @@ _dma_mark_dirty(const dma_controller_t *dma, const dma_memory_region_t *region,
 
     start = sg->offset / dma->dirty_pgsize;
     end = start + (sg->length / dma->dirty_pgsize) + (sg->length % dma->dirty_pgsize != 0) - 1;
+
     for (i = start; i <= end; i++) {
         region->dirty_bitmap[i / CHAR_BIT] |= 1 << (i % CHAR_BIT);
     }
@@ -246,11 +247,9 @@ dma_map_sg(dma_controller_t *dma, dma_sg_t *sg, struct iovec *iov,
         }
 
         if (sg->writeable) {
-            if (dma->dirty_pgsize > 0) {
-                _dma_mark_dirty(dma, region, sg);
-            }
             LIST_INSERT_HEAD(&dma->maps, sg, entry);
         }
+
         vfu_log(dma->vfu_ctx, LOG_DEBUG, "map %p-%p",
                 sg->dma_addr + sg->offset,
                 sg->dma_addr + sg->offset + sg->length);
@@ -266,13 +265,39 @@ dma_map_sg(dma_controller_t *dma, dma_sg_t *sg, struct iovec *iov,
 }
 
 static inline void
-dma_unmap_sg(dma_controller_t *dma, const dma_sg_t *sg,
-	     UNUSED struct iovec *iov, int cnt)
+dma_mark_sg_dirty(dma_controller_t *dma, dma_sg_t *sg, int cnt)
 {
+    dma_memory_region_t *region;
 
     assert(dma != NULL);
     assert(sg != NULL);
-    assert(iov != NULL);
+    assert(cnt > 0);
+
+    do {
+        if (sg->region >= dma->nregions) {
+            return;
+        }
+
+        region = &dma->regions[sg->region];
+
+        if (sg->writeable) {
+            if (dma->dirty_pgsize > 0) {
+                _dma_mark_dirty(dma, region, sg);
+            }
+        }
+
+        vfu_log(dma->vfu_ctx, LOG_DEBUG, "mark dirty %p-%p",
+                sg->dma_addr + sg->offset,
+                sg->dma_addr + sg->offset + sg->length);
+        sg++;
+    } while (--cnt > 0);
+}
+
+static inline void
+dma_unmap_sg(dma_controller_t *dma, dma_sg_t *sg, int cnt)
+{
+    assert(dma != NULL);
+    assert(sg != NULL);
     assert(cnt > 0);
 
     do {
@@ -289,9 +314,15 @@ dma_unmap_sg(dma_controller_t *dma, const dma_sg_t *sg,
             /* bad region */
             continue;
         }
+
         if (sg->writeable) {
             LIST_REMOVE(sg, entry);
+
+            if (dma->dirty_pgsize > 0) {
+                _dma_mark_dirty(dma, r, sg);
+            }
         }
+
         vfu_log(dma->vfu_ctx, LOG_DEBUG, "unmap %p-%p",
                 sg->dma_addr + sg->offset,
                 sg->dma_addr + sg->offset + sg->length);
diff --git a/lib/libvfio-user.c b/lib/libvfio-user.c
index 924d109..7d324aa 100644
--- a/lib/libvfio-user.c
+++ b/lib/libvfio-user.c
@@ -2038,7 +2038,7 @@ vfu_map_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, struct iovec *iov, int cnt,
 }
 
 EXPORT void
-vfu_unmap_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, struct iovec *iov, int cnt)
+vfu_mark_sg_dirty(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, int cnt)
 {
     if (unlikely(vfu_ctx->dma_unregister == NULL)) {
         return;
@@ -2046,7 +2046,20 @@ vfu_unmap_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg, struct iovec *iov, int cnt)
 
     quiesce_check_allowed(vfu_ctx);
 
-    return dma_unmap_sg(vfu_ctx->dma, sg, iov, cnt);
+    return dma_mark_sg_dirty(vfu_ctx->dma, sg, cnt);
+}
+
+EXPORT void
+vfu_unmap_sg(vfu_ctx_t *vfu_ctx, dma_sg_t *sg,
+             struct iovec *iov UNUSED, int cnt)
+{
+    if (unlikely(vfu_ctx->dma_unregister == NULL)) {
+        return;
+    }
+
+    quiesce_check_allowed(vfu_ctx);
+
+    return dma_unmap_sg(vfu_ctx->dma, sg, cnt);
 }
 
 static int
diff --git a/test/py/test_dirty_pages.py b/test/py/test_dirty_pages.py
index 8af325a..9a0de3b 100644
--- a/test/py/test_dirty_pages.py
+++ b/test/py/test_dirty_pages.py
@@ -318,6 +318,7 @@ def test_dirty_pages_get_modified():
     ret = vfu_map_sg(ctx, sg1, iovec1)
     assert ret == 0
 
+    # read only
     ret, sg2 = vfu_addr_to_sg(ctx, dma_addr=0x11000, length=0x1000,
                               prot=mmap.PROT_READ)
     assert ret == 1
@@ -325,52 +326,47 @@ def test_dirty_pages_get_modified():
     ret = vfu_map_sg(ctx, sg2, iovec2)
     assert ret == 0
 
-    global sg3, iovec3
-    ret, sg3 = vfu_addr_to_sg(ctx, dma_addr=0x14000, length=0x4000)
+    ret, sg3 = vfu_addr_to_sg(ctx, dma_addr=0x12000, length=0x1000)
     assert ret == 1
     iovec3 = iovec_t()
     ret = vfu_map_sg(ctx, sg3, iovec3)
     assert ret == 0
 
+    ret, sg4 = vfu_addr_to_sg(ctx, dma_addr=0x14000, length=0x4000)
+    assert ret == 1
+    iovec4 = iovec_t()
+    ret = vfu_map_sg(ctx, sg4, iovec4)
+    assert ret == 0
+
+    # not unmapped yet, dirty bitmap should be zero, but dirty maps will have
+    # been marked dirty still
     bitmap = get_dirty_page_bitmap()
-    assert bitmap == 0b11110001
+    assert bitmap == 0b00000000
 
-    # unmap segment, dirty bitmap should be the same
+    # unmap segments, dirty bitmap should be updated
     vfu_unmap_sg(ctx, sg1, iovec1)
+    vfu_unmap_sg(ctx, sg4, iovec4)
     bitmap = get_dirty_page_bitmap()
-    assert bitmap == 0b11110001
+    assert bitmap == 0b11110101
 
-    # check again, previously unmapped segment should be clean
+    # after another two unmaps, should just be one dirty page
+    vfu_unmap_sg(ctx, sg2, iovec2)
+    vfu_unmap_sg(ctx, sg3, iovec3)
     bitmap = get_dirty_page_bitmap()
-    assert bitmap == 0b11110000
-
+    assert bitmap == 0b00000100
 
-def stop_logging():
-    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
-                                    flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP)
-
-    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload)
+    # and should now be clear
+    bitmap = get_dirty_page_bitmap()
+    assert bitmap == 0b00000000
 
 
 def test_dirty_pages_stop():
-    stop_logging()
-
-    # one segment is still mapped, starting logging again and bitmap should be
-    # dirty
-    start_logging()
-    assert get_dirty_page_bitmap() == 0b11110000
-
-    # unmap segment, bitmap should still be dirty
-    vfu_unmap_sg(ctx, sg3, iovec3)
-    assert get_dirty_page_bitmap() == 0b11110000
-
-    # bitmap should be clear after it was unmapped before previous request for
-    # dirty pages
-    assert get_dirty_page_bitmap() == 0b00000000
-
     # FIXME we have a memory leak as we don't free dirty bitmaps when
     # destroying the context.
-    stop_logging()
+    payload = vfio_user_dirty_pages(argsz=len(vfio_user_dirty_pages()),
+                                    flags=VFIO_IOMMU_DIRTY_PAGES_FLAG_STOP)
+
+    msg(ctx, sock, VFIO_USER_DIRTY_PAGES, payload)
 
 
 def test_dirty_pages_start_with_quiesce():
@@ -402,6 +398,7 @@ def test_dirty_pages_bitmap_with_quiesce():
     iovec1 = iovec_t()
     ret = vfu_map_sg(ctx, sg1, iovec1)
     assert ret == 0
+    vfu_unmap_sg(ctx, sg1, iovec1)
 
     send_dirty_page_bitmap(busy=True)