diff --git a/components/spider-task-executor/src/error.rs b/components/spider-task-executor/src/error.rs index f25b4769e..f6a42a468 100644 --- a/components/spider-task-executor/src/error.rs +++ b/components/spider-task-executor/src/error.rs @@ -40,6 +40,10 @@ pub enum ExecutorError { #[error("task execution failed: {0}")] TaskError(#[from] TdlError), + /// The package's init function returned a [`TdlError`] during load. + #[error("failed to initialize TDL package: {0}")] + PackageInitError(TdlError), + /// The msgpack-encoded error payload returned by a failing task could not be decoded back into /// a [`TdlError`]. #[error("failed to deserialize error payload: {0}")] diff --git a/components/spider-task-executor/src/manager.rs b/components/spider-task-executor/src/manager.rs index 80194085b..ef3959daa 100644 --- a/components/spider-task-executor/src/manager.rs +++ b/components/spider-task-executor/src/manager.rs @@ -51,7 +51,9 @@ impl TdlPackage { /// 2. Look up [`SYM_GET_VERSION`], call it, and verify the returned [`Version`] is compatible /// with [`Version::SPIDER_TDL`]. /// 3. Look up [`SYM_GET_NAME`], call it, and decode the returned bytes as UTF-8. - /// 4. Look up [`SYM_EXECUTE`] and cache the fn pointer for per-task dispatch. + /// 4. Look up [`SYM_INIT`] (optional) and, if present, call it once; a returned [`TdlError`] + /// aborts the load. + /// 5. Look up [`SYM_EXECUTE`] and cache the fn pointer for per-task dispatch. /// /// # Returns /// @@ -63,10 +65,13 @@ impl TdlPackage { /// /// * [`ExecutorError::IncompatibleVersion`] if the package was built against an incompatible /// `spider-tdl` release. + /// * [`ExecutorError::PackageInitError`] if the package's init function returned a + /// [`TdlError`]. /// * Forwards [`Library::new`]'s return values on failure. /// * Forwards [`Library::get`]'s return values on failure for loading [`SYM_GET_VERSION`], /// [`SYM_GET_NAME`], or [`SYM_EXECUTE`]. /// * Forwards [`CCharArray::as_utf8`]'s return values on failure. + /// * Forwards [`rmp_serde::from_slice`]'s return values on failure. pub fn load(path: &Path) -> Result { // SAFETY: `Library::new` runs the dylib's initializers. Spider's design treats every TDL // package as trusted code installed by the operator, so this is the unsafety boundary for @@ -97,6 +102,19 @@ impl TdlPackage { }; let name = name_array.as_utf8()?.to_owned(); + // The init symbol is optional: packages built against an older `spider-tdl` may not export + // it, so a failed lookup is treated as "no init to run". + let init: Option = + unsafe { library.get::(SYM_INIT).ok().map(|symbol| *symbol) }; + if let Some(init) = init { + // SAFETY: see the SAFETY comment on the version lookup above. + let init_result = unsafe { init() }; + if let Err(error_bytes) = init_result.into_result() { + let err: TdlError = rmp_serde::from_slice(&error_bytes)?; + return Err(ExecutorError::PackageInitError(err)); + } + } + // SAFETY: see the SAFETY comment on the version lookup above. We deref the borrowed // `Symbol` to copy out the underlying fn pointer (`ExecuteFn` is `Copy`); the // pointer remains valid for as long as `library` stays loaded, which is the entire @@ -234,6 +252,9 @@ type GetVersionFn = unsafe extern "C" fn() -> Version; /// FFI signature of `__spider_tdl_package_get_name`. type GetNameFn = unsafe extern "C" fn() -> CCharArray<'static>; +/// FFI signature of `__spider_tdl_package_init`. +type InitFn = unsafe extern "C" fn() -> TaskExecutionResult; + /// FFI signature of `__spider_tdl_package_execute`. type ExecuteFn = unsafe extern "C" fn(CCharArray<'_>, CByteArray<'_>, CByteArray<'_>) -> TaskExecutionResult; @@ -244,6 +265,9 @@ const SYM_GET_VERSION: &[u8] = b"__spider_tdl_package_get_version\0"; /// FFI symbol name (NUL-terminated) for the package's declared name. const SYM_GET_NAME: &[u8] = b"__spider_tdl_package_get_name\0"; +/// FFI symbol name (NUL-terminated) for the package's optional init function. +const SYM_INIT: &[u8] = b"__spider_tdl_package_init\0"; + /// FFI symbol name (NUL-terminated) for the per-task dispatcher. const SYM_EXECUTE: &[u8] = b"__spider_tdl_package_execute\0"; diff --git a/components/spider-tdl/src/register.rs b/components/spider-tdl/src/register.rs index 846ef0d7a..3ed68dbb4 100644 --- a/components/spider-tdl/src/register.rs +++ b/components/spider-tdl/src/register.rs @@ -8,9 +8,10 @@ //! * A single assertion that rejects duplicate `NAME`s with a `const_eval` panic at build time. The //! hash map itself is built only once at runtime; uniqueness is enforced at compile time and is //! independent of the runtime structure. -//! * Three `extern "C"` entry points consumed by the package manager via `dlsym`: +//! * Four `extern "C"` entry points consumed by the package manager via `dlsym`: //! * `__spider_tdl_package_get_version` //! * `__spider_tdl_package_get_name` +//! * `__spider_tdl_package_init` //! * `__spider_tdl_package_execute` //! //! The helpers in this module are public only so they are reachable from macro expansions in @@ -19,7 +20,7 @@ use crate::TdlError; -/// Registers a TDL package's tasks and exports the three C-FFI entry points consumed by the task +/// Registers a TDL package's tasks and exports the four C-FFI entry points consumed by the task /// executor. /// /// Invoke once per package, at module scope: @@ -27,6 +28,7 @@ use crate::TdlError; /// ```ignore /// spider_tdl::register_tdl_package! { /// package_name: "complex-number", +/// init: my_init_fn, /// tasks: [add, sub, mul, div, always_fail], /// } /// ``` @@ -34,6 +36,9 @@ use crate::TdlError; /// Each entry in `tasks` must name a type that implements [`Task`](crate::Task) (typically a marker /// struct produced by the `#[task]` attribute macro). /// +/// The optional `init` field names a `fn() -> Result<(), TdlError>` that is run once when the +/// package is loaded. When omitted, it defaults to a no-op. +/// /// # Name Uniqueness /// /// All tasks registered within a single package must have distinct `NAME`s. This check happens in @@ -45,13 +50,41 @@ use crate::TdlError; /// package was compiled against. /// * `__spider_tdl_package_get_name` returns the package name passed to the macro, as a borrowed /// [`CCharArray`](crate::ffi::CCharArray). +/// * `__spider_tdl_package_init` runs the package's init function (or a no-op when none was +/// provided) and returns a [`TaskExecutionResult`](crate::ffi::TaskExecutionResult) carrying an +/// empty success buffer or the msgpack-encoded [`TdlError`](crate::TdlError). /// * `__spider_tdl_package_execute` dispatches a task by name for execution and returns a /// [`TaskExecutionResult`](crate::ffi::TaskExecutionResult). #[macro_export] macro_rules! register_tdl_package { + ( + package_name: $package_name:expr, + init: $init:path, + tasks: [$($task:path),* $(,)?] $(,)? + ) => { + $crate::register_tdl_package! { + @internal + package_name: $package_name, + init: $init, + tasks: [$($task),*], + } + }; ( package_name: $package_name:expr, tasks: [$($task:path),* $(,)?] $(,)? + ) => { + $crate::register_tdl_package! { + @internal + package_name: $package_name, + init: $crate::register::noop_package_init, + tasks: [$($task),*], + } + }; + ( + @internal + package_name: $package_name:expr, + init: $init:path, + tasks: [$($task:path),* $(,)?] $(,)? ) => { const __SPIDER_TDL_PACKAGE_NAME: &str = $package_name; @@ -88,6 +121,20 @@ macro_rules! register_tdl_package { $crate::ffi::CCharArray::from_utf8(__SPIDER_TDL_PACKAGE_NAME) } + #[unsafe(no_mangle)] + pub extern "C" fn __spider_tdl_package_init() -> $crate::ffi::TaskExecutionResult { + let init_fn: fn() -> ::std::result::Result<(), $crate::TdlError> = $init; + match init_fn() { + ::std::result::Result::Ok(()) => { + $crate::ffi::TaskExecutionResult::from_outputs(::std::vec::Vec::new()) + } + ::std::result::Result::Err(err) => { + let bytes = $crate::register::serialize_error_payload(&err); + $crate::ffi::TaskExecutionResult::from_error(bytes) + } + } + } + #[unsafe(no_mangle)] pub extern "C" fn __spider_tdl_package_execute( name: $crate::ffi::CCharArray<'_>, @@ -178,6 +225,12 @@ pub fn serialize_error_payload(err: &TdlError) -> Vec { rmp_serde::to_vec(err).expect("failed to serialize `TdlError` as msgpack") } +/// Default package init function used when a package registers no `init` hook. +#[doc(hidden)] +pub const fn noop_package_init() -> Result<(), TdlError> { + Ok(()) +} + #[cfg(test)] mod tests { use std::panic; @@ -214,6 +267,11 @@ mod tests { ); } + #[test] + fn noop_package_init_returns_ok() { + assert_eq!(noop_package_init(), Ok(())); + } + #[test] fn assert_unique_task_names_is_const_callable() { /// Confirms `assert_unique_task_names` is usable in a `const` context (i.e. compile-time). diff --git a/tests/huntsman/integration-test-tasks/src/lib.rs b/tests/huntsman/integration-test-tasks/src/lib.rs index 23bf00533..ebff39e03 100644 --- a/tests/huntsman/integration-test-tasks/src/lib.rs +++ b/tests/huntsman/integration-test-tasks/src/lib.rs @@ -1,6 +1,6 @@ //! Test TDL package used by the `task-executor` integration tests. //! -//! Exposes five tasks that exercise distinct executor code paths: +//! Exposes six tasks that exercise distinct executor code paths: //! //! * [`task_decl::fibonacci`] — basic compute + correctness. //! * [`task_decl::always_fail`] — in-task error reporting. @@ -11,6 +11,10 @@ //! serde cost, while the parent-side delta isolates IPC framing cost. //! * [`task_decl::assert_outputs_sum_zero`] — commit task: reads the job's task-graph outputs from //! its [`TaskContext`](spider_tdl::TaskContext) and asserts the `i64` outputs sum to zero. +//! * [`task_decl::assert_initialized`] — confirms the package's init hook ran on load. +//! +//! The package also registers an init hook ([`task_decl::package_init`]) that runs once when the +//! package is loaded. /// The constant sleep duration used by [`task_decl::sleep_and_echo`]. /// @@ -19,6 +23,8 @@ pub const INSTRUMENT_SLEEP_US: u64 = 50; mod task_decl { + use std::sync::atomic::AtomicBool; + use std::sync::atomic::Ordering; use std::thread::sleep; use std::time::Duration; @@ -28,6 +34,25 @@ mod task_decl { use crate::INSTRUMENT_SLEEP_US; + /// Package init hook. Records that init ran, or fails when `SPIDER_TEST_TDL_INIT_SHOULD_FAIL` + /// is set (used by the executor's init integration test). + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`TdlError::ExecutionError`] if `SPIDER_TEST_TDL_INIT_SHOULD_FAIL` is set in the + /// environment. + pub fn package_init() -> Result<(), TdlError> { + if std::env::var_os("SPIDER_TEST_TDL_INIT_SHOULD_FAIL").is_some() { + return Err(TdlError::ExecutionError( + "integration_test_tasks: init failure requested".to_owned(), + )); + } + INITIALIZED.store(true, Ordering::SeqCst); + Ok(()) + } + /// Computes the `index`-th Fibonacci number with a deliberately naive recursive /// implementation so the call has measurable CPU cost for the overhead benchmark. #[task(name = "fibonacci")] @@ -93,15 +118,38 @@ mod task_decl { } Ok(()) } + + /// Succeeds only if the package's init hook set the [`INITIALIZED`] flag on load. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`TdlError::ExecutionError`] if the init hook did not run before this task was dispatched. + #[task(name = "assert_initialized")] + pub fn assert_initialized(_ctx: TaskContext) -> Result<(), TdlError> { + if INITIALIZED.load(Ordering::SeqCst) { + Ok(()) + } else { + Err(TdlError::ExecutionError( + "assert_initialized: package init hook did not run".to_owned(), + )) + } + } + + /// Set by [`package_init`] so [`assert_initialized`] can confirm the init hook ran on load. + static INITIALIZED: AtomicBool = AtomicBool::new(false); } spider_tdl::register_tdl_package! { package_name: "integration_test_tasks", + init: task_decl::package_init, tasks: [ task_decl::fibonacci, task_decl::always_fail, task_decl::always_panic, task_decl::sleep_and_echo, task_decl::assert_outputs_sum_zero, + task_decl::assert_initialized, ], } diff --git a/tests/huntsman/tdl-integration/Cargo.toml b/tests/huntsman/tdl-integration/Cargo.toml index 73d10a635..608ef917f 100644 --- a/tests/huntsman/tdl-integration/Cargo.toml +++ b/tests/huntsman/tdl-integration/Cargo.toml @@ -12,6 +12,10 @@ path = "src/lib.rs" name = "complex" path = "tests/complex.rs" +[[test]] +name = "init" +path = "tests/init.rs" + [dev-dependencies] anyhow = "1.0.98" huntsman-complex-types = { path = "../../../examples/huntsman/complex/types" } diff --git a/tests/huntsman/tdl-integration/tests/init.rs b/tests/huntsman/tdl-integration/tests/init.rs new file mode 100644 index 000000000..1bc615d6c --- /dev/null +++ b/tests/huntsman/tdl-integration/tests/init.rs @@ -0,0 +1,84 @@ +//! End-to-end tests for the TDL package init hook against the `integration-test-tasks` cdylib. + +use spider_core::types::id::JobId; +use spider_core::types::id::ResourceGroupId; +use spider_core::types::id::TaskId; +use spider_core::types::io::TaskInputsSerializer; +use spider_task_executor::ExecutorError; +use spider_task_executor::TdlPackageManager; +use spider_tdl::TaskContext; +use spider_tdl::TdlError; + +/// # Returns +/// +/// The absolute path of the staged `integration-test-tasks` cdylib, derived from the +/// `SPIDER_TDL_PACKAGE_DIR` environment variable. +fn package_path() -> std::path::PathBuf { + const PACKAGE_NAME: &str = "integration_test_tasks"; + let dir = std::env::var_os("SPIDER_TDL_PACKAGE_DIR") + .map(std::path::PathBuf::from) + .expect("`SPIDER_TDL_PACKAGE_DIR` not set"); + dir.join(PACKAGE_NAME).join(format!("lib{PACKAGE_NAME}.so")) +} + +/// # Returns +/// +/// An encoded task context for testing. +fn encode_ctx() -> Vec { + let ctx = TaskContext::new( + JobId::random(), + TaskId::Index(0), + 1, + ResourceGroupId::random(), + None, + ) + .expect("failed to build `TaskContext`"); + rmp_serde::to_vec(&ctx).expect("failed to serialize `TaskContext`") +} + +/// # Returns +/// +/// A wire-format-encoded empty input. +fn encode_no_inputs() -> Vec { + TaskInputsSerializer::new().release() +} + +#[test] +#[ignore = "requires `integration-test-tasks` cdylib"] +fn init_hook_runs_before_task_dispatch() -> anyhow::Result<()> { + let path = package_path(); + let mut manager = TdlPackageManager::new(); + manager.load(&path)?; + let pkg = manager + .get("integration_test_tasks") + .expect("package should be loaded"); + pkg.execute_task("assert_initialized", &encode_ctx(), &encode_no_inputs())?; + Ok(()) +} + +#[test] +#[ignore = "requires `integration-test-tasks` cdylib"] +fn failing_init_aborts_load() -> anyhow::Result<()> { + const ENV_SPIDER_TEST_TDL_INIT_SHOULD_FAIL: &str = "SPIDER_TEST_TDL_INIT_SHOULD_FAIL"; + // SAFETY: `cargo nextest` (see taskfiles/test.yaml) runs each test in its own process + // (https://nexte.st/docs/design/why-process-per-test/), so this can't leak into other tests; + // and no other thread exists in this process yet, so nothing reads the environment + // concurrently. Under plain `cargo test` this would be unsound and could fail unrelated tests. + // The init hook failing with this var set is the behavior under test. + unsafe { std::env::set_var(ENV_SPIDER_TEST_TDL_INIT_SHOULD_FAIL, "1") }; + let path = package_path(); + let mut manager = TdlPackageManager::new(); + let err = manager + .load(&path) + .expect_err("load should fail when init errors"); + // SAFETY: see above. + unsafe { std::env::remove_var(ENV_SPIDER_TEST_TDL_INIT_SHOULD_FAIL) }; + let ExecutorError::PackageInitError(TdlError::ExecutionError(msg)) = &err else { + panic!("unexpected error: {err:?}"); + }; + anyhow::ensure!( + msg.contains("init failure requested"), + "unexpected message: {msg}" + ); + Ok(()) +}