Thanks to visit codestin.com
Credit goes to docs.rs

rustyscript/
worker.rs

1//! Provides a worker thread that can be used to run javascript code in a separate thread through a channel pair
2//! It also provides a default worker implementation that can be used without any additional setup:
3//! ```rust
4//! use rustyscript::{Error, worker::{Worker, DefaultWorker, DefaultWorkerOptions}};
5//! use std::time::Duration;
6//!
7//! fn main() -> Result<(), Error> {
8//!     let worker = DefaultWorker::new(DefaultWorkerOptions {
9//!         default_entrypoint: None,
10//!         timeout: Duration::from_secs(5),
11//!         ..Default::default()
12//!     })?;
13//!
14//!     let result: i32 = worker.eval("5 + 5".to_string())?;
15//!     assert_eq!(result, 10);
16//!     Ok(())
17//! }
18use std::{
19    cell::RefCell,
20    rc::Rc,
21    sync::mpsc::{channel, Receiver, Sender},
22    thread::{spawn, JoinHandle},
23};
24
25use crate::{Error, RuntimeOptions};
26
27/// A pool of worker threads that can be used to run javascript code in parallel
28/// Uses a round-robin strategy to distribute work between workers
29/// Each worker is an independent runtime instance
30pub struct WorkerPool<W>
31where
32    W: InnerWorker,
33{
34    workers: Vec<Rc<RefCell<Worker<W>>>>,
35    next_worker: usize,
36    options: W::RuntimeOptions,
37}
38
39impl<W> WorkerPool<W>
40where
41    W: InnerWorker,
42{
43    /// Create a new worker pool with the specified number of workers
44    ///
45    /// # Errors
46    /// Can fail if a runtime cannot be initialized (usually due to extension issues)
47    pub fn new(options: W::RuntimeOptions, n_workers: u32) -> Result<Self, Error> {
48        crate::init_platform(n_workers, true);
49        let mut workers = Vec::with_capacity(n_workers as usize + 1);
50        for _ in 0..n_workers {
51            workers.push(Rc::new(RefCell::new(Worker::new(options.clone())?)));
52        }
53
54        Ok(Self {
55            workers,
56            next_worker: 0,
57            options,
58        })
59    }
60
61    /// Returns the runtime options used by the workers in the pool
62    #[must_use]
63    pub fn options(&self) -> &W::RuntimeOptions {
64        &self.options
65    }
66
67    /// Stop all workers in the pool and wait for them to finish
68    pub fn shutdown(self) {
69        for worker in self.workers {
70            worker.borrow_mut().shutdown();
71        }
72    }
73
74    /// Get the number of workers in the pool
75    #[must_use]
76    pub fn len(&self) -> usize {
77        self.workers.len()
78    }
79
80    /// Check if the pool is empty
81    /// This will be true if the pool has no workers
82    /// This can happen if the pool was created with 0 workers
83    /// Which is not particularly useful, but is allowed
84    #[must_use]
85    pub fn is_empty(&self) -> bool {
86        self.workers.is_empty()
87    }
88
89    /// Get a worker by its index in the pool
90    #[must_use]
91    pub fn worker_by_id(&self, id: usize) -> Option<Rc<RefCell<Worker<W>>>> {
92        Some(Rc::clone(self.workers.get(id)?))
93    }
94
95    /// Get the next worker in the pool
96    pub fn next_worker(&mut self) -> Rc<RefCell<Worker<W>>> {
97        let worker = &self.workers[self.next_worker];
98        self.next_worker = (self.next_worker + 1) % self.workers.len();
99        Rc::clone(worker)
100    }
101
102    /// Send a request to the next worker in the pool
103    /// This will block the current thread until the response is received
104    ///
105    /// # Errors
106    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
107    pub fn send_and_await(&mut self, query: W::Query) -> Result<W::Response, Error> {
108        self.next_worker().borrow().send_and_await(query)
109    }
110
111    /// Evaluate a string of non-ecma javascript code in a separate thread
112    /// The code is evaluated in a new runtime instance, which is then destroyed
113    /// Returns a handle to the thread that is running the code
114    #[must_use = "The returned thread handle will return a Result<T, Error> when joined"]
115    pub fn eval_in_thread<T>(code: String) -> std::thread::JoinHandle<Result<T, Error>>
116    where
117        T: serde::de::DeserializeOwned + Send + 'static,
118    {
119        deno_core::JsRuntime::init_platform(None, true);
120        std::thread::spawn(move || {
121            let mut runtime = crate::Runtime::new(RuntimeOptions::default())?;
122            runtime.eval(&code)
123        })
124    }
125}
126
127/// A worker thread that can be used to run javascript code in a separate thread
128/// Contains a channel pair for communication, and a single runtime instance
129///
130/// This worker is generic over an implementation of the [`InnerWorker`] trait
131/// This allows flexibility in the runtime used by the worker, as well as the types of queries and responses that can be used
132///
133/// For a simple worker that uses the default runtime, see [`DefaultWorker`]
134pub struct Worker<W>
135where
136    W: InnerWorker,
137{
138    handle: Option<JoinHandle<()>>,
139    tx: Option<Sender<W::Query>>,
140    rx: Receiver<W::Response>,
141}
142
143impl<W> Worker<W>
144where
145    W: InnerWorker,
146{
147    /// Create a new worker instance
148    ///
149    /// # Errors
150    /// Can fail if the runtime cannot be initialized (usually due to extension issues)
151    pub fn new(options: W::RuntimeOptions) -> Result<Self, Error> {
152        let (qtx, qrx) = channel();
153        let (rtx, rrx) = channel();
154        let (init_tx, init_rx) = channel::<Option<Error>>();
155
156        let handle = spawn(move || {
157            let rx = qrx;
158            let tx = rtx;
159            let itx = init_tx;
160
161            let runtime = match W::init_runtime(options) {
162                Ok(rt) => rt,
163                Err(e) => {
164                    itx.send(Some(e)).ok(); // Stopping anyway, so no need to check for errors
165                    return;
166                }
167            };
168
169            if itx.send(None).is_ok() {
170                W::thread(runtime, rx, tx);
171            }
172        });
173
174        let worker = Self {
175            handle: Some(handle),
176            tx: Some(qtx),
177            rx: rrx,
178        };
179
180        // Wait for initialization to complete
181        match init_rx.recv() {
182            Ok(None) => Ok(worker),
183
184            // Initialization failed
185            Ok(Some(e)) => Err(e),
186
187            // Parser crashed on startup
188            _ => {
189                let Some(handle) = worker.handle else {
190                    return Err(Error::Runtime(
191                        "Could not start runtime thread: Worker handle missing".to_string(),
192                    ));
193                };
194
195                // Attempt to join the thread to get the error message
196                let Err(e) = handle.join() else {
197                    return Err(Error::Runtime("Could not start runtime thread".to_string()));
198                };
199
200                // Get the actual error message - String, &str, or default message
201                let e = if let Some(e) = e.downcast_ref::<String>() {
202                    e.clone()
203                } else if let Some(e) = e.downcast_ref::<&str>() {
204                    (*e).to_string()
205                } else {
206                    "Could not start runtime thread".to_string()
207                };
208
209                // Remove everything after the words 'Stack backtrace'
210                let e = match e.split("Stack backtrace").next() {
211                    Some(e) => e.trim(),
212                    None => &e,
213                }
214                .to_string();
215
216                Err(Error::Runtime(e))
217            }
218        }
219    }
220
221    /// Stop the worker and wait for it to finish
222    /// Stops by destroying the sender, which will cause the thread to exit the loop and finish
223    ///
224    /// WARNING: If implementing a custom `thread` function, make sure to handle rx failures gracefully
225    ///          Otherwise this will block indefinitely
226    pub fn shutdown(&mut self) {
227        if let (Some(tx), Some(hnd)) = (self.tx.take(), self.handle.take()) {
228            // We can stop the thread by destroying the sender
229            // This will cause the thread to exit the loop and finish
230            drop(tx);
231            hnd.join().ok();
232        }
233    }
234
235    /// Send a request to the worker
236    /// This will not block the current thread
237    ///
238    /// # Errors
239    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
240    pub fn send(&self, query: W::Query) -> Result<(), Error> {
241        match &self.tx {
242            None => return Err(Error::WorkerHasStopped),
243            Some(tx) => tx,
244        }
245        .send(query)
246        .map_err(|e| Error::Runtime(e.to_string()))
247    }
248
249    /// Receive a response from the worker
250    /// This will block the current thread until a response is received
251    ///
252    /// # Errors
253    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
254    pub fn receive(&self) -> Result<W::Response, Error> {
255        self.rx.recv().map_err(|e| Error::Runtime(e.to_string()))
256    }
257
258    /// Try to receive a response from the worker without blocking
259    /// This will return `Ok(None)` if no response is available
260    ///
261    /// # Errors
262    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
263    pub fn try_receive(&self) -> Result<Option<W::Response>, Error> {
264        match self.rx.try_recv() {
265            Ok(v) => Ok(Some(v)),
266            Err(e) => match e {
267                std::sync::mpsc::TryRecvError::Empty => Ok(None),
268                std::sync::mpsc::TryRecvError::Disconnected => Err(Error::Runtime(e.to_string())),
269            },
270        }
271    }
272
273    /// Send a request to the worker and wait for a response
274    /// This will block the current thread until a response is received
275    /// Will return an error if the worker has stopped or panicked
276    ///
277    /// # Errors
278    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
279    pub fn send_and_await(&self, query: W::Query) -> Result<W::Response, Error> {
280        self.send(query)?;
281        self.receive()
282    }
283
284    /// Consume the worker and wait for the thread to finish
285    ///
286    /// WARNING: If implementing a custom `thread` function, make sure to handle rx failures gracefully
287    ///          Otherwise this will block indefinitely
288    ///
289    /// # Errors
290    /// Will return an error if the worker has already been stopped, or if the worker thread panicked
291    pub fn join(mut self) -> Result<(), Error> {
292        self.shutdown();
293        match self.handle {
294            Some(hnd) => hnd
295                .join()
296                .map_err(|_| Error::Runtime("Worker thread panicked".to_string())),
297            None => Ok(()),
298        }
299    }
300}
301
302/// An implementation of the worker trait for a specific runtime
303/// This allows flexibility in the runtime used by the worker
304/// As well as the types of queries and responses that can be used
305///
306/// Implement this trait for a specific runtime to use it with the worker
307/// For an example implementation, see [`DefaultWorker`]
308pub trait InnerWorker
309where
310    Self: Send,
311    <Self as InnerWorker>::RuntimeOptions: std::marker::Send + 'static + Clone,
312    <Self as InnerWorker>::Query: std::marker::Send + 'static,
313    <Self as InnerWorker>::Response: std::marker::Send + 'static,
314{
315    /// The type of runtime used by this worker
316    /// This can just be `rustyscript::Runtime` if you don't need to use a custom runtime
317    type Runtime;
318
319    /// The type of options that can be used to initialize the runtime
320    /// Cannot be `rustyscript::RuntimeOptions` because it is not `Send`
321    type RuntimeOptions;
322
323    /// The type of query that can be sent to the worker
324    /// This should be an enum that contains all possible queries
325    type Query;
326
327    /// The type of response that can be received from the worker
328    /// This should be an enum that contains all possible responses
329    type Response;
330
331    /// Initialize the runtime used by the worker
332    /// This should return a new instance of the runtime that will respond to queries
333    ///
334    /// # Errors
335    /// Can fail if the runtime cannot be initialized (usually due to extension issues)
336    fn init_runtime(options: Self::RuntimeOptions) -> Result<Self::Runtime, Error>;
337
338    /// Handle a query sent to the worker
339    /// Must always return a response of some kind
340    fn handle_query(runtime: &mut Self::Runtime, query: Self::Query) -> Self::Response;
341
342    /// The main thread function that will be run by the worker
343    /// This should handle all incoming queries and send responses back
344    fn thread(mut runtime: Self::Runtime, rx: Receiver<Self::Query>, tx: Sender<Self::Response>) {
345        loop {
346            let Ok(msg) = rx.recv() else {
347                break;
348            };
349
350            let response = Self::handle_query(&mut runtime, msg);
351            if tx.send(response).is_err() {
352                break;
353            }
354        }
355    }
356}
357
358/// A worker implementation that uses the default runtime
359/// This is the simplest way to use the worker, as it requires no additional setup
360/// It attempts to provide as much functionality as possible from the standard runtime
361///
362/// Please note that it uses `serde_json::Value` for queries and responses, which comes with a performance cost
363/// For a more performant worker, or to use extensions and/or loader caches, you'll need to implement your own worker
364pub struct DefaultWorker(Worker<DefaultWorker>);
365impl InnerWorker for DefaultWorker {
366    type Runtime = (
367        crate::Runtime,
368        std::collections::HashMap<deno_core::ModuleId, crate::ModuleHandle>,
369    );
370    type RuntimeOptions = DefaultWorkerOptions;
371    type Query = DefaultWorkerQuery;
372    type Response = DefaultWorkerResponse;
373
374    fn init_runtime(options: Self::RuntimeOptions) -> Result<Self::Runtime, Error> {
375        let runtime = crate::Runtime::new(crate::RuntimeOptions {
376            default_entrypoint: options.default_entrypoint,
377            timeout: options.timeout,
378            shared_array_buffer_store: options.shared_array_buffer_store,
379            startup_snapshot: options.startup_snapshot,
380            ..Default::default()
381        })?;
382        let modules = std::collections::HashMap::new();
383        Ok((runtime, modules))
384    }
385
386    fn handle_query(runtime: &mut Self::Runtime, query: Self::Query) -> Self::Response {
387        let (runtime, modules) = runtime;
388        match query {
389            DefaultWorkerQuery::Eval(code) => match runtime.eval(&code) {
390                Ok(v) => Self::Response::Value(v),
391                Err(e) => Self::Response::Error(e),
392            },
393
394            DefaultWorkerQuery::LoadMainModule(module) => {
395                match runtime.load_modules(&module, vec![]) {
396                    Ok(handle) => {
397                        let id = handle.id();
398                        modules.insert(id, handle);
399                        Self::Response::ModuleId(id)
400                    }
401                    Err(e) => Self::Response::Error(e),
402                }
403            }
404
405            DefaultWorkerQuery::LoadModule(module) => match runtime.load_module(&module) {
406                Ok(handle) => {
407                    let id = handle.id();
408                    modules.insert(id, handle);
409                    Self::Response::ModuleId(id)
410                }
411                Err(e) => Self::Response::Error(e),
412            },
413
414            DefaultWorkerQuery::CallEntrypoint(id, args) => match modules.get(&id) {
415                Some(handle) => match runtime.call_entrypoint(handle, &args) {
416                    Ok(v) => Self::Response::Value(v),
417                    Err(e) => Self::Response::Error(e),
418                },
419                None => Self::Response::Error(Error::Runtime("Module not found".to_string())),
420            },
421
422            DefaultWorkerQuery::CallFunction(id, name, args) => {
423                let handle = if let Some(id) = id {
424                    match modules.get(&id) {
425                        Some(handle) => Some(handle),
426                        None => {
427                            return Self::Response::Error(Error::Runtime(
428                                "Module not found".to_string(),
429                            ))
430                        }
431                    }
432                } else {
433                    None
434                };
435
436                match runtime.call_function(handle, &name, &args) {
437                    Ok(v) => Self::Response::Value(v),
438                    Err(e) => Self::Response::Error(e),
439                }
440            }
441
442            DefaultWorkerQuery::GetValue(id, name) => {
443                let handle = if let Some(id) = id {
444                    match modules.get(&id) {
445                        Some(handle) => Some(handle),
446                        None => {
447                            return Self::Response::Error(Error::Runtime(
448                                "Module not found".to_string(),
449                            ))
450                        }
451                    }
452                } else {
453                    None
454                };
455
456                match runtime.get_value(handle, &name) {
457                    Ok(v) => Self::Response::Value(v),
458                    Err(e) => Self::Response::Error(e),
459                }
460            }
461        }
462    }
463}
464impl DefaultWorker {
465    /// Create a new worker instance
466    ///
467    /// # Errors
468    /// Can fail if the runtime cannot be initialized (usually due to extension issues)
469    pub fn new(options: DefaultWorkerOptions) -> Result<Self, Error> {
470        Worker::new(options).map(Self)
471    }
472
473    /// Get a reference to the underlying worker instance
474    #[must_use]
475    pub fn as_worker(&self) -> &Worker<DefaultWorker> {
476        &self.0
477    }
478
479    /// Evaluate a string of javascript code
480    /// Returns the result of the evaluation
481    ///
482    /// # Errors
483    /// Can fail a runtime error occurs during evaluation, or if the return value cannot be deserialized into the requested type
484    pub fn eval<T>(&self, code: String) -> Result<T, Error>
485    where
486        T: serde::de::DeserializeOwned,
487    {
488        match self.0.send_and_await(DefaultWorkerQuery::Eval(code))? {
489            DefaultWorkerResponse::Value(v) => Ok(crate::serde_json::from_value(v)?),
490            DefaultWorkerResponse::Error(e) => Err(e),
491            _ => Err(Error::Runtime(
492                "Unexpected response from the worker".to_string(),
493            )),
494        }
495    }
496
497    /// Load a module into the worker as the main module
498    /// Returns the module id of the loaded module
499    ///
500    /// # Errors
501    /// Can fail if execution of the module fails
502    pub fn load_main_module(&self, module: crate::Module) -> Result<deno_core::ModuleId, Error> {
503        match self
504            .0
505            .send_and_await(DefaultWorkerQuery::LoadMainModule(module))?
506        {
507            DefaultWorkerResponse::ModuleId(id) => Ok(id),
508            DefaultWorkerResponse::Error(e) => Err(e),
509            _ => Err(Error::Runtime(
510                "Unexpected response from the worker".to_string(),
511            )),
512        }
513    }
514
515    /// Load a module into the worker as a side module
516    /// Returns the module id of the loaded module
517    ///
518    /// # Errors
519    /// Can fail if execution of the module fails
520    pub fn load_module(&self, module: crate::Module) -> Result<deno_core::ModuleId, Error> {
521        match self
522            .0
523            .send_and_await(DefaultWorkerQuery::LoadModule(module))?
524        {
525            DefaultWorkerResponse::ModuleId(id) => Ok(id),
526            DefaultWorkerResponse::Error(e) => Err(e),
527            _ => Err(Error::Runtime(
528                "Unexpected response from the worker".to_string(),
529            )),
530        }
531    }
532
533    /// Call the entrypoint function in a module
534    /// Returns the result of the function call
535    /// The module id must be the id of a module loaded with `load_main_module` or `load_module`
536    ///
537    /// # Errors
538    /// Can fail the module is not found, if there is no entrypoint function, if the entrypoint function returns an error,
539    /// Or if the return value cannot be deserialized into the requested type
540    pub fn call_entrypoint<T>(
541        &self,
542        id: deno_core::ModuleId,
543        args: Vec<crate::serde_json::Value>,
544    ) -> Result<T, Error>
545    where
546        T: serde::de::DeserializeOwned,
547    {
548        match self
549            .0
550            .send_and_await(DefaultWorkerQuery::CallEntrypoint(id, args))?
551        {
552            DefaultWorkerResponse::Value(v) => {
553                crate::serde_json::from_value(v).map_err(Error::from)
554            }
555            DefaultWorkerResponse::Error(e) => Err(e),
556            _ => Err(Error::Runtime(
557                "Unexpected response from the worker".to_string(),
558            )),
559        }
560    }
561
562    /// Call a function in a module
563    /// Returns the result of the function call
564    /// The module id must be the id of a module loaded with `load_main_module` or `load_module`
565    ///
566    /// # Errors
567    /// Can fail if the function is not found, if the function returns an error,
568    /// Or if the return value cannot be deserialized into the requested type
569    pub fn call_function<T>(
570        &self,
571        module_context: Option<deno_core::ModuleId>,
572        name: String,
573        args: Vec<crate::serde_json::Value>,
574    ) -> Result<T, Error>
575    where
576        T: serde::de::DeserializeOwned,
577    {
578        match self
579            .0
580            .send_and_await(DefaultWorkerQuery::CallFunction(module_context, name, args))?
581        {
582            DefaultWorkerResponse::Value(v) => {
583                crate::serde_json::from_value(v).map_err(Error::from)
584            }
585            DefaultWorkerResponse::Error(e) => Err(e),
586            _ => Err(Error::Runtime(
587                "Unexpected response from the worker".to_string(),
588            )),
589        }
590    }
591
592    /// Get a value from a module
593    /// The module id must be the id of a module loaded with `load_main_module` or `load_module`
594    ///
595    /// # Errors
596    /// Can fail if the value is not found or if the value cannot be deserialized into the requested type
597    pub fn get_value<T>(
598        &self,
599        module_context: Option<deno_core::ModuleId>,
600        name: String,
601    ) -> Result<T, Error>
602    where
603        T: serde::de::DeserializeOwned,
604    {
605        match self
606            .0
607            .send_and_await(DefaultWorkerQuery::GetValue(module_context, name))?
608        {
609            DefaultWorkerResponse::Value(v) => {
610                crate::serde_json::from_value(v).map_err(Error::from)
611            }
612            DefaultWorkerResponse::Error(e) => Err(e),
613            _ => Err(Error::Runtime(
614                "Unexpected response from the worker".to_string(),
615            )),
616        }
617    }
618}
619impl AsRef<Worker<DefaultWorker>> for DefaultWorker {
620    fn as_ref(&self) -> &Worker<DefaultWorker> {
621        &self.0
622    }
623}
624
625/// Options for the default worker
626#[derive(Default, Clone)]
627pub struct DefaultWorkerOptions {
628    /// The default entrypoint function to use if none is registered
629    pub default_entrypoint: Option<String>,
630
631    /// The timeout to use for the runtime
632    pub timeout: std::time::Duration,
633
634    /// Optional snapshot to load into the runtime
635    /// This will reduce load times, but requires the same extensions to be loaded
636    /// as when the snapshot was created
637    pub startup_snapshot: Option<&'static [u8]>,
638
639    /// Optional shared array buffer store to use for the runtime
640    /// Allows data-sharing between runtimes across threads
641    pub shared_array_buffer_store: Option<deno_core::SharedArrayBufferStore>,
642}
643
644/// Query types for the default worker
645#[derive(Debug, Clone)]
646pub enum DefaultWorkerQuery {
647    /// Evaluates a string of javascript code
648    Eval(String),
649
650    /// Loads a module into the worker as the main module
651    LoadMainModule(crate::Module),
652
653    /// Loads a module into the worker as a side module
654    LoadModule(crate::Module),
655
656    /// Calls an entrypoint function in a module
657    CallEntrypoint(deno_core::ModuleId, Vec<crate::serde_json::Value>),
658
659    /// Calls a function in a module
660    CallFunction(
661        Option<deno_core::ModuleId>,
662        String,
663        Vec<crate::serde_json::Value>,
664    ),
665
666    /// Gets a value from a module
667    GetValue(Option<deno_core::ModuleId>, String),
668}
669
670/// Response types for the default worker
671#[derive(Debug, Clone)]
672pub enum DefaultWorkerResponse {
673    /// A successful response with a value
674    Value(crate::serde_json::Value),
675
676    /// A successful response with a module id
677    ModuleId(deno_core::ModuleId),
678
679    /// A successful response with no value
680    Ok(()),
681
682    /// An error response
683    Error(Error),
684}