Support querying migration progress

include/libvfio-user.h

@@ -624,6 +624,39 @@ typedef struct {
      */
     ssize_t (*write_data)(vfu_ctx_t *vfu_ctx, void *buf, uint64_t count);
 
+    /*
+     * Function that is called for retrieving the estimated data length that
+     * needs to be read out to complete stop copy.
+     *
+     * The function must return the estimated amount of data, or -1 on error,
+     * setting errno.
+     */
+    ssize_t (*get_data_size)(vfu_ctx_t *vfu_ctx);
+
+    /*
+     * Function that is called to retrieve an estimate of the current data
+     * sizes remaining to be transfered. The estimate is split into two
+     * categories. The callback should fill in both of these values with device
+     * specific information.
+     *
+     * initial_bytes indicates the amount of initial precopy data available
+     * from the device. This value should be non-zero after a transition to
+     * PRE_COPY and decrease as migration data is read out from the server. It
+     * is recommended for clients to only transition from PRE_COPY to STOP_COPY
+     * after this field has reached zero.
+     *
+     * dirty_bytes tracks the state changes relative to data previously
+     * retrieved. This field should start at zero on after a transition to
+     * PRE_COPY, and may increase or decrease as migration state is read or as
+     * internal device state changes.
+     *
+     * This callback is guaranteed to only be called when the device is in
+     * PRE_COPY.
+     *
+     * The callback should return -1 on error, setting errno.
+     */
+    int (*get_precopy_info)(vfu_ctx_t *vfu_ctx, uint64_t *initial_bytes,
+                            uint64_t *dirty_bytes);
 } vfu_migration_callbacks_t;
 
 int

include/vfio-user.h

@@ -70,6 +70,7 @@ enum vfio_user_command {
     VFIO_USER_DEVICE_FEATURE            = 16,
     VFIO_USER_MIG_DATA_READ             = 17,
     VFIO_USER_MIG_DATA_WRITE            = 18,
+    VFIO_USER_MIG_GET_PRECOPY_INFO      = 19,
     VFIO_USER_MAX,
 };
 
@@ -265,6 +266,16 @@ struct vfio_user_device_feature_migration {
 _Static_assert(sizeof(struct vfio_user_device_feature_migration) == 8,
                "bad vfio_user_device_feature_migration size");
 
+/* Analogous to struct vfio_device_feature_mig_data_size */
+struct vfio_user_device_feature_mig_data_size {
+        uint64_t stop_copy_length;
+} __attribute__((packed));
+#ifndef VFIO_DEVICE_FEATURE_MIG_DATA_SIZE
+#define VFIO_DEVICE_FEATURE_MIG_DATA_SIZE 9
+#endif
+_Static_assert(sizeof(struct vfio_user_device_feature_migration) == 8,
+               "bad vfio_user_device_feature_mig_state size");
+
 /* Analogous to struct vfio_device_feature_mig_state */
 struct vfio_user_device_feature_mig_state {
     uint32_t    device_state;
@@ -273,7 +284,7 @@ struct vfio_user_device_feature_mig_state {
 #ifndef VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE
 #define VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE 2
 #endif
-_Static_assert(sizeof(struct vfio_user_device_feature_migration) == 8,
+_Static_assert(sizeof(struct vfio_user_device_feature_mig_state) == 8,
                "bad vfio_user_device_feature_mig_state size");
 
 /* Analogous to enum vfio_device_mig_state */
@@ -295,6 +306,14 @@ struct vfio_user_mig_data {
     uint8_t     data[];
 } __attribute__((packed));
 
+/* Analogous to struct vfio_precopy_info */
+struct vfio_user_precopy_info {
+        uint32_t argsz;
+        uint32_t flags;
+        uint64_t initial_bytes;
+        uint64_t dirty_bytes;
+} __attribute__((packed));
+
 #ifdef __cplusplus
 }
 #endif

lib/libvfio-user.c

@@ -914,6 +914,7 @@ device_feature_flags_supported(vfu_ctx_t *vfu_ctx, uint32_t feature)
 
     switch (feature) {
         case VFIO_DEVICE_FEATURE_MIGRATION:
+        case VFIO_DEVICE_FEATURE_MIG_DATA_SIZE:
         case VFIO_DEVICE_FEATURE_DMA_LOGGING_REPORT:
             return VFIO_DEVICE_FEATURE_GET | VFIO_DEVICE_FEATURE_PROBE;
         case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
@@ -934,6 +935,7 @@ is_migration_feature(uint32_t feature)
     switch (feature) {
         case VFIO_DEVICE_FEATURE_MIGRATION:
         case VFIO_DEVICE_FEATURE_MIG_DEVICE_STATE:
+        case VFIO_DEVICE_FEATURE_MIG_DATA_SIZE:
             return true;
     }
 
@@ -1001,6 +1003,17 @@ handle_migration_device_feature_get(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg,
             state->device_state = migration_get_state(vfu_ctx);
             return 0;
         }
+
+        case VFIO_DEVICE_FEATURE_MIG_DATA_SIZE: {
+            struct vfio_user_device_feature_mig_data_size *data_size =
+                (void *)res->data;
+            ssize_t ret = migration_get_data_size(vfu_ctx);
+            if (ret < 0) {
+                return ret;
+            }
+            data_size->stop_copy_length = ret;
+            return 0;
+        }
 
         default:
             vfu_log(vfu_ctx, LOG_ERR, "invalid flags for migration GET (%d)",
@@ -1378,6 +1391,10 @@ handle_request(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
         ret = handle_mig_data_write(vfu_ctx, msg);
         break;
 
+    case VFIO_USER_MIG_GET_PRECOPY_INFO:
+        ret = handle_mig_get_precopy_info(vfu_ctx, msg);
+        break;
+
     default:
         msg->processed_cmd = false;
         vfu_log(vfu_ctx, LOG_ERR, "bad command %d", msg->hdr.cmd);

lib/migration.c

@@ -285,6 +285,12 @@ migration_set_state(vfu_ctx_t *vfu_ctx, uint32_t device_state)
     return ret;
 }
 
+ssize_t
+migration_get_data_size(vfu_ctx_t *vfu_ctx)
+{
+    return vfu_ctx->migration->callbacks.get_data_size(vfu_ctx);
+}
+
 ssize_t
 handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
 {
@@ -406,6 +412,44 @@ handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
     return 0;
 }
 
+int
+handle_mig_get_precopy_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg)
+{
+    struct migration *migr = vfu_ctx->migration;
+
+    if (migr->state != VFIO_USER_DEVICE_STATE_PRE_COPY) {
+        vfu_log(vfu_ctx, LOG_ERR, "bad migration state to get precopy info: %d",
+                migr->state);
+        return ERROR_INT(EINVAL);
+    }
+
+    msg->out.iov.iov_len = sizeof(struct vfio_user_precopy_info);
+    msg->out.iov.iov_base = calloc(1, msg->out.iov.iov_len);
+
+    if (msg->out.iov.iov_base == NULL) {
+        return ERROR_INT(ENOMEM);
+    }
+
+    struct vfio_user_precopy_info *res = msg->out.iov.iov_base;
+    res->argsz = sizeof(struct vfio_user_precopy_info);
+
+    uint64_t initial_bytes, dirty_bytes;
+    ssize_t ret = migr->callbacks.get_precopy_info(vfu_ctx, &initial_bytes,
+                                                   &dirty_bytes);
+
+    if (ret < 0) {
+        vfu_log(vfu_ctx, LOG_ERR, "get_precopy_info callback failed, errno=%d",
+                errno);
+        iov_free(&msg->out.iov);
+        return ret;
+    }
+
+    res->initial_bytes = initial_bytes;
+    res->dirty_bytes = dirty_bytes;
+
+    return 0;
+}
+
 bool
 MOCK_DEFINE(device_is_stopped_and_copying)(struct migration *migr)
 {

lib/migration.h

@@ -53,12 +53,18 @@ migration_get_state(vfu_ctx_t *vfu_ctx);
 ssize_t
 migration_set_state(vfu_ctx_t *vfu_ctx, uint32_t device_state);
 
+ssize_t
+migration_get_data_size(vfu_ctx_t *vfu_ctx);
+
 ssize_t
 handle_mig_data_read(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
 ssize_t
 handle_mig_data_write(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
 
+int
+handle_mig_get_precopy_info(vfu_ctx_t *vfu_ctx, vfu_msg_t *msg);
+
 bool
 migration_available(vfu_ctx_t *vfu_ctx);
 

lib/migration_priv.h

@@ -47,4 +47,4 @@ MOCK_DECLARE(int, state_trans_notify, vfu_ctx_t *vfu_ctx,
 
 #endif
 
-/* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */
+/* ex: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab: */

samples/gpio-pci-idio-16.c

@@ -52,6 +52,7 @@ _log(vfu_ctx_t *vfu_ctx UNUSED, UNUSED int level, char const *msg)
 }
 
 static int pin[16];
+bool initial_dirty;
 bool dirty = true;
 
 static ssize_t
@@ -119,6 +120,10 @@ static int
 migration_device_state_transition(vfu_ctx_t *vfu_ctx, vfu_migr_state_t state)
 {
     vfu_log(vfu_ctx, LOG_DEBUG, "migration: transition to state %d", state);
+    if (state == VFU_MIGR_STATE_PRE_COPY) {
+        initial_dirty = true;
+        dirty = true;
+    }
     return 0;
 }
 
@@ -130,6 +135,7 @@ migration_read_data(UNUSED vfu_ctx_t *vfu_ctx, void *buf, uint64_t size)
     if (dirty) {
         memcpy(buf, &pin, sizeof(pin));
         dirty = false;
+        initial_dirty = false;
         return sizeof(pin);
     }
 
@@ -144,6 +150,23 @@ migration_write_data(UNUSED vfu_ctx_t *vfu_ctx, void *buf, uint64_t size)
     return 0;
 }
 
+static ssize_t
+migration_get_data_size(UNUSED vfu_ctx_t *vfu_ctx)
+{
+    return dirty ? sizeof(pin) : 0;
+}
+
+static int
+migration_get_precopy_info(UNUSED vfu_ctx_t *vfu_ctx, uint64_t *initial_bytes,
+                           uint64_t *dirty_bytes)
+{
+    *initial_bytes = initial_dirty ? sizeof(pin) : 0;
+    *dirty_bytes = dirty ? sizeof(pin) : 0;
+    *dirty_bytes -= *initial_bytes;
+
+    return 0;
+}
+
 static void
 dma_register(UNUSED vfu_ctx_t *vfu_ctx, UNUSED vfu_dma_info_t *info)
 {
@@ -168,7 +191,9 @@ main(int argc, char *argv[])
         .version = VFU_MIGR_CALLBACKS_VERS,
         .transition = &migration_device_state_transition,
         .read_data = &migration_read_data,
-        .write_data = &migration_write_data
+        .write_data = &migration_write_data,
+        .get_data_size = &migration_get_data_size,
+        .get_precopy_info = &migration_get_precopy_info,
     };
 
     while ((opt = getopt(argc, argv, "MRv")) != -1) {

test/py/libvfio_user.py

@@ -580,6 +580,8 @@ class vfio_user_bitmap_range(Structure):
 transition_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, c.c_int, use_errno=True)
 read_data_cb_t = c.CFUNCTYPE(c.c_ssize_t, c.c_void_p, c.c_void_p, c.c_uint64)
 write_data_cb_t = c.CFUNCTYPE(c.c_ssize_t, c.c_void_p, c.c_void_p, c.c_uint64)
+get_precopy_info_cb_t = c.CFUNCTYPE(c.c_int, c.c_void_p, c.c_void_p, c.c_void_p)
+get_data_size_cb_t = c.CFUNCTYPE(c.c_ssize_t, c.c_void_p)
 
 
 class vfu_migration_callbacks_t(Structure):
@@ -588,6 +590,8 @@ class vfu_migration_callbacks_t(Structure):
         ("transition", transition_cb_t),
         ("read_data", read_data_cb_t),
         ("write_data", write_data_cb_t),
+        ("get_data_size", get_data_size_cb_t),
+        ("get_precopy_info", get_precopy_info_cb_t),
     ]
 
 

test/py/test_migration.py

@@ -122,6 +122,19 @@ def migr_write_data_cb(_ctx, buf, count):
     return count
 
 
+@get_data_size_cb_t
+def migr_get_data_size(_ctx):
+    global read_data
+
+    return len(read_data) if read_data else 0;
+
+
+@get_precopy_info_cb_t
+def migr_get_precopy_info(_ctx, init, dirty):
+    # TODO implement this for tests
+    return 0
+
+
 def setup_fail_callbacks(errno):
     global callbacks_errno
     callbacks_errno = errno
@@ -160,6 +173,8 @@ def test_migration_setup():
     cbs.transition = migr_trans_cb
     cbs.read_data = migr_read_data_cb
     cbs.write_data = migr_write_data_cb
+    cbs.get_data_size = migr_get_data_size
+    cbs.get_precopy_info = migr_get_precopy_info
 
     ret = vfu_setup_device_migration_callbacks(ctx, cbs)
     assert ret < 0, "do not allow old callbacks version"