allow device to specify data_offset when resuming (#272)

Handling data_offset and data_size internally is wrong: we can't simply
assume that the migration data should be appending to the migration
region, devices might have their own requirements.

This also requires a way for the  device to return the data_offset, we
do this by making the prepare_data callback applicable in resume state.

Also, allow migration read/write callabcks to return errors.

Signed-off-by: Thanos Makatos <thanos.makatos@nutanix.com>

3 files changed, 45 insertions, 49 deletions

diff --git a/include/libvfio-user.h b/include/libvfio-user.h
index 3ee22ec..75b5943 100644
--- a/include/libvfio-user.h
+++ b/include/libvfio-user.h
@@ -381,35 +381,44 @@ typedef struct {
     __u64 (*get_pending_bytes)(vfu_ctx_t *vfu_ctx);
 
     /*
-     * Function that is called to instruct the device to prepare migration data.
-     * The function must return only after migration data are available at the
-     * specified offset.
+     * Function that is called to instruct the device to prepare migration data
+     * to be read when in pre-copy or stop-and-copy state, and to prepare for
+     * receiving migration data when in resuming state.
+     *
+     * When in pre-copy and stop-and-copy state, the function must return only
+     * after migration data are available at the specified offset.
+     *
+     * When in resuming state, @offset must be set to where migration data must
+     * written. @size points to NULL.
      */
     int (*prepare_data)(vfu_ctx_t *vfu_ctx, __u64 *offset, __u64 *size);
 
     /*
-     * Function that is called to read migration data. offset and size can
-     * be any subrange on the offset and size previously returned by
-     * prepare_data. The function must return the amount of data read. This
-     * function can be called even if the migration data can be memory mapped.
+     * Function that is called to read migration data. offset and size can be
+     * any subrange on the offset and size previously returned by prepare_data.
+     * The function must return the amount of data read or -errno on error.
+     * This function can be called even if the migration data can be memory
+     * mapped.
      *
      * Does this mean that reading data_offset/data_size updates the values?
      */
-    size_t (*read_data)(vfu_ctx_t *vfu_ctx, void *buf,
-                        __u64 count, __u64 offset);
+    ssize_t (*read_data)(vfu_ctx_t *vfu_ctx, void *buf,
+                         __u64 count, __u64 offset);
 
     /* Callbacks for restoring device state */
 
     /*
+     * Fuction that is called for writing previously stored device state. The
+     * function must return the amount of data written or -errno on error.
+     */
+    ssize_t (*write_data)(vfu_ctx_t *vfu_ctx, void *buf, __u64 count,
+                          __u64 offset);
+
+    /*
      * Function that is called when client has written some previously stored
      * device state.
      */
-    int (*data_written)(vfu_ctx_t *vfu_ctx,
-                        __u64 count, __u64 offset);
-
-    /* Fuction that is called for writing previously stored device state. */
-    size_t (*write_data)(vfu_ctx_t *vfu_ctx, void *buf,
-                         __u64 count, __u64 offset);
+    int (*data_written)(vfu_ctx_t *vfu_ctx, __u64 count);
 
 } vfu_migration_callbacks_t;
 
diff --git a/lib/migration.c b/lib/migration.c
index 3577cad..2f532e7 100644
--- a/lib/migration.c
+++ b/lib/migration.c
@@ -350,12 +350,16 @@ handle_data_offset(vfu_ctx_t *vfu_ctx, struct migration *migr,
         return ret;
     case VFIO_DEVICE_STATE_RESUMING:
         if (is_write) {
+            /* TODO writing to read-only registers should be simply ignored */
             vfu_log(vfu_ctx, LOG_ERR, "bad write to migration data_offset");
             return -EINVAL;
-        } else {
-            *offset = migr->info.data_offset + sizeof(struct vfio_device_migration_info);
-            return 0;
         }
+        ret = migr->callbacks.prepare_data(vfu_ctx, offset, NULL);
+        if (ret < 0) {
+            return ret;
+        }
+        *offset += sizeof(struct vfio_device_migration_info);
+        return 0;
     }
     /* TODO improve error message */
     vfu_log(vfu_ctx, LOG_ERR,
@@ -389,18 +393,12 @@ static ssize_t
 handle_data_size_when_resuming(vfu_ctx_t *vfu_ctx, struct migration *migr,
                                __u64 size, bool is_write)
 {
-    int ret = 0;
-
     assert(migr != NULL);
 
     if (is_write) {
-        ret = migr->callbacks.data_written(vfu_ctx, size, migr->info.data_offset);
-        if (ret >= 0) {
-            migr->info.data_size = size;
-            migr->info.data_offset += size;
-        }
+        return  migr->callbacks.data_written(vfu_ctx, size);
     }
-    return ret;
+    return 0;
 }
 
 static ssize_t
@@ -416,21 +414,16 @@ handle_data_size(vfu_ctx_t *vfu_ctx, struct migration *migr,
     case VFIO_DEVICE_STATE_SAVING:
     case VFIO_DEVICE_STATE_RUNNING | VFIO_DEVICE_STATE_SAVING:
         ret = handle_data_size_when_saving(vfu_ctx, migr, is_write);
-        break;
+        if (ret == 0 && !is_write) {
+            *size = migr->iter.size;
+        }
+        return ret;
     case VFIO_DEVICE_STATE_RESUMING:
-        ret = handle_data_size_when_resuming(vfu_ctx, migr, *size, is_write);
-        break;
-    default:
-        /* TODO improve error message */
-        vfu_log(vfu_ctx, LOG_ERR, "bad access to data_size");
-        ret = -EINVAL;
+        return handle_data_size_when_resuming(vfu_ctx, migr, *size, is_write);
     }
-
-    if (ret == 0 && !is_write) {
-        *size = migr->iter.size;
-    }
-
-    return ret;
+    /* TODO improve error message */
+    vfu_log(vfu_ctx, LOG_ERR, "bad access to data_size");
+    return -EINVAL;
 }
 
 static ssize_t
diff --git a/samples/server.c b/samples/server.c
index ad5ce34..82b22ff 100644
--- a/samples/server.c
+++ b/samples/server.c
@@ -330,16 +330,11 @@ migration_prepare_data(vfu_ctx_t *vfu_ctx, __u64 *offset, __u64 *size)
     } else {
         assert(false); /* FIXME fail gracefully */
     }
-
-    /*
-     * Don't provide all migration data in one go in order to make it a bit
-     * more interesting.
-     */
-    *size = MIN(server_data->migration.pending_bytes, server_data->migration.migr_data_len / 4);
+    *size = server_data->migration.pending_bytes;
     return 0;
 }
 
-static size_t
+static ssize_t
 migration_read_data(vfu_ctx_t *vfu_ctx, void *buf, __u64 size, __u64 offset)
 {
     struct server_data *server_data = vfu_get_private(vfu_ctx);
@@ -390,7 +385,7 @@ migration_read_data(vfu_ctx_t *vfu_ctx, void *buf, __u64 size, __u64 offset)
     return size;
 }
 
-static size_t
+static ssize_t
 migration_write_data(vfu_ctx_t *vfu_ctx, void *data, __u64 size, __u64 offset)
 {
     struct server_data *server_data = vfu_get_private(vfu_ctx);
@@ -430,8 +425,7 @@ migration_write_data(vfu_ctx_t *vfu_ctx, void *data, __u64 size, __u64 offset)
 
 
 static int
-migration_data_written(UNUSED vfu_ctx_t *vfu_ctx, UNUSED __u64 count,
-                       UNUSED __u64 offset)
+migration_data_written(UNUSED vfu_ctx_t *vfu_ctx, UNUSED __u64 count)
 {
     /*
      * We apply migration state directly in the migration_write_data callback,