fix: minor changes to samples according to John's review

Signed-off-by: William Henderson <william.henderson@nutanix.com>

2 files changed, 38 insertions, 44 deletions

diff --git a/samples/client.c b/samples/client.c
index 6890b81..8779c03 100644
--- a/samples/client.c
+++ b/samples/client.c
@@ -531,7 +531,7 @@ set_migration_state(int sock, uint32_t state)
             .iov_len = sizeof(change_state)
         }
     };
-    void* response = malloc(sizeof(req) + sizeof(change_state));
+    void *response = malloc(sizeof(req) + sizeof(change_state));
 
     if (response == NULL) {
         return -1;
@@ -804,7 +804,7 @@ get_dirty_bitmap(int sock, struct vfio_user_dma_map *dma_region)
 
     size_t size = sizeof(*res) + sizeof(*report) + bitmap_size;
 
-    void* data = calloc(1, size);
+    void *data = calloc(1, size);
     assert(data != NULL);
 
     res = data;
@@ -877,10 +877,10 @@ do_migrate(int sock, size_t nr_iters, size_t max_iter_size,
             err(EXIT_FAILURE, "failed to read migration data");
         }
 
-        // We know we've finished transferring data when less is returned
-        // than is requested.
-        if ((size_t)ret < migr_iter[i].iov_len) {
-            migr_iter[i].iov_len = ret;
+        migr_iter[i].iov_len = ret;
+
+        // We know we've finished transferring data when we read 0 bytes.
+        if (ret == 0) {
             is_more = false;
         }
     }
@@ -1164,7 +1164,7 @@ int main(int argc, char *argv[])
     struct vfio_user_device_feature_dma_logging_control *dirty_pages_control;
     size_t dirty_pages_size = sizeof(*dirty_pages_feature) +
                                sizeof(*dirty_pages_control);
-    void* dirty_pages = malloc(dirty_pages_size);
+    void *dirty_pages = malloc(dirty_pages_size);
     dirty_pages_feature = dirty_pages;
     dirty_pages_control = dirty_pages + sizeof(*dirty_pages_feature);
 
diff --git a/samples/server.c b/samples/server.c
index 3a02246..fdce13a 100644
--- a/samples/server.c
+++ b/samples/server.c
@@ -62,7 +62,7 @@ struct server_data {
     size_t bar1_size;
     struct dma_regions regions[NR_DMA_REGIONS];
     struct {
-        uint64_t pending_bytes;
+        uint64_t bytes_transferred;
         vfu_migr_state_t state;
     } migration;
 };
@@ -132,10 +132,6 @@ bar1_access(vfu_ctx_t *vfu_ctx, char * const buf,
     }
 
     if (is_write) {
-        if (server_data->migration.state == VFU_MIGR_STATE_PRE_COPY) {
-            /* dirty the whole thing */
-            server_data->migration.pending_bytes = server_data->bar1_size;
-        }
         memcpy(server_data->bar1 + offset, buf, count);
     } else {
         memcpy(buf, server_data->bar1, count);
@@ -274,20 +270,24 @@ migration_device_state_transition(vfu_ctx_t *vfu_ctx, vfu_migr_state_t state)
             if (setitimer(ITIMER_REAL, &new, NULL) != 0) {
                 err(EXIT_FAILURE, "failed to disable timer");
             }
-            server_data->migration.pending_bytes = server_data->bar1_size + sizeof(time_t); /* FIXME BAR0 region size */
+            server_data->migration.bytes_transferred = 0;
             break;
         case VFU_MIGR_STATE_PRE_COPY:
-            server_data->migration.pending_bytes = server_data->bar1_size;
+            server_data->migration.bytes_transferred = 0;
             break;
         case VFU_MIGR_STATE_STOP:
             /* FIXME should gracefully fail */
-            assert(server_data->migration.pending_bytes == 0);
+            if (server_data->migration.state == VFU_MIGR_STATE_STOP_AND_COPY) {
+                assert(server_data->migration.bytes_transferred ==
+                       server_data->bar1_size + sizeof(time_t));
+            }
             break;
         case VFU_MIGR_STATE_RESUME:
-            server_data->migration.pending_bytes = server_data->bar1_size + sizeof(time_t);
+            server_data->migration.bytes_transferred = 0;
             break;
         case VFU_MIGR_STATE_RUNNING:
-            assert(server_data->migration.pending_bytes == 0);
+            assert(server_data->migration.bytes_transferred ==
+                   server_data->bar1_size + sizeof(time_t));
             ret = arm_timer(vfu_ctx, server_data->bar0);
             if (ret < 0) {
                 return ret;
@@ -312,18 +312,18 @@ migration_read_data(vfu_ctx_t *vfu_ctx, void *buf, uint64_t size)
      * more complex state tracking which exceeds the scope of this sample.
      */
 
-    if (server_data->migration.pending_bytes == 0 || size == 0) {
-        vfu_log(vfu_ctx, LOG_DEBUG, "no data left to read");
-        return 0;
-    }
-
     uint32_t total_read = server_data->bar1_size;
 
     if (server_data->migration.state == VFU_MIGR_STATE_STOP_AND_COPY) {
         total_read += sizeof(server_data->bar0);
     }
 
-    uint32_t read_start = total_read - server_data->migration.pending_bytes;
+    if (server_data->migration.bytes_transferred == total_read || size == 0) {
+        vfu_log(vfu_ctx, LOG_DEBUG, "no data left to read");
+        return 0;
+    }
+
+    uint32_t read_start = server_data->migration.bytes_transferred;
     uint32_t read_end = MIN(read_start + size, total_read); // exclusive
     assert(read_end > read_start);
 
@@ -357,7 +357,7 @@ migration_read_data(vfu_ctx_t *vfu_ctx, void *buf, uint64_t size)
         memcpy(buf, &server_data->bar0 + read_start, bytes_read);
     }
 
-    server_data->migration.pending_bytes -= bytes_read;
+    server_data->migration.bytes_transferred += bytes_read;
 
     return bytes_read;
 }
@@ -371,13 +371,13 @@ migration_write_data(vfu_ctx_t *vfu_ctx, void *data, uint64_t size)
     assert(server_data != NULL);
     assert(data != NULL);
 
-    if (server_data->migration.pending_bytes == 0 || size == 0) {
+    uint32_t total_write = server_data->bar1_size + sizeof(server_data->bar0);
+
+    if (server_data->migration.bytes_transferred == total_write || size == 0) {
         return 0;
     }
 
-    uint32_t total_write = server_data->bar1_size + sizeof(server_data->bar0);
-
-    uint32_t write_start = total_write - server_data->migration.pending_bytes;
+    uint32_t write_start = server_data->migration.bytes_transferred;
     uint32_t write_end = MIN(write_start + size, total_write); // exclusive
     assert(write_end > write_start);
 
@@ -385,15 +385,17 @@ migration_write_data(vfu_ctx_t *vfu_ctx, void *data, uint64_t size)
 
     if (write_end <= server_data->bar1_size) {
         // case 1: entire write lies within bar1
-        // TODO check the following is always allowed
 
         memcpy(server_data->bar1 + write_start, buf, bytes_written);
-    } else if (
-        write_start < server_data->bar1_size // starts in bar1
-        && write_end > server_data->bar1_size // ends in bar0
-    ) {
-        // case 2: part of the write in bar1 and part of the write in bar0
-        // TODO check the following is always allowed
+    } else if (write_start >= server_data->bar1_size) {
+        // case 2: entire write lies within bar0
+
+        write_start -= server_data->bar1_size;
+        write_end -= server_data->bar1_size;
+
+        memcpy(&server_data->bar0 + write_start, buf, bytes_written);
+    } else {
+        // case 3: part of the write in bar1 and part of the write in bar0
 
         uint32_t length_in_bar1 = server_data->bar1_size - write_start;
         uint32_t length_in_bar0 = write_end - server_data->bar1_size;
@@ -401,17 +403,9 @@ migration_write_data(vfu_ctx_t *vfu_ctx, void *data, uint64_t size)
         
         memcpy(server_data->bar1 + write_start, buf, length_in_bar1);
         memcpy(&server_data->bar0, buf + length_in_bar1, length_in_bar0);
-    } else if (write_start >= server_data->bar1_size) {
-        // case 3: entire write lies within bar0
-        // TODO check the following is always allowed
-
-        write_start -= server_data->bar1_size;
-        write_end -= server_data->bar1_size;
-
-        memcpy(&server_data->bar0 + write_start, buf, bytes_written);
     }
 
-    server_data->migration.pending_bytes -= bytes_written;
+    server_data->migration.bytes_transferred += bytes_written;
 
     return bytes_written;
 }