1use super::{FileOptions, ServiceOptions};
2use std::cell::RefCell;
3use std::fs::File;
4use std::io::{self, Read, Write};
5
6pub struct ServiceWorker {
18 output: File,
19 input: io::Stdin,
20 options: ServiceOptions,
21}
22
23pub 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 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 pub fn set_message_handler(new_handler: Box<dyn Handler>) {
53 HANDLER.with(|handler| handler.replace(Some(new_handler)));
54 }
55
56 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 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}