Rust high-level pre/post migration callbacks

[Paolo Bonzini <pbonzini@redhat.com>]

Hi,

based on the low-level sketch in Zhao and my presentation,
I would like to propose this more high-level implementation
of pre/post migration callbacks.

Instead of dealing with pre/post callbacks, devices implement a
snapshot/restore mechanism; this way, C code sees a simplified
picture and does not have to deal with Rust concepts such as
Mutex<>.

Using it is very easy, you can just declare your state like:

     regs: Migratable<Mutex<MyDeviceRegisters>>

If a pure snapshot is possible, implementing the new trait
is also simple:

impl_vmstate_struct!(MyDeviceRegisters, ...);

impl ToMigrationState for MyDeviceRegisters {
     type Migrated = Self;
     fn to_migration_state(&self) ->
         Result<Box<Self>, ...> {
         Ok(Box::new(self.clone()))
     }

     fn restore_migrated_state_mut(&mut self, source: &Self,
         _version_id: u8) -> Result<(), migration::InvalidError> {
         *self = source;
         Ok(())
     }
}

I'm really bad at writing Rust code with the correct syntax
from the get-go, but I'll try anyway.

new traits:

/// Enables QEMU migration support for types that may be wrapped in
/// synchronization primitives (like `Mutex`) that the C migration
/// code cannot directly handle. The trait provides methods to
/// extract essential state for migration and restore it after
/// migration completes.
///
/// On top of extracting data from synchronization wrappers during save
/// and restoring it during load, it's also possible to convert
/// runtime representations to migration-safe formats.
trait ToMigrationState {
     type Migrated: Default + VMState;
     fn to_migration_state(&self) ->
         Result<Box<Self::Migrated>, migration::InvalidError>;
     fn restore_migrated_state_mut(&mut self, source: &Self::Migrated,
         version_id: u8) -> Result<(), migration::InvalidError>;
}

/// Extension trait for types that support migration state restoration
/// through interior mutability.
///
/// This trait extends `ToMigrationState` for types that can restore
/// their state without requiring mutable access.  While user structs
/// will generally use `ToMigrationState`, the device will have multiple
/// references and therefore the device struct has to employ an interior
/// mutability wrapper like `Mutex`, `RefCell`, or `BqlRefCell`.  In
/// turn, wrappers implementing this trait can be used within `Migratable<T>`,
/// which makes no assumptions on how to achieve mutable access to the
/// run-time state.
trait ToMigrationStateShared: ToMigrationState {
     fn restore_migrated_state(&self, source: &Self::Migrated) ->
         Result<(), migration::InvalidError>;
}


with implementations for wrapper types like:

impl<T> ToMigrationState for Mutex<T: ToMigrationState> {
     type Migrated = T::Migrated;
     fn to_migration_state(&self) ->
         Result<Box<Self::Migrated>, migration::InvalidError> {
         self.lock().to_migration_state()
     }
     ...
}

impl<T> ToMigrationStateShared for Mutex<T: ToMigrationState> {
     fn restore_migrated_state(&self, source: &Self::Migrated,
         version_id: u8) -> Result<(), migration::InvalidError>{
         self.lock().restore_migrated_state_mut(source, version_id)
     }
}

impl<T> ToMigrationState for BqlRefCell<T: ToMigrationState> {
     type Migrated = T::Migrated;
     fn to_migration_state(&self) ->
         Result<Box<Self::Migrated>, migration::InvalidError> {
         self.borrow().to_migration_state()
     }
     ...
}

impl<T> ToMigrationStateShared for BqlRefCell<T: ToMigrationState> {
     fn restore_migrated_state(&self, source: &Self::Migrated,
         version_id: u8) ->Result<(), migration::InvalidError> {
         self.borrow_mut().restore_migrated_state_mut(source, version_id)
     }
}

new struct maps the above trait to the C-style callbacks:

/// A wrapper that bridges Rust types with QEMU's C-based migration system.
///
/// `Migratable<T>` enables QEMU migration support for Rust types that implement
/// `ToMigrationState`, as long as they are wrapped with an interior mutability
/// like `Mutex` or `BqlRefCell`.  It provides translation functionality as well
/// as access to synchronization primitives that the C code cannot directly handle.
///
/// This wrapper acts as a transparent proxy during normal operation
/// (via `Deref`/`DerefMut`), while handling state extraction and restoration
/// around migration.
pub struct<T: ToMigrationStateShared> Migratable {
     runtime_state: T,
     // C vmstate does not support NULL pointers, so no Option<Box<>>
     // Actually a BqlCell<*mut T::Migrated>, but keeping it simple
     // for now.
     migration_state: *mut T::Migrated
};

unsafe impl<T> Send for Migratable<T: Send> {}
unsafe impl<T> Sync for Migratable<T: Sync> {}

// just return runtime_state
impl<T> Deref for Migratable<T: ToMigrationStateShared> {
     type Migrated = T;
     ...
}
impl<T> DerefMut for Migratable<T: ToMigrationStateShared> {
     ...
}

impl Migratable {
     fn pre_save(...) -> ... {
         self.migration_state = Box::into_raw(self.0.to_migration_state()?);
     }

     fn post_save(...) -> ... {
         drop(Box::from_raw(self.migration_state.replace(ptr::null_mut()));
     }

     fn pre_load(...) -> ... {
         self.migration_state = Box::into_raw(Box::default());
     }

     fn post_load(...) -> ... {
         let state = Box::from_raw(self.migration_state.replace(ptr::null_mut());
         self.0.restore_migrated_state(state, version_id)
     }
}

unsafe impl VMState for Migratable<T: ToMigrationStateShared> {
     const BASE: bindings::VMStateField = {
         static VMSD: &$crate::bindings::VMStateDescription =
             VMStateDescriptionBuilder::<Self>::new()
                 .version_id(T::VMSD.version_id)
                 .minimum_version_id(T::VMSD.minimum_version_id)
                 .priority(T::VMSD.priority)
                 .pre_load(Self::pre_load)
                 .post_load(Self::post_load)
                 .pre_save(Self::pre_save)
                 .post_save(Self::post_save)
                 .fields(vmstate_fields! {
                     vmstate_of!(Migratable<T>, migration_state)
                 }
                 .build();

         bindings::VMStateField {
             vmsd: addr_of!(*VMSD),
             size: size_of::<Migratable<T>>(),
             flags: bindings::VMStateFlags::VMS_STRUCT,
             ..common::Zeroable::ZERO
         }
     };
}

This is just a sketch but should give the idea.

Paolo