cubecl_common/device/handle/

base.rs

use crate::device::{DeviceId, DeviceService, ServerUtilitiesHandle};
use alloc::boxed::Box;
use alloc::string::String;
use core::any::Any;
#[cfg(feature = "std")]
use cubecl_environment::future::channel::oneshot;

/// An error returned by a blocking device call.
///
/// When a task panics on a device's runner thread, the panic payload is
/// captured and carried here instead of being logged and dropped.
pub struct CallError {
    failure: Failure,
}

/// Why a device call produced no result.
enum Failure {
    /// The task panicked; the payload is the one `catch_unwind` captured.
    Panicked(Box<dyn Any + Send>),
    /// The service was closed ([`DeviceHandle::close`](super::DeviceHandle::close)) before
    /// the task reached it.
    Closed,
    /// The runner channel disconnected before the task could produce a result.
    Disconnected,
}

impl CallError {
    // Only the `channel` backend (active when `multi_threading` is set) produces these
    // errors.
    /// Builds an error from a panic payload captured by `catch_unwind`.
    #[cfg_attr(not(multi_threading), allow(dead_code))]
    pub(crate) fn from_panic(payload: Box<dyn Any + Send>) -> Self {
        Self {
            failure: Failure::Panicked(payload),
        }
    }

    /// Builds an error for a runner channel that disconnected before producing
    /// a result (no panic payload available).
    #[cfg_attr(not(multi_threading), allow(dead_code))]
    pub(crate) fn disconnected() -> Self {
        Self {
            failure: Failure::Disconnected,
        }
    }

    /// Builds an error for a task that reached a closed service.
    #[cfg_attr(not(multi_threading), allow(dead_code))]
    pub(crate) fn closed() -> Self {
        Self {
            failure: Failure::Closed,
        }
    }

    /// Whether the service was closed before the task reached it, which only happens as the
    /// process exits: the call is refused rather than run on a runtime being torn down.
    pub fn is_closed(&self) -> bool {
        matches!(self.failure, Failure::Closed)
    }

    /// Returns the panic message if the failure was a panic whose payload is a
    /// string (the common case for `panic!`, indexing, `unwrap`, ...).
    pub fn message(&self) -> Option<&str> {
        let Failure::Panicked(payload) = &self.failure else {
            return None;
        };
        if let Some(s) = payload.downcast_ref::<&'static str>() {
            Some(*s)
        } else if let Some(s) = payload.downcast_ref::<String>() {
            Some(s.as_str())
        } else {
            None
        }
    }

    /// Consumes the error and returns the captured panic payload, if any.
    ///
    /// The payload can be handed to [`std::panic::resume_unwind`] to re-raise
    /// the original panic on the caller's thread.
    pub fn into_panic(self) -> Option<Box<dyn Any + Send>> {
        match self.failure {
            Failure::Panicked(payload) => Some(payload),
            Failure::Closed | Failure::Disconnected => None,
        }
    }

    /// Re-raises the failure on the current thread, diverging.
    #[track_caller]
    pub fn resume(self) -> ! {
        match self.failure {
            #[cfg(feature = "std")]
            Failure::Panicked(payload) => std::panic::resume_unwind(payload),
            // std::panic::resume_unwind doesn't exist in a no_std build
            #[cfg(not(feature = "std"))]
            Failure::Panicked(_payload) => {
                panic!("a device task panicked but its payload cannot be re-raised without `std`")
            }
            Failure::Closed => panic!("the device service was closed, as the process exits"),
            Failure::Disconnected => {
                panic!("device runner channel disconnected before producing a result")
            }
        }
    }
}

/// Extension trait for re-raising a [`CallError`] on the caller's thread.
pub trait CallResultExt<R> {
    /// Returns the success value, or re-raises the original panic captured in the
    /// [`CallError`] via [`CallError::resume`].
    ///
    /// Use this instead of [`Result::unwrap`] to panics with the original payload itself.
    fn unwrap_or_resume(self) -> R;
}

impl<R> CallResultExt<R> for Result<R, CallError> {
    #[track_caller]
    fn unwrap_or_resume(self) -> R {
        match self {
            Ok(value) => value,
            Err(err) => err.resume(),
        }
    }
}

/// A call [spawned](super::DeviceHandle::spawn) on a device that may still be running.
#[must_use = "the call's result is lost unless it is joined"]
pub struct SpawnedCall<R> {
    state: SpawnedState<R>,
}

enum SpawnedState<R> {
    /// The result is already known.
    Done(Result<R, CallError>),
    #[cfg(feature = "std")]
    Running(oneshot::Receiver<Result<R, CallError>>),
}

impl<R> SpawnedCall<R> {
    pub(crate) fn done(result: Result<R, CallError>) -> Self {
        Self {
            state: SpawnedState::Done(result),
        }
    }

    #[cfg(feature = "std")]
    pub(crate) fn running(receiver: oneshot::Receiver<Result<R, CallError>>) -> Self {
        Self {
            state: SpawnedState::Running(receiver),
        }
    }

    /// Blocks until the call returns, failing as
    /// [`submit_blocking`](super::DeviceHandle::submit_blocking) does.
    pub fn join(self) -> Result<R, CallError> {
        match self.state {
            SpawnedState::Done(result) => result,
            #[cfg(feature = "std")]
            SpawnedState::Running(receiver) => receiver
                .recv()
                .unwrap_or_else(|_| Err(CallError::disconnected())),
        }
    }
}

impl core::fmt::Debug for CallError {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match (&self.failure, self.message()) {
            (_, Some(message)) => write!(
                f,
                "CallError(task panicked on device runner thread: {message})"
            ),
            (Failure::Panicked(_), None) => f.write_str(
                "CallError(task panicked on device runner thread with a non-string payload)",
            ),
            (Failure::Closed, None) => {
                f.write_str("CallError(the device service was closed, as the process exits)")
            }
            (Failure::Disconnected, None) => f.write_str(
                "CallError(device runner channel disconnected before producing a result)",
            ),
        }
    }
}

#[derive(new, Clone, Debug)]
/// Error when creating a [`DeviceService`].
pub struct ServiceCreationError {
    #[allow(dead_code)] // Debug uses it.
    reason: alloc::string::String,
}

#[derive(new, Clone, Debug)]
/// Error when closing a [`DeviceService`]: the service may still be serving.
pub struct ServiceCloseError {
    #[allow(dead_code)] // Debug uses it.
    reason: alloc::string::String,
}

/// How a service's state is reached: the transport under a [`DeviceHandle`].
///
/// A transport holds no service type. It hands a task the boxed state as
/// `&mut dyn Any`; the handle above it knows what that is.
pub trait DeviceHandleSpec: Sized + Clone {
    /// If functions block the current thread even if they are non-blocking.
    const BLOCKING: bool;

    /// Registers `service` for the device and returns a transport to it.
    ///
    /// If a runner thread for this `device_id` does not exist, it will be spawned.
    fn insert<S: DeviceService>(
        device_id: DeviceId,
        service: S,
    ) -> Result<Self, ServiceCreationError>;

    /// Creates or retrieves a transport for the given device ID.
    ///
    /// If a runner thread for this `device_id` does not exist, it will be spawned.
    fn new<S: DeviceService>(device_id: DeviceId) -> Self;

    /// Retrieves the device ID for this handle.
    fn device_id(&self) -> DeviceId;

    /// Retrieves the server utilities for this thread.
    fn utilities(&self) -> ServerUtilitiesHandle;

    /// Doesn't flush the service state, but flushes any task enqueued in the communication
    /// channel.
    ///
    /// # Notes
    ///
    /// This is a no-op for blocking handles.
    fn flush_queue(&self);

    fn submit_blocking<'a, R: Send, T: FnOnce(&mut dyn Any) -> R + Send + 'a>(
        &self,
        task: T,
    ) -> Result<R, CallError>;

    fn submit<T: FnOnce(&mut dyn Any) + Send + 'static>(&self, task: T);

    /// Runs `task` without waiting for it to return.
    ///
    /// A transport whose calls run on the caller's thread has run it by the time this returns.
    fn spawn<R: Send + 'static, T: FnOnce(&mut dyn Any) -> R + Send + 'static>(
        &self,
        task: T,
    ) -> SpawnedCall<R>;

    fn exclusive<R: Send, T: FnOnce() -> R + Send>(&self, task: T) -> Result<R, CallError>;

    fn shutdown(device_id: DeviceId) {
        let _ = device_id;
    }

    /// Stops serving `S` on `device_id` though handles to it remain, dropping it where it
    /// lives; later calls to it answer [`CallError::is_closed`], and it is not started again.
    ///
    /// A transport whose calls run on the caller's thread has no runner to close it on, and
    /// keeps serving: only the channel transport serves a native runtime.
    fn close<S: DeviceService>(device_id: DeviceId) -> Result<(), ServiceCloseError> {
        let _ = device_id;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::CallError;
    use alloc::boxed::Box;
    use core::any::Any;

    /// D.1 — A disconnected error carries neither a message nor a payload, which is
    /// what distinguishes it from a non-string panic (both have `message() == None`).
    #[test]
    fn test_disconnected_has_no_message_or_payload() {
        let err = CallError::disconnected();
        assert_eq!(err.message(), None);
        assert!(err.into_panic().is_none());
    }

    /// A closed service is not a panic and not a dead runner: callers map exactly this case to
    /// an error, and re-raise the other two.
    #[test]
    fn test_closed_is_neither_a_panic_nor_a_disconnect() {
        let err = CallError::closed();
        assert!(err.is_closed());
        assert!(!CallError::disconnected().is_closed());
        assert!(alloc::format!("{err:?}").contains("closed"));
        assert!(err.into_panic().is_none());
    }

    /// D.2 — Debug of a string payload includes the panic message.
    #[test]
    fn test_debug_string_payload_includes_message() {
        let err = CallError::from_panic(Box::new("oops") as Box<dyn Any + Send>);
        let debug = alloc::format!("{err:?}");
        assert!(debug.contains("task panicked on device runner thread"));
        assert!(debug.contains("oops"));
    }

    /// D.3 — Debug of a non-string payload reports it as such.
    #[test]
    fn test_debug_non_string_payload() {
        let err = CallError::from_panic(Box::new(123u8) as Box<dyn Any + Send>);
        let debug = alloc::format!("{err:?}");
        assert!(debug.contains("non-string payload"), "got: {debug}");
    }

    /// D.4 — Debug of a disconnected error mentions the disconnection.
    #[test]
    fn test_debug_disconnected() {
        let err = CallError::disconnected();
        let debug = alloc::format!("{err:?}");
        assert!(debug.contains("disconnected"), "got: {debug}");
    }
}