runa/crates/runa-rune/src/execution.rs

use std::cell::Cell;
use std::future::{Future, poll_fn};
use std::num::NonZeroUsize;
use std::pin::Pin;
use std::rc::Rc;
use std::task::{Context, Poll};

use rune::runtime::{Args, VmError};
use rune::{FromValue, Vm};

use crate::ExperimentResult;

#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)]
pub struct InvocationProgress {
    pub vm_steps: u64,
    pub instruction_pointer: usize,
    pub call_depth: usize,
}

/// A worker-local execution slice; native callbacks must remain nonblocking.
#[derive(Debug)]
pub enum Slice<T> {
    Completed(T),
    Awaiting,
    BudgetExhausted,
    Failed(Box<dyn std::error::Error>),
}

type ExecutionFuture<T> = Pin<Box<dyn Future<Output = ExperimentResult<T>>>>;

pub struct Invocation<T> {
    future: Option<ExecutionFuture<T>>,
    stepped: Rc<Cell<bool>>,
    progress: Rc<Cell<InvocationProgress>>,
}

impl<T: FromValue + 'static> Invocation<T> {
    pub fn start(mut vm: Vm, name: &str, args: impl Args) -> Result<Self, VmError> {
        let execution = vm.execute([name], args)?;
        let entrypoint = execution.vm().ip();
        let mut execution = execution.into_owned();
        // Rune 0.14.2's into_owned transfers the stack but resets the VM IP.
        execution.vm_mut().set_ip(entrypoint);
        let stepped = Rc::new(Cell::new(false));
        let marker = Rc::clone(&stepped);
        let progress = Rc::new(Cell::new(InvocationProgress::default()));
        let inspector = Rc::clone(&progress);
        let future = Box::pin(async move {
            loop {
                let result = execution.async_step().await;
                let vm = execution.vm();
                let previous = inspector.get();
                inspector.set(InvocationProgress {
                    vm_steps: previous.vm_steps.saturating_add(1),
                    instruction_pointer: vm.ip(),
                    call_depth: vm.call_frames().len(),
                });
                if let Some(value) = result.into_result()? {
                    return Ok(rune::from_value(value)?);
                }
                marker.set(true);
                let mut yielded = false;
                poll_fn(|cx| {
                    if yielded {
                        Poll::Ready(())
                    } else {
                        yielded = true;
                        cx.waker().wake_by_ref();
                        Poll::Pending
                    }
                })
                .await;
            }
        });
        Ok(Self {
            future: Some(future),
            stepped,
            progress,
        })
    }

    /// Get instruction telemetry from the most recently completed VM step.
    pub fn progress(&self) -> InvocationProgress {
        self.progress.get()
    }

    /// Bound VM steps per call, preserving the exact pending async step.
    pub fn poll_slice(&mut self, cx: &mut Context<'_>, budget: NonZeroUsize) -> Slice<T> {
        let Some(future) = self.future.as_mut() else {
            return Slice::Failed("invocation has already finished".into());
        };
        for _ in 0..budget.get() {
            self.stepped.set(false);
            match future.as_mut().poll(cx) {
                Poll::Ready(result) => {
                    self.future = None;
                    return match result {
                        Ok(value) => Slice::Completed(value),
                        Err(error) => Slice::Failed(error),
                    };
                }
                Poll::Pending if !self.stepped.get() => return Slice::Awaiting,
                Poll::Pending => {}
            }
        }
        Slice::BudgetExhausted
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::CompiledProgram;
    use rune::Module;
    use std::sync::{Arc, Mutex};
    use std::task::{Wake, Waker};

    use std::sync::atomic::{AtomicBool, Ordering};

    #[test]
    fn shutdown_invocation_preserves_state_argument() -> ExperimentResult<()> {
        let program = CompiledProgram::compile(
            "pub fn handle(ctx, state, sender, message) { let n = 0; while n < 20 { n += 1; } state + message } pub fn on_shutdown(state) { state + 100 }",
        )?;
        assert_eq!(
            rune::from_value::<i64>(program.vm().call(["on_shutdown"], (1i64,))?)?,
            101
        );
        let mut invocation = Invocation::<i64>::start(program.vm(), "on_shutdown", (1i64,))?;
        let mut cx = Context::from_waker(Waker::noop());
        loop {
            match invocation.poll_slice(&mut cx, NonZeroUsize::new(32).unwrap()) {
                Slice::Completed(value) => {
                    assert_eq!(value, 101);
                    break;
                }
                Slice::Failed(error) => return Err(error),
                _ => {}
            }
        }
        Ok(())
    }

    struct Signal {
        waker: Mutex<Option<Waker>>,
        ready: AtomicBool,
        dropped: AtomicBool,
    }
    struct PendingOperation(Arc<Signal>);

    impl Future for PendingOperation {
        type Output = i64;

        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<i64> {
            if self.0.ready.load(Ordering::SeqCst) {
                Poll::Ready(7)
            } else {
                *self.0.waker.lock().unwrap() = Some(cx.waker().clone());
                Poll::Pending
            }
        }
    }

    impl Drop for PendingOperation {
        fn drop(&mut self) {
            self.0.dropped.store(true, Ordering::SeqCst);
        }
    }

    struct WakeCounter(std::sync::atomic::AtomicUsize);
    impl Wake for WakeCounter {
        fn wake(self: Arc<Self>) {
            self.0.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
        }
    }

    #[test]
    fn budgets_await_wake_resume_and_cancel() -> ExperimentResult<()> {
        let signal = Arc::new(Signal {
            waker: Mutex::new(None),
            ready: AtomicBool::new(false),
            dropped: AtomicBool::new(false),
        });
        let mut module = Module::new();
        let native_signal = Arc::clone(&signal);
        module
            .function("wait", move || PendingOperation(Arc::clone(&native_signal)))
            .build()?;
        let program = CompiledProgram::compile_with_module(
            "pub async fn handle() { let n = 0; while n < 10 { n += 1; } let x = wait().await; while n < 20 { n += 1; } n + x }",
            module,
        )?;
        let counter = Arc::new(WakeCounter(std::sync::atomic::AtomicUsize::new(0)));
        let waker = Waker::from(Arc::clone(&counter));
        let mut cx = Context::from_waker(&waker);
        let budget = NonZeroUsize::new(1).unwrap();
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        assert!(matches!(
            invocation.poll_slice(&mut cx, budget),
            Slice::BudgetExhausted
        ));
        for _ in 0..1_000 {
            if matches!(invocation.poll_slice(&mut cx, budget), Slice::Awaiting) {
                break;
            }
        }
        let registered = signal
            .waker
            .lock()
            .unwrap()
            .take()
            .expect("await registered waker");
        assert!(matches!(
            invocation.poll_slice(&mut cx, budget),
            Slice::Awaiting
        ));
        signal.ready.store(true, Ordering::SeqCst);
        registered.wake();
        assert!(counter.0.load(std::sync::atomic::Ordering::Relaxed) > 0);
        let mut completed = false;
        let mut exhausted_after_await = false;
        for _ in 0..1_000 {
            match invocation.poll_slice(&mut cx, budget) {
                Slice::Completed(value) => {
                    assert_eq!(value, 27);
                    completed = true;
                    break;
                }
                Slice::Failed(error) => return Err(error),
                Slice::BudgetExhausted => exhausted_after_await = true,
                Slice::Awaiting => return Err("ready operation remained pending".into()),
            }
        }
        assert!(completed);
        assert!(exhausted_after_await);

        signal.ready.store(false, Ordering::SeqCst);
        signal.dropped.store(false, Ordering::SeqCst);
        let mut cancelled = Invocation::<i64>::start(program.vm(), "handle", ())?;
        for _ in 0..1_000 {
            if matches!(cancelled.poll_slice(&mut cx, budget), Slice::Awaiting) {
                break;
            }
        }
        assert!(!signal.dropped.load(Ordering::SeqCst));
        drop(cancelled);
        assert!(signal.dropped.load(Ordering::SeqCst));
        Ok(())
    }

    #[test]
    fn errors_and_finished_invocations_have_typed_outcomes() -> ExperimentResult<()> {
        let program = CompiledProgram::compile("pub fn handle() { 1 / 0 }")?;
        let mut invocation = Invocation::<i64>::start(program.vm(), "handle", ())?;
        let mut cx = Context::from_waker(Waker::noop());
        assert!(matches!(
            invocation.poll_slice(&mut cx, NonZeroUsize::new(100).unwrap()),
            Slice::Failed(_)
        ));
        assert!(matches!(
            invocation.poll_slice(&mut cx, NonZeroUsize::new(100).unwrap()),
            Slice::Failed(_)
        ));
        assert!(Invocation::<i64>::start(program.vm(), "missing", ()).is_err());
        Ok(())
    }
}