wasi_worker/
service.rs

1use super::{FileOptions, ServiceOptions};
2use std::cell::RefCell;
3use std::fs::File;
4use std::io::{self, Read, Write};
5
6/// Connects Rust Handler with browser service worker via WASI filesystem.
7///
8/// ServiceWorker is a singleton which holds input and output file handles and
9/// owns woker via Handler trait. Worker is supposedly reactive, usually operating
10/// on incoming events (on_message) and posting messages to main browser application
11/// via ServiceWorker::post_message().
12///
13/// Note: ServiceWorker supposed to operate in single threaded environment
14/// like a browser service worker.
15///
16/// TODO: it requires cleaning of filesystem, add drop implementation
17pub struct ServiceWorker {
18    output: File,
19    input: io::Stdin,
20    options: ServiceOptions,
21}
22
23/// Handler for incoming messages via ServiceWorker
24pub trait Handler {
25    fn on_message(&self, msg: &[u8]) -> std::io::Result<()>;
26}
27
28thread_local! {
29  static SERVICE: RefCell<Option<ServiceWorker>> = RefCell::new(None);
30  static HANDLER: RefCell<Option<Box<dyn Handler>>> = RefCell::new(None);
31}
32
33impl ServiceWorker {
34    /// Initialize ServiceWorker instance.
35    /// ServiceWorker operates as singleton, all struct methods are static.
36    /// Unless initialized all methods will result in error io::ErrorKind::NotConnected.
37    pub fn initialize(options: ServiceOptions) -> io::Result<()> {
38        let output = match &options.output {
39            FileOptions::File(path) => File::create(path)?,
40        };
41        let sw = ServiceWorker {
42            output,
43            input: io::stdin(),
44            options,
45        };
46        SERVICE.with(|service| service.replace(Some(sw)));
47        Ok(())
48    }
49
50    /// Message handler is required to process incoming messages.
51    /// Please note, there is no queue therefore messages received before handler initialized will be lost.
52    pub fn set_message_handler(new_handler: Box<dyn Handler>) {
53        HANDLER.with(|handler| handler.replace(Some(new_handler)));
54    }
55
56    /// This method is a trigger
57    /// This is workaround while we don't have wasi::poll_oneoff,
58    /// ideally we shall just poll and wait for FD_READ event.
59    pub fn on_message() -> io::Result<usize> {
60        let mut buf: [u8; 1000] = [0; 1000];
61        let len = SERVICE.with(|service| {
62            if let Some(sw) = &mut *service.borrow_mut() {
63                sw.input.read(&mut buf)
64            } else {
65                Err(io::Error::new(
66                    io::ErrorKind::ConnectionRefused,
67                    "Cannot borrow service mutably",
68                ))
69            }
70        })?;
71        HANDLER.with(|handler| {
72            if let Some(handler) = &*handler.borrow() {
73                handler.on_message(&buf[0..len])?;
74                Ok(len)
75            } else {
76                Err(io::Error::new(
77                    io::ErrorKind::NotConnected,
78                    "Worker was not initialized",
79                ))
80            }
81        })
82    }
83
84    /// Post message to external consumers
85    ///
86    /// Example usage:
87    /// ```
88    /// use wasi_worker::ServiceWorker;
89    /// ServiceWorker::post_message(b"mymesage");
90    /// ```
91    pub fn post_message(msg: &[u8]) -> std::io::Result<()> {
92        SERVICE.with(|service| {
93            if let Some(sw) = &mut *service.borrow_mut() {
94                sw.output.write_all(msg)
95            } else {
96                Err(io::Error::new(
97                    io::ErrorKind::NotConnected,
98                    "Service was not initialized",
99                ))
100            }
101        })
102    }
103
104    pub fn kill() -> () {
105        SERVICE.with(|service| service.replace(None));
106        HANDLER.with(|handler| handler.replace(None));
107    }
108}
109
110impl Drop for ServiceWorker {
111    fn drop(&mut self) {
112        if self.options.cleanup {
113            let clr = match &self.options.output {
114                FileOptions::File(output) => std::fs::remove_file(output),
115            };
116            match clr {
117                Ok(_) => (),
118                Err(err) => eprintln!("Failed to remove file {}", err),
119            }
120        }
121    }
122}