// Copyright 2019 Intel Corporation. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::fs::{File, remove_file};
use std::io::ErrorKind;
use std::ops::Deref;
use std::os::unix::io::AsRawFd;
use std::sync::atomic::{AtomicBool, AtomicU8, Ordering};
use std::sync::{Arc, Barrier, Mutex};
use std::{io, result};

use anyhow::anyhow;
use event_monitor::event;
use log::{error, info, warn};
use seccompiler::SeccompAction;
use serde::{Deserialize, Serialize};
use thiserror::Error;
use vhost::Error as VhostError;
use vhost::vhost_user::message::{
    VhostUserInflight, VhostUserProtocolFeatures, VhostUserVirtioFeatures,
};
use vhost::vhost_user::{Error as VhostUserError, FrontendReqHandler, VhostUserFrontendReqHandler};
use virtio_queue::{Error as QueueError, Queue};
use vm_memory::guest_memory::Error as MmapError;
use vm_memory::mmap::MmapRegionError;
use vm_memory::{Address, GuestAddressSpace, GuestMemory, GuestMemoryAtomic};
use vm_migration::protocol::MemoryRangeTable;
use vm_migration::{MigratableError, Pausable, Snapshot};
use vmm_sys_util::eventfd::EventFd;
use vu_common_ctrl::VhostUserHandle;

use crate::seccomp_filters::Thread;
use crate::{
    ActivateError, EPOLL_HELPER_EVENT_LAST, EpollHelper, EpollHelperError, EpollHelperHandler,
    GuestMemoryMmap, GuestRegionMmap, VIRTIO_F_IN_ORDER, VIRTIO_F_NOTIFICATION_DATA,
    VIRTIO_F_ORDER_PLATFORM, VIRTIO_F_RING_EVENT_IDX, VIRTIO_F_RING_INDIRECT_DESC,
    VIRTIO_F_VERSION_1, VirtioCommon, VirtioInterrupt,
};

pub mod blk;
pub mod fs;
pub mod generic_vhost_user;
pub mod net;
pub mod vu_common_ctrl;

pub use self::blk::Blk;
pub use self::fs::*;
pub use self::generic_vhost_user::GenericVhostUser;
pub use self::net::Net;
pub use self::vu_common_ctrl::VhostUserConfig;

#[derive(Error, Debug)]
pub enum Error {
    #[error("Failed accepting connection")]
    AcceptConnection(#[source] io::Error),
    #[error("Invalid available address")]
    AvailAddress,
    #[error("Queue number  is not correct")]
    BadQueueNum,
    #[error("Failed binding vhost-user socket")]
    BindSocket(#[source] io::Error),
    #[error("Creating eventfd failed")]
    CreateEventFd(#[source] io::Error),
    #[error("Cloning kill eventfd failed")]
    CloneKillEventFd(#[source] io::Error),
    #[error("Invalid descriptor table address")]
    DescriptorTableAddress,
    #[error("Signal used queue failed")]
    FailedSignalingUsedQueue(#[source] io::Error),
    #[error("Failed to read vhost eventfd")]
    MemoryRegions(#[source] MmapError),
    #[error("Failed removing socket path")]
    RemoveSocketPath(#[source] io::Error),
    #[error("Failed to create frontend")]
    VhostUserCreateFrontend(#[source] VhostError),
    #[error("Failed to open vhost device")]
    VhostUserOpen(#[source] VhostError),
    #[error("Connection to socket failed")]
    VhostUserConnect(#[source] VhostError),
    #[error("Backend disconnected for socket {0}")]
    BackendDisconnected(String),
    #[error("Get features failed")]
    VhostUserGetFeatures(#[source] VhostError),
    #[error("Get queue max number failed")]
    VhostUserGetQueueMaxNum(#[source] VhostError),
    #[error("Get protocol features failed")]
    VhostUserGetProtocolFeatures(#[source] VhostError),
    #[error("Get vring base failed")]
    VhostUserGetVringBase(#[source] VhostError),
    #[error("Vhost-user Backend not support vhost-user protocol")]
    VhostUserProtocolNotSupport,
    #[error("Set owner failed")]
    VhostUserSetOwner(#[source] VhostError),
    #[error("Reset owner failed")]
    VhostUserResetOwner(#[source] VhostError),
    #[error("Set features failed")]
    VhostUserSetFeatures(#[source] VhostError),
    #[error("Set protocol features failed")]
    VhostUserSetProtocolFeatures(#[source] VhostError),
    #[error("Set mem table failed")]
    VhostUserSetMemTable(#[source] VhostError),
    #[error("Set vring num failed")]
    VhostUserSetVringNum(#[source] VhostError),
    #[error("Set vring addr failed")]
    VhostUserSetVringAddr(#[source] VhostError),
    #[error("Set vring base failed")]
    VhostUserSetVringBase(#[source] VhostError),
    #[error("Set vring call failed")]
    VhostUserSetVringCall(#[source] VhostError),
    #[error("Set vring kick failed")]
    VhostUserSetVringKick(#[source] VhostError),
    #[error("Set vring enable failed")]
    VhostUserSetVringEnable(#[source] VhostError),
    #[error("Failed to create vhost eventfd")]
    VhostIrqCreate(#[source] io::Error),
    #[error("Failed to read vhost eventfd")]
    VhostIrqRead(#[source] io::Error),
    #[error("Failed to read vhost eventfd")]
    VhostUserMemoryRegion(#[source] MmapError),
    #[error("Failed to create the frontend request handler from backend")]
    FrontendReqHandlerCreation(#[source] VhostUserError),
    #[error("Set backend request fd failed")]
    VhostUserSetBackendRequestFd(#[source] vhost::Error),
    #[error("Add memory region failed")]
    VhostUserAddMemReg(#[source] VhostError),
    #[error("Failed getting the configuration")]
    VhostUserGetConfig(#[source] VhostError),
    #[error("Failed setting the configuration")]
    VhostUserSetConfig(#[source] VhostError),
    #[error("Failed getting inflight shm log")]
    VhostUserGetInflight(#[source] VhostError),
    #[error("Failed setting inflight shm log")]
    VhostUserSetInflight(#[source] VhostError),
    #[error("Failed setting the log base")]
    VhostUserSetLogBase(#[source] VhostError),
    #[error("Invalid used address")]
    UsedAddress,
    #[error("Invalid features provided from vhost-user backend")]
    InvalidFeatures,
    #[error("Missing file descriptor for the region")]
    MissingRegionFd,
    #[error("Missing IrqFd")]
    MissingIrqFd,
    #[error("Failed getting the available index")]
    GetAvailableIndex(#[source] QueueError),
    #[error("Migration is not supported by this vhost-user device")]
    MigrationNotSupported,
    #[error("Failed creating memfd")]
    MemfdCreate(#[source] io::Error),
    #[error("Failed truncating the file size to the expected size")]
    SetFileSize(#[source] io::Error),
    #[error("Failed to set the seals on the file")]
    SetSeals(#[source] io::Error),
    #[error("Failed creating new mmap region")]
    NewMmapRegion(#[source] MmapRegionError),
    #[error("Could not find the shm log region")]
    MissingShmLogRegion,
    #[error("Failed setting device state fd")]
    VhostUserSetDeviceStateFd(#[source] VhostError),
    #[error("Failed checking device state")]
    VhostUserCheckDeviceState(#[source] VhostError),
    #[error("Failed saving/restoring backend state")]
    SaveRestoreBackendState(#[source] io::Error),
    #[error("Vring bases count ({0}) does not match queue count ({1})")]
    VringBasesCountMismatch(usize, usize),
    #[error("Backend state and vring bases must both be present or both be absent")]
    InconsistentBackendState,
    #[error("Failed to create timerfd")]
    TimerFdCreate(#[source] io::Error),
    #[error("Failed to arm timerfd")]
    TimerFdArm(#[source] io::Error),
    #[error("Failed waiting on timerfd")]
    TimerFdWait(#[source] io::Error),
    #[error("Failed to create epoll instance")]
    EpollCreate(#[source] io::Error),
    #[error("Failed to add fd to epoll")]
    EpollCtl(#[source] io::Error),
    #[error("Failed waiting on epoll")]
    EpollWait(#[source] io::Error),
    #[error("Aborted vhost-user connect: kill event received")]
    ConnectKilled,
}
type Result<T> = result::Result<T, Error>;

fn io_error_is_connection_lost(error: &io::Error) -> bool {
    matches!(
        error.kind(),
        ErrorKind::BrokenPipe
            | ErrorKind::ConnectionAborted
            | ErrorKind::ConnectionReset
            | ErrorKind::NotConnected
            | ErrorKind::UnexpectedEof
    )
}

fn vhost_user_error_is_transport_lost(error: &VhostUserError) -> bool {
    match error {
        VhostUserError::Disconnected
        | VhostUserError::PartialMessage
        | VhostUserError::SocketBroken(_)
        | VhostUserError::SocketConnect(_) => true,
        VhostUserError::SocketError(e) => io_error_is_connection_lost(e),
        _ => false,
    }
}

fn vhost_error_is_transport_lost(error: &VhostError) -> bool {
    match error {
        VhostError::VhostUserProtocol(e) => vhost_user_error_is_transport_lost(e),
        VhostError::IOError(e) => io_error_is_connection_lost(e),
        _ => false,
    }
}

impl Error {
    fn is_transport_lost(&self) -> bool {
        match self {
            Error::VhostUserConnect(_) | Error::BackendDisconnected(_) => true,
            Error::VhostUserCreateFrontend(e)
            | Error::VhostUserOpen(e)
            | Error::VhostUserGetFeatures(e)
            | Error::VhostUserGetQueueMaxNum(e)
            | Error::VhostUserGetProtocolFeatures(e)
            | Error::VhostUserGetVringBase(e)
            | Error::VhostUserSetOwner(e)
            | Error::VhostUserResetOwner(e)
            | Error::VhostUserSetFeatures(e)
            | Error::VhostUserSetProtocolFeatures(e)
            | Error::VhostUserSetMemTable(e)
            | Error::VhostUserSetVringNum(e)
            | Error::VhostUserSetVringAddr(e)
            | Error::VhostUserSetVringBase(e)
            | Error::VhostUserSetVringCall(e)
            | Error::VhostUserSetVringKick(e)
            | Error::VhostUserSetVringEnable(e)
            | Error::VhostUserSetBackendRequestFd(e)
            | Error::VhostUserAddMemReg(e)
            | Error::VhostUserGetConfig(e)
            | Error::VhostUserSetConfig(e)
            | Error::VhostUserGetInflight(e)
            | Error::VhostUserSetInflight(e)
            | Error::VhostUserSetLogBase(e)
            | Error::VhostUserSetDeviceStateFd(e)
            | Error::VhostUserCheckDeviceState(e) => vhost_error_is_transport_lost(e),
            Error::FrontendReqHandlerCreation(e) => vhost_user_error_is_transport_lost(e),
            _ => false,
        }
    }
}

pub const DEFAULT_VIRTIO_FEATURES: u64 = (1 << VIRTIO_F_RING_INDIRECT_DESC)
    | (1 << VIRTIO_F_RING_EVENT_IDX)
    | (1 << VIRTIO_F_VERSION_1)
    | (1 << VIRTIO_F_IN_ORDER)
    | (1 << VIRTIO_F_ORDER_PLATFORM)
    | (1 << VIRTIO_F_NOTIFICATION_DATA)
    | VhostUserVirtioFeatures::PROTOCOL_FEATURES.bits()
    | VhostUserVirtioFeatures::LOG_ALL.bits();

const HUP_CONNECTION_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 1;
const BACKEND_REQ_EVENT: u16 = EPOLL_HELPER_EVENT_LAST + 2;

#[derive(Default)]
pub struct Inflight {
    pub info: VhostUserInflight,
    pub fd: Option<File>,
}

pub struct VhostUserEpollHandler<S: VhostUserFrontendReqHandler> {
    pub vu: Arc<Mutex<VhostUserHandle>>,
    pub mem: GuestMemoryAtomic<GuestMemoryMmap>,
    pub kill_evt: EventFd,
    pub pause_evt: EventFd,
    pub queues: Vec<(usize, Queue, EventFd)>,
    pub virtio_interrupt: Arc<dyn VirtioInterrupt>,
    pub acked_features: u64,
    pub acked_protocol_features: u64,
    pub socket_path: String,
    pub server: bool,
    pub backend_req_handler: Option<FrontendReqHandler<S>>,
    pub inflight: Option<Inflight>,
    /// Flag set by the worker when the vhost-user backend is no longer reachable.
    pub disconnected: Arc<AtomicBool>,
}

impl<S: VhostUserFrontendReqHandler> VhostUserEpollHandler<S> {
    pub fn run(
        &mut self,
        paused: &AtomicBool,
        paused_sync: &Barrier,
    ) -> result::Result<(), EpollHelperError> {
        let mut helper = EpollHelper::new(&self.kill_evt, &self.pause_evt)?;
        helper.add_event_custom(
            self.vu.lock().unwrap().socket_handle().as_raw_fd(),
            HUP_CONNECTION_EVENT,
            epoll::Events::EPOLLHUP,
        )?;

        if let Some(backend_req_handler) = &self.backend_req_handler {
            helper.add_event(backend_req_handler.as_raw_fd(), BACKEND_REQ_EVENT)?;
        }

        helper.run(paused, paused_sync, self)?;

        Ok(())
    }

    fn reconnect(&mut self, helper: &mut EpollHelper) -> result::Result<(), EpollHelperError> {
        let result = self.reconnect_inner(helper);
        if result.is_err() {
            // If reconnect fails, mark disconnected to avoid repeated failed socket calls.
            self.disconnected.store(true, Ordering::Relaxed);
        }
        result
    }

    fn reconnect_inner(
        &mut self,
        helper: &mut EpollHelper,
    ) -> result::Result<(), EpollHelperError> {
        helper.del_event_custom(
            self.vu.lock().unwrap().socket_handle().as_raw_fd(),
            HUP_CONNECTION_EVENT,
            epoll::Events::EPOLLHUP,
        )?;

        let mut vhost_user = match VhostUserHandle::connect_vhost_user(
            self.server,
            &self.socket_path,
            self.queues.len() as u64,
            true,
            Some(&self.kill_evt),
        ) {
            Ok(vu) => vu,
            // Kill event fired during the connect retry loop; abandon the
            // reconnect attempt. The EpollHelper observes the same kill
            // event and will tear down on its next iteration.
            Err(Error::ConnectKilled) => return Ok(()),
            Err(e) => {
                return Err(EpollHelperError::IoError(io::Error::other(format!(
                    "failed connecting vhost-user backend for socket {}: {e:?}",
                    self.socket_path
                ))));
            }
        };

        let queues = self
            .queues
            .iter()
            .map(|(i, q, e)| (*i, vm_virtio::clone_queue(q), e.try_clone().unwrap()))
            .collect::<Vec<_>>();
        // Initialize the backend
        vhost_user
            .reinitialize_vhost_user(
                self.mem.memory().deref(),
                &queues,
                self.virtio_interrupt.as_ref(),
                self.acked_features,
                self.acked_protocol_features,
                &self.backend_req_handler,
                self.inflight.as_mut(),
            )
            .map_err(|e| {
                EpollHelperError::IoError(io::Error::other(format!(
                    "failed reconnecting vhost-user backend: {e:?}"
                )))
            })?;

        helper.add_event_custom(
            vhost_user.socket_handle().as_raw_fd(),
            HUP_CONNECTION_EVENT,
            epoll::Events::EPOLLHUP,
        )?;

        // Update vhost-user reference
        let mut vu = self.vu.lock().unwrap();
        *vu = vhost_user;

        Ok(())
    }
}

impl<S: VhostUserFrontendReqHandler> EpollHelperHandler for VhostUserEpollHandler<S> {
    fn handle_event(
        &mut self,
        helper: &mut EpollHelper,
        event: &epoll::Event,
    ) -> result::Result<(), EpollHelperError> {
        let ev_type = event.data as u16;
        let result = match ev_type {
            HUP_CONNECTION_EVENT => {
                info!(
                    "vhost-user backend for socket {} disconnected, attempting reconnection",
                    self.socket_path
                );
                self.reconnect(helper).map_err(|e| {
                    EpollHelperError::HandleEvent(anyhow!(
                        "failed to reconnect vhost-user backend for socket {}: {e:?}",
                        self.socket_path
                    ))
                })
            }
            BACKEND_REQ_EVENT => {
                if let Some(backend_req_handler) = self.backend_req_handler.as_mut() {
                    backend_req_handler
                        .handle_request()
                        .map(|_| ())
                        .map_err(|e| {
                            EpollHelperError::HandleEvent(anyhow!(
                                "Failed to handle request from vhost-user backend: {e:?}"
                            ))
                        })
                } else {
                    Ok(())
                }
            }
            _ => Err(EpollHelperError::HandleEvent(anyhow!(
                "Unknown event for vhost-user thread"
            ))),
        };

        // If the backend hits and error it is unusable from this point on.
        if result.is_err() {
            self.disconnected.store(true, Ordering::Relaxed);
        }
        result
    }
}

/// Common snapshot state for all vhost-user device types.
///
/// Generic over `C` which is the device-specific config type
/// (e.g. VirtioBlockConfig, VirtioFsConfig, VirtioNetConfig).
/// Devices without a config type use `()`.
#[derive(Default, Serialize, Deserialize)]
pub struct VhostUserState<C> {
    pub avail_features: u64,
    pub acked_features: u64,
    pub config: C,
    pub acked_protocol_features: u64,
    pub vu_num_queues: usize,
    #[serde(default)]
    pub backend_req_support: bool,
    #[serde(default)]
    pub vring_bases: Option<Vec<u64>>,
    #[serde(default)]
    pub backend_state: Option<Vec<u8>>,
}

impl<C> VhostUserState<C> {
    pub fn validate(&self) -> Result<()> {
        if self.backend_state.is_some() != self.vring_bases.is_some() {
            return Err(Error::InconsistentBackendState);
        }
        Ok(())
    }
}

#[derive(Default)]
pub struct VhostUserCommon {
    pub virtio_common: VirtioCommon,
    pub vu: Option<Arc<Mutex<VhostUserHandle>>>,
    pub guest_memory: Option<GuestMemoryAtomic<GuestMemoryMmap>>,
    pub acked_protocol_features: u64,
    pub socket_path: String,
    pub vu_num_queues: usize,
    pub migration_started: bool,
    pub server: bool,
    pub vring_bases: Option<Vec<u64>>,
    /// Indicates that the backend is no longer reachable. Shared with EPollHandler.
    pub disconnected: Arc<AtomicBool>,
    saved_dirty_log: Option<MemoryRangeTable>,
    dirty_logging: bool,
}

impl VhostUserCommon {
    #[expect(clippy::too_many_arguments)]
    pub fn activate<T: VhostUserFrontendReqHandler>(
        &mut self,
        mem: GuestMemoryAtomic<GuestMemoryMmap>,
        queues: &[(usize, Queue, EventFd)],
        interrupt_cb: Arc<dyn VirtioInterrupt>,
        acked_features: u64,
        backend_req_handler: Option<FrontendReqHandler<T>>,
        kill_evt: EventFd,
        pause_evt: EventFd,
    ) -> result::Result<VhostUserEpollHandler<T>, ActivateError> {
        self.guest_memory = Some(mem.clone());

        if self.disconnected.load(Ordering::Relaxed) {
            warn!(
                "Not activating disconnected vhost-user device for socket {}",
                self.socket_path
            );
            return Err(ActivateError::BadActivate);
        }

        let mut inflight: Option<Inflight> =
            if self.acked_protocol_features & VhostUserProtocolFeatures::INFLIGHT_SHMFD.bits() != 0
            {
                Some(Inflight::default())
            } else {
                None
            };

        if self.vu.is_none() {
            error!("Missing vhost-user handle for socket {}", self.socket_path);
            return Err(ActivateError::BadActivate);
        }
        let vu = self.vu.as_ref().unwrap();
        let queues = queues
            .iter()
            .map(|(i, q, e)| (*i, vm_virtio::clone_queue(q), e.try_clone().unwrap()))
            .collect::<Vec<_>>();
        let vring_bases = self.vring_bases.take();
        vu.lock()
            .unwrap()
            .setup_vhost_user(
                &mem.memory(),
                &queues,
                interrupt_cb.as_ref(),
                acked_features,
                &backend_req_handler,
                inflight.as_mut(),
                vring_bases.as_deref(),
            )
            .map_err(ActivateError::VhostUserSetup)?;

        Ok(VhostUserEpollHandler {
            vu: vu.clone(),
            mem,
            kill_evt,
            pause_evt,
            queues,
            virtio_interrupt: interrupt_cb,
            acked_features,
            acked_protocol_features: self.acked_protocol_features,
            socket_path: self.socket_path.clone(),
            server: self.server,
            backend_req_handler,
            inflight,
            disconnected: self.disconnected.clone(),
        })
    }

    /// Like `VirtioCommon::spawn_worker`, but on failure also runs
    /// `self.reset(id)` to tear down the vhost-user backend.
    #[expect(clippy::too_many_arguments)]
    pub fn spawn_worker<F>(
        &mut self,
        id: &str,
        seccomp_action: &SeccompAction,
        thread_type: Thread,
        exit_evt: &EventFd,
        device_status: Arc<AtomicU8>,
        interrupt_cb: Arc<dyn VirtioInterrupt>,
        f: F,
    ) -> result::Result<(), ActivateError>
    where
        F: FnOnce() -> result::Result<(), EpollHelperError> + Send + 'static,
    {
        if let Err(e) = self.virtio_common.spawn_worker(
            id,
            seccomp_action,
            thread_type,
            exit_evt,
            device_status,
            interrupt_cb,
            f,
        ) {
            self.reset(id);
            return Err(e);
        }
        Ok(())
    }

    pub fn reset(&mut self, id: &str) {
        // Resume the virtio thread if it was paused. Reset must always
        // converge to fresh state, so backend resume / reset failures are
        // logged but don't skip the rest of the teardown.
        if self.virtio_common.pause_evt.take().is_some()
            && let Err(e) = self.virtio_common.resume()
        {
            error!("Failed to resume paused device during reset: {e:?}");
        }

        // Skip reset_vhost_user if the backend is not connected.
        if !self.disconnected.load(Ordering::Relaxed)
            && let Some(vu) = &self.vu
            && let Err(e) = vu.lock().unwrap().reset_vhost_user()
        {
            if e.is_transport_lost() {
                warn!(
                    "Failed to reset vhost-user daemon for socket {}: {e:?}; \
                     marking device as disconnected",
                    self.socket_path
                );
                self.disconnected.store(true, Ordering::Relaxed);
            } else {
                warn!(
                    "Failed to reset vhost-user daemon for socket {}: {e:?}",
                    self.socket_path
                );
            }
        }

        self.virtio_common.reset();

        event!("virtio-device", "reset", "id", id);
    }

    fn memory_update_error(&self, source: Error) -> crate::Error {
        if self.acked_protocol_features & VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS.bits() != 0
        {
            crate::Error::VhostUserAddMemoryRegion(source)
        } else {
            crate::Error::VhostUserUpdateMemory(source)
        }
    }

    pub fn shutdown(&mut self) {
        // Join the workers so they drop their Arc<VhostUserHandle> and the
        // socket closes, letting the migration destination reconnect.
        self.virtio_common.wait_for_epoll_threads();

        // Remove socket path if needed
        if self.server {
            let _ = remove_file(&self.socket_path);
        }

        // Drop the vhost-user handle
        self.vu = None;
    }

    fn add_memory_region_internal(&self, region: &Arc<GuestRegionMmap>) -> Result<()> {
        if let Some(vu) = &self.vu {
            if self.acked_protocol_features & VhostUserProtocolFeatures::CONFIGURE_MEM_SLOTS.bits()
                != 0
            {
                return vu.lock().unwrap().add_memory_region(region);
            } else if let Some(guest_memory) = &self.guest_memory {
                return vu
                    .lock()
                    .unwrap()
                    .update_mem_table(guest_memory.memory().deref());
            }
        }

        Ok(())
    }

    pub fn add_memory_region(
        &mut self,
        region: &Arc<GuestRegionMmap>,
    ) -> result::Result<(), crate::Error> {
        if self.disconnected.load(Ordering::Relaxed) {
            warn!(
                "Skipping add memory region on disconnected dev with socket: {}",
                self.socket_path
            );
        }

        if let Err(e) = self.add_memory_region_internal(region) {
            if e.is_transport_lost() {
                warn!(
                    "Failed updating memory on vhost-user backend for socket {}: {e:?}; \
                     marking device as disconnected",
                    self.socket_path
                );
                self.disconnected.store(true, Ordering::Relaxed);
            } else {
                warn!(
                    "Failed updating memory on vhost-user backend for socket {}: {e:?}",
                    self.socket_path
                );
            }
            return Err(self.memory_update_error(e));
        }
        Ok(())
    }

    pub fn pause(&mut self) -> result::Result<(), MigratableError> {
        if self.disconnected.load(Ordering::Relaxed) {
            return Err(MigratableError::DeviceDisconnected(
                self.socket_path.clone(),
            ));
        }

        if let Some(vu) = &self.vu
            && let Err(e) = vu.lock().unwrap().pause_vhost_user()
        {
            if e.is_transport_lost() {
                self.disconnected.store(true, Ordering::Relaxed);
                return Err(MigratableError::DeviceDisconnected(
                    self.socket_path.clone(),
                ));
            }

            return Err(MigratableError::Pause(anyhow!(
                "Error pausing vhost-user backend for socket {}: {e:?}",
                self.socket_path
            )));
        }
        Ok(())
    }

    fn resume_internal(&mut self) -> result::Result<(), MigratableError> {
        // Skip the resume_vhost_user call if the backend is disconnected. Process the queue
        // interrupts to kick any paused workers.
        if self.disconnected.load(Ordering::Relaxed) {
            return Err(MigratableError::DeviceDisconnected(
                self.socket_path.clone(),
            ));
        }

        if let Some(vu) = &self.vu
            && let Err(e) = vu.lock().unwrap().resume_vhost_user()
        {
            if e.is_transport_lost() {
                self.disconnected.store(true, Ordering::Relaxed);
                return Err(MigratableError::DeviceDisconnected(
                    self.socket_path.clone(),
                ));
            }

            return Err(MigratableError::Resume(anyhow!(
                "Error resuming vhost-user backend for socket {}: {e:?}",
                self.socket_path
            )));
        }

        Ok(())
    }

    pub fn resume(&mut self) -> result::Result<(), MigratableError> {
        let ret = self.resume_internal();

        // Always run the interrupt loop so workers don't get stuck.
        for i in 0..self.vu_num_queues {
            self.virtio_common
                .trigger_interrupt(crate::VirtioInterruptType::Queue(i as u16))
                .ok();
        }

        ret
    }

    pub fn state<C: Default>(
        &self,
        config: C,
    ) -> result::Result<VhostUserState<C>, MigratableError> {
        let mut state = VhostUserState {
            avail_features: self.virtio_common.avail_features,
            acked_features: self.virtio_common.acked_features,
            config,
            acked_protocol_features: self.acked_protocol_features,
            vu_num_queues: self.vu_num_queues,
            ..Default::default()
        };

        if let Some(vu) = &self.vu {
            let mut vu_locked = vu.lock().unwrap();
            if vu_locked.supports_device_state() {
                let (backend_state, vring_bases) = vu_locked.save_backend_state().map_err(|e| {
                    MigratableError::Snapshot(anyhow!("Failed saving backend state: {e:?}"))
                })?;
                state.backend_state = Some(backend_state);
                state.vring_bases = Some(vring_bases);
            }
        }

        Ok(state)
    }

    pub fn snapshot<T>(&mut self, state: &T) -> result::Result<Snapshot, MigratableError>
    where
        T: Serialize,
    {
        let snapshot = Snapshot::new_from_state(state)?;

        if self.migration_started {
            // Local migration does not enable dirty logging.
            if self.dirty_logging {
                self.saved_dirty_log = Some(self.dirty_log()?);
            }
            self.shutdown();
        }

        Ok(snapshot)
    }

    pub fn start_dirty_log(&mut self) -> result::Result<(), MigratableError> {
        if let Some(vu) = &self.vu {
            if let Some(guest_memory) = &self.guest_memory {
                let last_ram_addr = guest_memory.memory().last_addr().raw_value();
                vu.lock()
                    .unwrap()
                    .start_dirty_log(last_ram_addr)
                    .map_err(|e| {
                        MigratableError::StartDirtyLog(anyhow!(
                            "Error starting migration for vhost-user backend: {e:?}"
                        ))
                    })?;
                self.dirty_logging = true;
                Ok(())
            } else {
                Err(MigratableError::StartDirtyLog(anyhow!(
                    "Missing guest memory"
                )))
            }
        } else {
            Ok(())
        }
    }

    pub fn stop_dirty_log(&mut self) -> result::Result<(), MigratableError> {
        if let Some(vu) = &self.vu {
            vu.lock().unwrap().stop_dirty_log().map_err(|e| {
                MigratableError::StopDirtyLog(anyhow!(
                    "Error stopping migration for vhost-user backend: {e:?}"
                ))
            })?;
        }

        self.dirty_logging = false;
        Ok(())
    }

    pub fn dirty_log(&mut self) -> result::Result<MemoryRangeTable, MigratableError> {
        if let Some(vu) = &self.vu {
            if let Some(guest_memory) = &self.guest_memory {
                let last_ram_addr = guest_memory.memory().last_addr().raw_value();
                vu.lock().unwrap().dirty_log(last_ram_addr).map_err(|e| {
                    MigratableError::DirtyLog(anyhow!(
                        "Error retrieving dirty ranges from vhost-user backend: {e:?}"
                    ))
                })
            } else {
                Err(MigratableError::DirtyLog(anyhow!("Missing guest memory")))
            }
        } else {
            Ok(self.saved_dirty_log.take().unwrap_or_default())
        }
    }

    pub fn start_migration(&mut self) -> result::Result<(), MigratableError> {
        self.migration_started = true;
        Ok(())
    }

    pub fn complete_migration(&mut self) -> result::Result<(), MigratableError> {
        self.migration_started = false;
        self.dirty_logging = false;

        // Make sure the device thread is killed in order to prevent from
        // reconnections to the socket.
        if let Some(workers) = self.virtio_common.workers.as_ref() {
            workers.signal_exit().map_err(|e| {
                MigratableError::CompleteMigration(anyhow!(
                    "Error killing vhost-user thread: {e:?}"
                ))
            })?;
        }

        // Drop the vhost-user handler to avoid further calls to fail because
        // the connection with the backend has been closed.
        self.vu = None;

        Ok(())
    }
}
