Switch to a single threaded tokio executor, remove some unsafe

Cargo.toml

@@ -49,7 +49,7 @@ status-line = "0.2.0"
 strum = { version = "0.26", features = ["derive"] }
 time = "0.3"
 thiserror = "1.0.26"
-tokio = { version = "1", features = ["rt", "rt-multi-thread", "time", "parking_lot", "signal"] }
+tokio = { version = "1", features = ["rt", "time", "parking_lot", "signal"] }
 tokio-stream = "0.1"
 tracing = "0.1"
 tracing-appender = "0.2"

src/exec/mod.rs

@@ -12,7 +12,9 @@ use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
 use std::time::Instant;
+use tokio::runtime::Builder;
 use tokio::signal::ctrl_c;
+use tokio::task::LocalSet;
 use tokio::time::MissedTickBehavior;
 use tokio_stream::wrappers::IntervalStream;
 
@@ -107,35 +109,40 @@ fn spawn_stream(
 ) -> Receiver<Result<WorkloadStats>> {
     let (tx, rx) = channel(1);
 
-    tokio::spawn(async move {
-        match rate {
-            Some(rate) => {
-                let stream = interval_stream(rate);
-                run_stream(
-                    stream,
-                    workload,
-                    iter_counter,
-                    concurrency,
-                    sampling,
-                    progress,
-                    tx,
-                )
-                .await
+    let rt = Builder::new_current_thread().enable_all().build().unwrap();
+    std::thread::spawn(move || {
+        let local = LocalSet::new();
+        local.spawn_local(async move {
+            match rate {
+                Some(rate) => {
+                    let stream = interval_stream(rate);
+                    run_stream(
+                        stream,
+                        workload,
+                        iter_counter,
+                        concurrency,
+                        sampling,
+                        progress,
+                        tx,
+                    )
+                    .await
+                }
+                None => {
+                    let stream = futures::stream::repeat_with(|| ());
+                    run_stream(
+                        stream,
+                        workload,
+                        iter_counter,
+                        concurrency,
+                        sampling,
+                        progress,
+                        tx,
+                    )
+                    .await
+                }
             }
-            None => {
-                let stream = futures::stream::repeat_with(|| ());
-                run_stream(
-                    stream,
-                    workload,
-                    iter_counter,
-                    concurrency,
-                    sampling,
-                    progress,
-                    tx,
-                )
-                .await
-            }
-        }
+        });
+        rt.block_on(local);
     });
     rx
 }

src/exec/workload.rs

@@ -177,17 +177,14 @@ impl Program {
     /// If execution fails, emits diagnostic messages, e.g. stacktrace to standard error stream.
     /// Also signals an error if the function execution succeeds, but the function returns
     /// an error value.
-    pub async fn async_call(
-        &self,
-        fun: &FnRef,
-        args: impl Args + Send,
-    ) -> Result<Value, LatteError> {
+    pub async fn async_call(&self, fun: &FnRef, args: impl Args) -> Result<Value, LatteError> {
         let handle_err = |e: VmError| {
             let mut out = StandardStream::stderr(ColorChoice::Auto);
             let _ = e.emit(&mut out, &self.sources);
             LatteError::ScriptExecError(fun.name.to_string(), e)
         };
-        let execution = self.vm().send_execute(fun.hash, args).map_err(handle_err)?;
+        let mut vm = self.vm();
+        let mut execution = vm.execute(fun.hash, args).map_err(handle_err)?;
         let result = execution
             .async_complete()
             .await

src/main.rs

@@ -507,15 +507,8 @@ fn edit_workload(workload: PathBuf) -> Result<()> {
     Ok(())
 }
 
-fn init_runtime(thread_count: usize) -> std::io::Result<Runtime> {
-    if thread_count == 1 {
-        Builder::new_current_thread().enable_all().build()
-    } else {
-        Builder::new_multi_thread()
-            .worker_threads(thread_count)
-            .enable_all()
-            .build()
-    }
+fn init_runtime() -> std::io::Result<Runtime> {
+    Builder::new_current_thread().enable_all().build()
 }
 
 fn setup_logging(run_id: &str, config: &AppConfig) -> Result<WorkerGuard> {
@@ -560,12 +553,7 @@ fn main() {
     };
 
     let command = config.command;
-    let thread_count = match &command {
-        Command::Run(cmd) => cmd.threads.get(),
-        Command::Load(cmd) => cmd.threads.get(),
-        _ => 1,
-    };
-    let runtime = init_runtime(thread_count);
+    let runtime = init_runtime();
     if let Err(e) = runtime.unwrap().block_on(async_main(run_id, command)) {
         eprintln!("error: {e}");
         exit(128);

src/scripting/context.rs

@@ -54,13 +54,9 @@ pub struct GlobalContext {
 // Needed, because Rune `Value` is !Send, as it may contain some internal pointers.
 // Therefore, it is not safe to pass a `Value` to another thread by cloning it, because
 // both objects could accidentally share some unprotected, `!Sync` data.
-// To make it safe, the same `Context` is never used by more than one thread at once, and
-// we make sure in `clone` to make a deep copy of the `data` field by serializing
+// To make it safe, we make sure in `clone` to make a deep copy of the `data` field by serializing
 // and deserializing it, so no pointers could get through.
-unsafe impl Send for LocalContext {}
-unsafe impl Sync for LocalContext {}
 unsafe impl Send for GlobalContext {}
-unsafe impl Sync for GlobalContext {}
 
 impl GlobalContext {
     pub fn new(session: scylla::Session, retry_strategy: RetryStrategy) -> Self {