clusterflux-public/crates/disasmer-sdk/src/lib.rs

use serde::{Deserialize, Serialize};

pub use disasmer_macros::{main, task};

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EntrypointDescriptor {
    pub name: &'static str,
    pub function: &'static str,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct TaskDescriptor {
    pub name: &'static str,
    pub function: &'static str,
    pub remotely_startable: bool,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RegisteredProgram {
    pub entrypoints: &'static [EntrypointDescriptor],
    pub tasks: &'static [TaskDescriptor],
}

impl RegisteredProgram {
    pub fn select_entrypoint(&self, name: &str) -> Option<EntrypointDescriptor> {
        self.entrypoints
            .iter()
            .copied()
            .find(|entrypoint| entrypoint.name == name)
    }

    pub fn task(&self, name: &str) -> Option<TaskDescriptor> {
        self.tasks.iter().copied().find(|task| task.name == name)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct SourceSnapshot {
    pub digest: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Blob {
    pub digest: String,
}

#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct Artifact {
    pub id: String,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct EnvRef {
    pub name: &'static str,
}

impl EnvRef {
    pub const fn new_static(name: &'static str) -> Self {
        Self { name }
    }
}

#[macro_export]
macro_rules! env {
    ($name:literal) => {
        $crate::EnvRef::new_static($name)
    };
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TaskArgKind {
    SmallSerialized,
    Handle,
}

pub trait TaskArg: Serialize {
    fn task_arg_kind(&self) -> TaskArgKind {
        TaskArgKind::SmallSerialized
    }
}

macro_rules! small_task_arg {
    ($($ty:ty),+ $(,)?) => {
        $(
            impl TaskArg for $ty {}
        )+
    };
}

small_task_arg!(
    (),
    bool,
    char,
    String,
    i8,
    i16,
    i32,
    i64,
    isize,
    u8,
    u16,
    u32,
    u64,
    usize,
    f32,
    f64,
);

impl TaskArg for SourceSnapshot {
    fn task_arg_kind(&self) -> TaskArgKind {
        TaskArgKind::Handle
    }
}

impl TaskArg for Blob {
    fn task_arg_kind(&self) -> TaskArgKind {
        TaskArgKind::Handle
    }
}

impl TaskArg for Artifact {
    fn task_arg_kind(&self) -> TaskArgKind {
        TaskArgKind::Handle
    }
}

impl<T> TaskArg for Option<T> where T: TaskArg {}

impl<T> TaskArg for Vec<T> where T: TaskArg {}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TaskArgBudget {
    pub max_inline_bytes: usize,
}

impl Default for TaskArgBudget {
    fn default() -> Self {
        Self {
            max_inline_bytes: 64 * 1024,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TaskArgValidation {
    pub inline_bytes: usize,
    pub kind: TaskArgKind,
}

#[derive(Clone, Debug, PartialEq, Eq)]
pub enum TaskArgError {
    Serialization(String),
    TooLarge { size: usize, limit: usize },
    HostOnly { type_name: &'static str },
}

impl std::fmt::Display for TaskArgError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Serialization(error) => write!(f, "task argument could not serialize: {error}"),
            Self::TooLarge { size, limit } => write!(
                f,
                "task argument is {size} bytes; inline task arguments are limited to {limit} bytes, use SourceSnapshot, Blob, Artifact, or VFS handles"
            ),
            Self::HostOnly { type_name } => write!(
                f,
                "task boundary value `{type_name}` is host-only; use small serialized data or handles"
            ),
        }
    }
}

impl std::error::Error for TaskArgError {}

pub fn validate_task_arg<T>(
    value: &T,
    budget: TaskArgBudget,
) -> Result<TaskArgValidation, TaskArgError>
where
    T: TaskArg + ?Sized,
{
    let bytes = serde_json::to_vec(value)
        .map_err(|error| TaskArgError::Serialization(error.to_string()))?;
    let kind = value.task_arg_kind();
    if kind == TaskArgKind::SmallSerialized && bytes.len() > budget.max_inline_bytes {
        return Err(TaskArgError::TooLarge {
            size: bytes.len(),
            limit: budget.max_inline_bytes,
        });
    }
    Ok(TaskArgValidation {
        inline_bytes: bytes.len(),
        kind,
    })
}

pub fn reject_host_only_task_arg<T: ?Sized>() -> TaskArgError {
    TaskArgError::HostOnly {
        type_name: std::any::type_name::<T>(),
    }
}

pub mod spawn {
    use std::sync::atomic::{AtomicU64, Ordering};
    use std::sync::{Mutex, OnceLock};

    use crate::{validate_task_arg, EnvRef, TaskArg, TaskArgBudget, TaskArgError};

    static NEXT_THREAD_ID: AtomicU64 = AtomicU64::new(1);
    static RUNTIME_THREADS: OnceLock<Mutex<Vec<RuntimeSpawnEvent>>> = OnceLock::new();

    #[derive(Clone, Debug, PartialEq, Eq)]
    pub struct RuntimeSpawnEvent {
        pub virtual_thread_id: u64,
        pub name: &'static str,
        pub env: Option<EnvRef>,
        pub debugger_visible: bool,
    }

    fn runtime_threads() -> &'static Mutex<Vec<RuntimeSpawnEvent>> {
        RUNTIME_THREADS.get_or_init(|| Mutex::new(Vec::new()))
    }

    fn register_runtime_thread(
        virtual_thread_id: u64,
        name: &'static str,
        env: Option<EnvRef>,
    ) -> RuntimeSpawnEvent {
        let event = RuntimeSpawnEvent {
            virtual_thread_id,
            name,
            env,
            debugger_visible: true,
        };
        runtime_threads().lock().unwrap().push(event.clone());
        event
    }

    pub fn drain_runtime_spawn_events() -> Vec<RuntimeSpawnEvent> {
        runtime_threads().lock().unwrap().drain(..).collect()
    }

    pub fn runtime_spawn_events() -> Vec<RuntimeSpawnEvent> {
        runtime_threads().lock().unwrap().clone()
    }

    pub fn task<F, R>(entry: F) -> TaskBuilder<F, R>
    where
        F: FnOnce() -> R,
        R: TaskArg,
    {
        TaskBuilder {
            entry: Some(entry),
            env: None,
            name: "task",
        }
    }

    pub fn task_with_arg<A, F, R>(arg: A, entry: F) -> TaskWithArgBuilder<A, F, R>
    where
        A: TaskArg,
        F: FnOnce(A) -> R,
        R: TaskArg,
    {
        TaskWithArgBuilder {
            arg: Some(arg),
            entry: Some(entry),
            env: None,
            name: "task",
            arg_budget: TaskArgBudget::default(),
        }
    }

    pub struct TaskBuilder<F, R>
    where
        F: FnOnce() -> R,
        R: TaskArg,
    {
        entry: Option<F>,
        env: Option<EnvRef>,
        name: &'static str,
    }

    impl<F, R> TaskBuilder<F, R>
    where
        F: FnOnce() -> R,
        R: TaskArg,
    {
        pub fn env(mut self, env: EnvRef) -> Self {
            self.env = Some(env);
            self
        }

        pub fn name(mut self, name: &'static str) -> Self {
            self.name = name;
            self
        }

        pub async fn start(mut self) -> TaskHandle<R> {
            let id = NEXT_THREAD_ID.fetch_add(1, Ordering::SeqCst);
            let runtime_event = register_runtime_thread(id, self.name, self.env);
            let entry = self.entry.take().expect("task entry used once");
            TaskHandle {
                virtual_thread_id: id,
                name: self.name,
                env: self.env,
                debugger_visible: runtime_event.debugger_visible,
                result: Some(entry()),
            }
        }
    }

    pub struct TaskWithArgBuilder<A, F, R>
    where
        A: TaskArg,
        F: FnOnce(A) -> R,
        R: TaskArg,
    {
        arg: Option<A>,
        entry: Option<F>,
        env: Option<EnvRef>,
        name: &'static str,
        arg_budget: TaskArgBudget,
    }

    impl<A, F, R> TaskWithArgBuilder<A, F, R>
    where
        A: TaskArg,
        F: FnOnce(A) -> R,
        R: TaskArg,
    {
        pub fn env(mut self, env: EnvRef) -> Self {
            self.env = Some(env);
            self
        }

        pub fn name(mut self, name: &'static str) -> Self {
            self.name = name;
            self
        }

        pub fn arg_budget(mut self, arg_budget: TaskArgBudget) -> Self {
            self.arg_budget = arg_budget;
            self
        }

        pub async fn start(mut self) -> Result<TaskHandle<R>, TaskArgError> {
            let arg = self.arg.take().expect("task argument used once");
            validate_task_arg(&arg, self.arg_budget)?;
            let id = NEXT_THREAD_ID.fetch_add(1, Ordering::SeqCst);
            let runtime_event = register_runtime_thread(id, self.name, self.env);
            let entry = self.entry.take().expect("task entry used once");
            Ok(TaskHandle {
                virtual_thread_id: id,
                name: self.name,
                env: self.env,
                debugger_visible: runtime_event.debugger_visible,
                result: Some(entry(arg)),
            })
        }
    }

    pub struct TaskHandle<R> {
        virtual_thread_id: u64,
        name: &'static str,
        env: Option<EnvRef>,
        debugger_visible: bool,
        result: Option<R>,
    }

    impl<R> TaskHandle<R>
    where
        R: TaskArg,
    {
        pub fn virtual_thread_id(&self) -> u64 {
            self.virtual_thread_id
        }

        pub fn name(&self) -> &'static str {
            self.name
        }

        pub fn env(&self) -> Option<EnvRef> {
            self.env
        }

        pub fn debugger_visible(&self) -> bool {
            self.debugger_visible
        }

        pub async fn join(mut self) -> R {
            self.result.take().expect("task joined once")
        }
    }
}

#[cfg(test)]
mod tests {
    use futures_executor::block_on;

    #[test]
    fn env_macro_creates_logical_environment_reference() {
        let env = crate::env!("linux");

        assert_eq!(env.name, "linux");
    }

    #[test]
    fn spawn_task_start_join_returns_small_result() {
        let result = block_on(async {
            let handle = crate::spawn::task(|| 42)
                .name("compile linux")
                .env(crate::env!("linux"))
                .start()
                .await;
            assert_eq!(handle.name(), "compile linux");
            assert_eq!(handle.env().unwrap().name, "linux");
            assert!(handle.debugger_visible());
            handle.join().await
        });

        assert_eq!(result, 42);
    }

    #[test]
    fn spawn_task_start_registers_debugger_visible_runtime_thread() {
        let handle = block_on(async {
            crate::spawn::task(|| 7_u32)
                .name("sdk-runtime-thread-test")
                .env(crate::env!("linux"))
                .start()
                .await
        });
        let events = crate::spawn::runtime_spawn_events();
        let event = events
            .iter()
            .find(|event| event.virtual_thread_id == handle.virtual_thread_id())
            .expect("spawn runtime event should exist for task handle");

        assert!(handle.debugger_visible());
        assert!(event.debugger_visible);
        assert_eq!(event.name, "sdk-runtime-thread-test");
        assert_eq!(event.env.unwrap().name, "linux");
        assert_eq!(block_on(handle.join()), 7);
    }

    #[test]
    fn task_arg_validation_allows_handles_and_rejects_oversized_inline_values() {
        let artifact = crate::Artifact {
            id: "artifact://build/app".to_owned(),
        };
        let artifact_validation = crate::validate_task_arg(
            &artifact,
            crate::TaskArgBudget {
                max_inline_bytes: 4,
            },
        )
        .unwrap();
        assert_eq!(artifact_validation.kind, crate::TaskArgKind::Handle);

        let bytes = vec![1_u8, 2, 3, 4, 5];
        let error = crate::validate_task_arg(
            &bytes,
            crate::TaskArgBudget {
                max_inline_bytes: 4,
            },
        )
        .unwrap_err();
        assert!(matches!(error, crate::TaskArgError::TooLarge { .. }));
    }

    #[test]
    fn spawn_task_with_arg_rejects_large_inline_argument_before_dispatch() {
        let dispatched = std::cell::Cell::new(false);
        let result = block_on(async {
            crate::spawn::task_with_arg(vec![1_u8, 2, 3, 4, 5], |_| {
                dispatched.set(true);
                42_u32
            })
            .arg_budget(crate::TaskArgBudget {
                max_inline_bytes: 4,
            })
            .start()
            .await
        });

        assert!(matches!(result, Err(crate::TaskArgError::TooLarge { .. })));
        assert!(!dispatched.get());
    }

    #[test]
    fn spawn_task_with_arg_allows_runtime_handles_under_inline_budget() {
        let artifact = crate::Artifact {
            id: "artifact://build/app".to_owned(),
        };
        let result = block_on(async {
            let handle = crate::spawn::task_with_arg(artifact, |artifact| artifact.id)
                .name("publish artifact")
                .env(crate::env!("linux"))
                .arg_budget(crate::TaskArgBudget {
                    max_inline_bytes: 4,
                })
                .start()
                .await
                .unwrap();

            assert_eq!(handle.name(), "publish artifact");
            handle.join().await
        });

        assert_eq!(result, "artifact://build/app");
    }

    #[test]
    fn host_only_task_arg_error_names_rejected_type() {
        let error = crate::reject_host_only_task_arg::<*const u8>();

        assert!(error.to_string().contains("*const u8"));
        assert!(error.to_string().contains("host-only"));
    }
}