1use std::collections::HashMap;
4use std::io;
5#[cfg(all(test, unix))]
6use std::io::Read;
7use std::path::{Path, PathBuf};
8use std::process::{Child, Command, Stdio};
9use std::sync::{
10 Arc, Mutex,
11 atomic::{AtomicBool, Ordering},
12 mpsc,
13};
14use std::time::Duration;
15
16use tokio_util::sync::CancellationToken;
17
18use crate::runtime::{BoxFuture, HOOK_NAMES};
19use crate::{
20 Hook, HookContext, HookResult, IsolatedContext, PluginConfig, PluginError, PluginHandle,
21 PluginMetadata, PluginResult, PluginRuntime, SandboxConfig, Value,
22};
23
24mod protocol;
25use protocol::protocol_error as protocol;
26use protocol::{
27 Failure, Reply, Request, Response, WorkerLog, decode, decode_reply, encode, read_frame,
28 write_frame,
29};
30#[cfg(all(test, unix))]
31use protocol::{MAX_LOG_BYTES, WireValue};
32
33const STARTUP_TIMEOUT_MS: u64 = 10_000;
34const MAX_WORKERS: usize = 32;
35static ACTIVE_WORKERS: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
36
37struct WorkerPermit;
38impl WorkerPermit {
39 fn acquire() -> PluginResult<Self> {
40 let mut count = ACTIVE_WORKERS.load(Ordering::Acquire);
43 loop {
44 if count >= MAX_WORKERS {
45 return Err(PluginError::ConfigError {
46 message: "maximum of 32 concurrent plugin workers reached".into(),
47 });
48 }
49 match ACTIVE_WORKERS.compare_exchange_weak(
50 count,
51 count + 1,
52 Ordering::AcqRel,
53 Ordering::Acquire,
54 ) {
55 Ok(_) => return Ok(Self),
56 Err(actual) => count = actual,
57 }
58 }
59 }
60}
61impl Drop for WorkerPermit {
62 fn drop(&mut self) {
63 ACTIVE_WORKERS.fetch_sub(1, Ordering::AcqRel);
64 }
65}
66
67struct Job {
68 bytes: Vec<u8>,
69 reply: mpsc::SyncSender<io::Result<Vec<u8>>>,
70}
71struct Worker {
72 permit: Mutex<Option<WorkerPermit>>,
73 child: Mutex<Child>,
74 requests: mpsc::SyncSender<Job>,
75 busy: AtomicBool,
76 alive: AtomicBool,
77}
78
79fn terminate_child(child: &mut Child) {
80 #[cfg(unix)]
81 {
82 unsafe {
85 libc::kill(-(child.id() as i32), libc::SIGKILL);
86 }
87 }
88 let _ = child.kill();
89 let _ = child.wait();
90}
91
92struct StartingWorker(Option<Child>);
94impl Drop for StartingWorker {
95 fn drop(&mut self) {
96 if let Some(child) = &mut self.0 {
97 terminate_child(child);
98 }
99 }
100}
101
102type IoTask = Box<dyn FnOnce() + Send + 'static>;
103
104impl Worker {
105 fn spawn(executable: &Path, language: &str) -> PluginResult<Arc<Self>> {
106 Self::spawn_with_io_thread(executable, language, |task| {
107 std::thread::Builder::new()
108 .name("gravityfile-plugin-ipc".into())
109 .spawn(task)
110 .map(|_| ())
111 })
112 }
113
114 fn spawn_with_io_thread(
115 executable: &Path,
116 language: &str,
117 start_thread: impl FnOnce(IoTask) -> io::Result<()>,
118 ) -> PluginResult<Arc<Self>> {
119 let mut command = Command::new(executable);
120 command.args(["__plugin-worker", language]);
121 Self::spawn_command(command, start_thread)
122 }
123
124 fn spawn_command(
125 mut command: Command,
126 start_thread: impl FnOnce(IoTask) -> io::Result<()>,
127 ) -> PluginResult<Arc<Self>> {
128 let permit = WorkerPermit::acquire()?;
129 command
130 .env_clear()
131 .current_dir(std::env::temp_dir())
132 .stdin(Stdio::piped())
133 .stdout(Stdio::piped())
134 .stderr(Stdio::null());
135 #[cfg(unix)]
136 {
137 use std::os::unix::process::CommandExt;
138 command.process_group(0);
139 }
140 let mut starting = StartingWorker(Some(command.spawn()?));
141 let child = starting.0.as_mut().expect("newly spawned worker");
142 let mut input = child
143 .stdin
144 .take()
145 .ok_or_else(|| protocol("worker stdin unavailable"))?;
146 let mut output = child
147 .stdout
148 .take()
149 .ok_or_else(|| protocol("worker stdout unavailable"))?;
150 let (requests, jobs) = mpsc::sync_channel::<Job>(1);
151 start_thread(Box::new(move || {
154 while let Ok(job) = jobs.recv() {
155 let result =
156 write_frame(&mut input, &job.bytes).and_then(|_| read_frame(&mut output));
157 let failed = result.is_err();
158 let _ = job.reply.send(result);
159 if failed {
160 break;
161 }
162 }
163 }))?;
164 Ok(Arc::new(Self {
165 permit: Mutex::new(Some(permit)),
166 child: Mutex::new(starting.0.take().expect("worker startup ownership")),
167 requests,
168 busy: AtomicBool::new(false),
169 alive: AtomicBool::new(true),
170 }))
171 }
172
173 fn terminate(&self) {
174 if !self.alive.swap(false, Ordering::AcqRel) {
175 return;
176 }
177 let mut child = self.child.lock().unwrap_or_else(|error| error.into_inner());
178 terminate_child(&mut child);
179 self.permit
180 .lock()
181 .unwrap_or_else(|error| error.into_inner())
182 .take();
183 }
184
185 fn request(self: &Arc<Self>, request: Request, timeout_ms: u64) -> PluginResult<Response> {
186 let _lease = RequestLease::acquire(self, false)?;
187 self.request_reserved(request, timeout_ms)
188 }
189
190 async fn request_async(
191 self: &Arc<Self>,
192 request: Request,
193 timeout_ms: u64,
194 ) -> PluginResult<Response> {
195 self.request_async_with_dispatch(request, timeout_ms, |worker, request, timeout_ms| {
196 worker.request_reserved(request, timeout_ms)
197 })
198 .await
199 }
200
201 async fn request_async_with_dispatch(
202 self: &Arc<Self>,
203 request: Request,
204 timeout_ms: u64,
205 dispatch: impl FnOnce(Arc<Self>, Request, u64) -> PluginResult<Response> + Send + 'static,
206 ) -> PluginResult<Response> {
207 let mut lease = RequestLease::acquire(self, true)?;
208 let worker = Arc::clone(self);
209 let deadline = tokio::time::Instant::now() + Duration::from_millis(timeout_ms);
214 let task =
215 tokio_util::task::AbortOnDropHandle::new(tokio::task::spawn_blocking(move || {
216 dispatch(worker, request, timeout_ms)
217 }));
218 let result = tokio::time::timeout_at(deadline, task)
219 .await
220 .map_err(|_| PluginError::Timeout {
221 name: "worker".into(),
222 timeout_ms,
223 })?
224 .map_err(|error| protocol(error.to_string()))?;
225 lease.terminate_on_drop = false;
226 result
227 }
228
229 fn request_reserved(&self, request: Request, timeout_ms: u64) -> PluginResult<Response> {
230 self.request_reserved_with_encoding(request, timeout_ms, encode)
231 }
232
233 fn request_reserved_with_encoding(
234 &self,
235 request: Request,
236 timeout_ms: u64,
237 encode_request: impl FnOnce(Request) -> PluginResult<Vec<u8>>,
238 ) -> PluginResult<Response> {
239 let started = std::time::Instant::now();
240 let budget = Duration::from_millis(timeout_ms);
241 if !self.alive.load(Ordering::Acquire) {
242 return Err(protocol("plugin worker terminated; reload the plugin"));
243 }
244 let expected = request.expected_response();
245 let bytes = encode_request(request)?;
246 if started.elapsed() >= budget {
249 self.terminate();
250 return Err(PluginError::Timeout {
251 name: "worker".into(),
252 timeout_ms,
253 });
254 }
255 let (reply, result) = mpsc::sync_channel(1);
256 if let Err(error) = self.requests.try_send(Job { bytes, reply }) {
257 self.terminate();
260 return Err(protocol(match error {
261 mpsc::TrySendError::Disconnected(_) => "plugin worker channel closed",
262 mpsc::TrySendError::Full(_) => "plugin worker channel unexpectedly full",
263 }));
264 }
265 match result.recv_timeout(budget.saturating_sub(started.elapsed())) {
266 Ok(Ok(bytes)) => match decode_reply(&bytes) {
267 Ok(reply) => {
268 if let Ok(response) = &reply.result
269 && response.kind() != expected
270 {
271 self.terminate();
272 return Err(protocol(format!(
273 "unexpected plugin response: expected {expected:?}, received {:?}",
274 response.kind()
275 )));
276 }
277 if !reply.log.is_empty() {
278 tracing::info!(target: "plugin_worker", output = %reply.log, truncated = reply.log_truncated, "Plugin log");
279 }
280 match reply.result {
281 Ok(response) => Ok(response),
282 Err(error) => {
283 self.terminate();
284 if let Some(timeout_ms) = error.timeout_ms {
285 Err(PluginError::Timeout {
286 name: "worker".into(),
287 timeout_ms,
288 })
289 } else if error.cancelled {
290 Err(PluginError::Cancelled {
291 name: "worker".into(),
292 })
293 } else {
294 Err(protocol(error.message))
295 }
296 }
297 }
298 }
299 Err(error) => {
300 self.terminate();
301 Err(error)
302 }
303 },
304 Ok(Err(error)) => {
305 self.terminate();
306 Err(error.into())
307 }
308 Err(mpsc::RecvTimeoutError::Timeout) => {
309 self.terminate();
310 Err(PluginError::Timeout {
311 name: "worker".into(),
312 timeout_ms,
313 })
314 }
315 Err(error) => {
316 self.terminate();
317 Err(protocol(error.to_string()))
318 }
319 }
320 }
321}
322struct RequestLease {
325 worker: Arc<Worker>,
326 terminate_on_drop: bool,
327}
328
329impl RequestLease {
330 fn acquire(worker: &Arc<Worker>, terminate_on_drop: bool) -> PluginResult<Self> {
331 worker
332 .busy
333 .compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
334 .map_err(|_| protocol("plugin already has an active request"))?;
335 Ok(Self {
336 worker: Arc::clone(worker),
337 terminate_on_drop,
338 })
339 }
340}
341
342impl Drop for RequestLease {
343 fn drop(&mut self) {
344 if self.terminate_on_drop {
345 self.worker.terminate();
346 }
347 self.worker.busy.store(false, Ordering::Release);
348 }
349}
350
351impl Drop for Worker {
352 fn drop(&mut self) {
353 self.terminate();
354 }
355}
356
357struct Loaded {
358 worker: Arc<Worker>,
359 remote: usize,
360 metadata: PluginMetadata,
361 hooks: Vec<String>,
362}
363
364pub struct ProcessRuntime {
367 executable: PathBuf,
368 language: &'static str,
369 config: PluginConfig,
370 plugins: HashMap<PluginHandle, Loaded>,
371 next: usize,
372}
373impl ProcessRuntime {
374 pub fn new(executable: PathBuf) -> PluginResult<Self> {
377 Ok(Self {
378 executable: dunce::canonicalize(&executable)?,
379 language: "lua",
380 config: PluginConfig::default(),
381 plugins: HashMap::new(),
382 next: 0,
383 })
384 }
385 fn loaded(&self, handle: PluginHandle) -> PluginResult<&Loaded> {
386 self.plugins
387 .get(&handle)
388 .ok_or_else(|| protocol("unknown plugin handle"))
389 }
390}
391
392impl Drop for ProcessRuntime {
393 fn drop(&mut self) {
394 for plugin in self.plugins.values() {
395 plugin.worker.terminate();
396 }
397 }
398}
399
400impl PluginRuntime for ProcessRuntime {
401 fn name(&self) -> &'static str {
402 self.language
403 }
404 fn file_extensions(&self) -> &'static [&'static str] {
405 &[".lua"]
406 }
407 fn init(&mut self, config: &PluginConfig) -> PluginResult<()> {
408 if config.default_timeout_ms == 0 {
409 return Err(PluginError::ConfigError {
410 message: "worker timeout must be nonzero".into(),
411 });
412 }
413 self.config = config.clone();
414 Ok(())
415 }
416 fn load_plugin(&mut self, id: &str, source: &Path) -> PluginResult<PluginHandle> {
417 let source = dunce::canonicalize(source)?;
418 let worker = Worker::spawn(&self.executable, self.language)?;
419 worker.request(Request::Init(self.config.clone()), STARTUP_TIMEOUT_MS)?;
420 let response = worker.request(
421 Request::Load {
422 name: id.into(),
423 source,
424 },
425 self.config.default_timeout_ms,
426 )?;
427 let Response::Loaded {
428 handle: remote,
429 metadata,
430 hooks,
431 } = response
432 else {
433 return Err(protocol("unexpected load response"));
434 };
435 let handle = PluginHandle::new(self.next);
436 self.next = self
437 .next
438 .checked_add(1)
439 .ok_or_else(|| protocol("plugin handle space exhausted"))?;
440 self.plugins.insert(
441 handle,
442 Loaded {
443 worker,
444 remote,
445 metadata: *metadata,
446 hooks,
447 },
448 );
449 Ok(handle)
450 }
451 fn unload_plugin(&mut self, handle: PluginHandle) -> PluginResult<()> {
452 let loaded = self
453 .plugins
454 .remove(&handle)
455 .ok_or_else(|| protocol("unknown plugin handle"))?;
456 loaded.worker.terminate();
457 Ok(())
458 }
459 fn get_metadata(&self, handle: PluginHandle) -> Option<&PluginMetadata> {
460 self.plugins.get(&handle).map(|p| &p.metadata)
461 }
462 fn has_hook(&self, handle: PluginHandle, name: &str) -> bool {
463 self.plugins
464 .get(&handle)
465 .is_some_and(|p| p.hooks.iter().any(|hook| hook == name))
466 }
467 fn call_hook_sync(
468 &self,
469 handle: PluginHandle,
470 hook: &Hook,
471 ctx: &HookContext,
472 ) -> PluginResult<HookResult> {
473 let plugin = self.loaded(handle)?;
474 let timeout = if hook.name() == "on_render" {
475 self.config.default_timeout_ms.min(50)
476 } else {
477 self.config.default_timeout_ms
478 };
479 hook_response(plugin.worker.request(
480 Request::Hook {
481 handle: plugin.remote,
482 hook: hook.clone(),
483 context: ctx.into(),
484 },
485 timeout,
486 )?)
487 }
488 fn call_hook_async<'a>(
489 &'a self,
490 handle: PluginHandle,
491 hook: &'a Hook,
492 ctx: &'a HookContext,
493 ) -> BoxFuture<'a, PluginResult<HookResult>> {
494 Box::pin(async move {
495 let plugin = self.loaded(handle)?;
496 let worker = Arc::clone(&plugin.worker);
497 let request = Request::Hook {
498 handle: plugin.remote,
499 hook: hook.clone(),
500 context: ctx.into(),
501 };
502 let timeout = self.config.default_timeout_ms;
503 hook_response(worker.request_async(request, timeout).await?)
504 })
505 }
506 fn call_method<'a>(
507 &'a self,
508 handle: PluginHandle,
509 method: &'a str,
510 args: Vec<Value>,
511 ) -> BoxFuture<'a, PluginResult<Value>> {
512 Box::pin(async move {
513 let plugin = self.loaded(handle)?;
514 let request = Request::Method {
515 handle: plugin.remote,
516 name: method.into(),
517 args: args.into_iter().map(Into::into).collect(),
518 };
519 let worker = Arc::clone(&plugin.worker);
520 let timeout = self.config.default_timeout_ms;
521 value_response(worker.request_async(request, timeout).await?)
522 })
523 }
524 fn create_isolated_context(
525 &self,
526 sandbox: &SandboxConfig,
527 ) -> PluginResult<Box<dyn IsolatedContext>> {
528 if sandbox.timeout_ms == 0 {
529 return Err(PluginError::ConfigError {
530 message: "worker timeout must be nonzero".into(),
531 });
532 }
533 let worker = Worker::spawn(&self.executable, self.language)?;
534 worker.request(Request::Isolate(sandbox.clone()), STARTUP_TIMEOUT_MS)?;
535 Ok(Box::new(ProcessContext {
536 worker,
537 timeout: sandbox.timeout_ms,
538 }))
539 }
540 fn loaded_plugins(&self) -> Vec<PluginHandle> {
541 let mut handles: Vec<_> = self.plugins.keys().copied().collect();
542 handles.sort_by_key(|h| h.id());
543 handles
544 }
545 fn shutdown(&mut self) -> PluginResult<()> {
546 for plugin in self.plugins.values() {
547 plugin.worker.terminate();
548 }
549 self.plugins.clear();
550 Ok(())
551 }
552}
553
554fn hook_response(response: Response) -> PluginResult<HookResult> {
555 match response {
556 Response::Hook {
557 handled,
558 prevent_default,
559 stop_propagation,
560 value,
561 error,
562 notices,
563 } => Ok(HookResult {
564 handled,
565 prevent_default,
566 stop_propagation,
567 value: value.map(Into::into),
568 error,
569 notices,
570 plugin: None,
571 }),
572 _ => Err(protocol("unexpected hook response")),
573 }
574}
575fn value_response(response: Response) -> PluginResult<Value> {
576 match response {
577 Response::Value(value) => Ok(value.into()),
578 _ => Err(protocol("unexpected value response")),
579 }
580}
581
582struct ProcessContext {
583 worker: Arc<Worker>,
584 timeout: u64,
585}
586impl Drop for ProcessContext {
587 fn drop(&mut self) {
588 self.worker.terminate();
589 }
590}
591impl ProcessContext {
592 async fn request(&self, request: Request, cancel: CancellationToken) -> PluginResult<Value> {
593 let timeout = self.timeout;
594 tokio::select! {
595 biased;
596 _ = cancel.cancelled() => Err(PluginError::Cancelled { name: "worker".into() }),
599 result = self.worker.request_async(request, timeout) => value_response(result?),
600 }
601 }
602}
603impl IsolatedContext for ProcessContext {
604 fn execute<'a>(
605 &'a self,
606 code: &'a [u8],
607 cancel: CancellationToken,
608 ) -> BoxFuture<'a, PluginResult<Value>> {
609 Box::pin(async move {
610 if code.len() > 1024 * 1024 {
611 return Err(protocol("plugin source exceeds 1 MiB"));
612 }
613 self.request(Request::Execute(code.to_vec()), cancel).await
614 })
615 }
616 fn call_function<'a>(
617 &'a self,
618 name: &'a str,
619 args: Vec<Value>,
620 cancel: CancellationToken,
621 ) -> BoxFuture<'a, PluginResult<Value>> {
622 Box::pin(self.request(
623 Request::Function {
624 name: name.into(),
625 args: args.into_iter().map(Into::into).collect(),
626 },
627 cancel,
628 ))
629 }
630 fn set_global(&mut self, name: &str, value: Value) -> PluginResult<()> {
631 self.worker
632 .request(
633 Request::Set {
634 name: name.into(),
635 value: value.into(),
636 },
637 self.timeout,
638 )
639 .map(|_| ())
640 }
641 fn get_global(&self, name: &str) -> PluginResult<Value> {
642 value_response(
643 self.worker
644 .request(Request::Get(name.into()), self.timeout)?,
645 )
646 }
647}
648
649pub fn run_worker(language: &str) -> PluginResult<()> {
652 let log = WorkerLog::default();
653 let output_log = log.clone();
654 tracing_subscriber::fmt()
655 .with_ansi(false)
656 .with_writer(move || output_log.clone())
657 .try_init()
658 .map_err(|error| protocol(error.to_string()))?;
659 let mut runtime: Box<dyn PluginRuntime> = match language {
660 "lua" => Box::new(crate::lua::LuaRuntime::new()?),
661 other => {
662 return Err(PluginError::RuntimeNotAvailable {
663 runtime: other.into(),
664 });
665 }
666 };
667 let executor = tokio::runtime::Builder::new_current_thread()
668 .enable_all()
669 .build()?;
670 let mut isolated: Option<Box<dyn IsolatedContext>> = None;
671 let mut input = io::stdin().lock();
672 let mut output = io::stdout().lock();
673 loop {
674 let bytes = match read_frame(&mut input) {
675 Ok(bytes) => bytes,
676 Err(error) if error.kind() == io::ErrorKind::UnexpectedEof => break,
677 Err(error) => return Err(error.into()),
678 };
679 let request: Request = decode(&bytes)?;
680 let shutdown = matches!(request, Request::Shutdown);
681 let result = execute_request(&mut *runtime, &mut isolated, &executor, request)
682 .map_err(Failure::from);
683 let (log_text, log_truncated) = log.take();
684 let bytes = match encode(Reply {
685 result,
686 log: log_text,
687 log_truncated,
688 }) {
689 Ok(bytes) => bytes,
690 Err(error) => encode(Reply {
691 result: Err(error.into()),
692 log: String::new(),
693 log_truncated: true,
694 })?,
695 };
696 write_frame(&mut output, &bytes)?;
697 if shutdown {
698 break;
699 }
700 }
701 executor.shutdown_background();
702 Ok(())
703}
704
705fn execute_request(
706 runtime: &mut dyn PluginRuntime,
707 isolated: &mut Option<Box<dyn IsolatedContext>>,
708 executor: &tokio::runtime::Runtime,
709 request: Request,
710) -> PluginResult<Response> {
711 match request {
712 Request::Init(config) => {
713 runtime.init(&config)?;
714 Ok(Response::Unit)
715 }
716 Request::Load { name, source } => {
717 let handle = runtime.load_plugin(&name, &source)?;
718 let metadata = runtime
719 .get_metadata(handle)
720 .ok_or_else(|| protocol("missing plugin metadata"))?
721 .clone();
722 let hooks = HOOK_NAMES
723 .iter()
724 .filter(|name| runtime.has_hook(handle, name))
725 .map(|name| (*name).to_string())
726 .collect();
727 Ok(Response::Loaded {
728 handle: handle.id(),
729 metadata: Box::new(metadata),
730 hooks,
731 })
732 }
733 Request::Hook {
734 handle,
735 hook,
736 context,
737 } => {
738 let result =
739 runtime.call_hook_sync(PluginHandle::new(handle), &hook, &context.into())?;
740 Ok(Response::Hook {
741 handled: result.handled,
742 prevent_default: result.prevent_default,
743 stop_propagation: result.stop_propagation,
744 value: result.value.map(Into::into),
745 error: result.error,
746 notices: result.notices,
747 })
748 }
749 Request::Method { handle, name, args } => Ok(Response::Value(
750 executor
751 .block_on(runtime.call_method(
752 PluginHandle::new(handle),
753 &name,
754 args.into_iter().map(Into::into).collect(),
755 ))?
756 .into(),
757 )),
758 Request::Isolate(sandbox) => {
759 *isolated = Some(runtime.create_isolated_context(&sandbox)?);
760 Ok(Response::Unit)
761 }
762 Request::Shutdown => {
763 runtime.shutdown()?;
764 Ok(Response::Unit)
765 }
766 request => {
767 let context = isolated
768 .as_mut()
769 .ok_or_else(|| protocol("isolated context is not initialized"))?;
770 match request {
771 Request::Execute(code) => Ok(Response::Value(
772 executor
773 .block_on(context.execute(&code, CancellationToken::new()))?
774 .into(),
775 )),
776 Request::Function { name, args } => Ok(Response::Value(
777 executor
778 .block_on(context.call_function(
779 &name,
780 args.into_iter().map(Into::into).collect(),
781 CancellationToken::new(),
782 ))?
783 .into(),
784 )),
785 Request::Set { name, value } => {
786 context.set_global(&name, value.into())?;
787 Ok(Response::Unit)
788 }
789 Request::Get(name) => Ok(Response::Value(context.get_global(&name)?.into())),
790 _ => Err(protocol("invalid isolated request")),
791 }
792 }
793 }
794}
795
796#[cfg(all(test, unix))]
798mod tests {
799 use super::*;
800
801 #[cfg(unix)]
802 #[test]
803 fn expired_encoding_budget_prevents_worker_dispatch() {
804 let root = tempfile::tempdir().unwrap();
805 let script = root.path().join("worker");
806 std::fs::write(&script, b"#!/bin/sh\ncat \"$0.response\"\nexec sleep 10\n").unwrap();
807 let bytes = encode(Reply {
808 result: Ok(Response::Unit),
809 log: String::new(),
810 log_truncated: false,
811 })
812 .unwrap();
813 let mut response = std::fs::File::create(root.path().join("worker.response")).unwrap();
814 write_frame(&mut response, &bytes).unwrap();
815 let worker = spawn_script(&script).unwrap();
816 let result = worker.request_reserved_with_encoding(Request::Shutdown, 20, |request| {
817 std::thread::sleep(Duration::from_millis(50));
818 encode(request)
819 });
820 assert!(
821 matches!(result, Err(PluginError::Timeout { timeout_ms: 20, .. })),
822 "expired encoding returned error {:?}",
823 result.err()
824 );
825 assert!(!worker.alive.load(Ordering::Acquire));
826 assert!(worker.permit.lock().unwrap().is_none());
827 assert!(worker.child.lock().unwrap().try_wait().unwrap().is_some());
828 }
829
830 #[cfg(unix)]
831 fn script_command(script: &Path) -> Command {
832 let mut command = Command::new("/bin/sh");
835 command.arg(script).args(["__plugin-worker", "lua"]);
836 command
837 }
838
839 #[cfg(unix)]
840 fn spawn_script(script: &Path) -> PluginResult<Arc<Worker>> {
841 Worker::spawn_command(script_command(script), |task| {
842 std::thread::Builder::new().spawn(task).map(|_| ())
843 })
844 }
845
846 #[cfg(unix)]
847 #[test]
848 fn startup_guard_closes_pipes_inherited_by_worker_descendants() {
849 use std::io::BufRead;
850 use std::os::unix::process::CommandExt;
851 let mut child = Command::new("/bin/sh")
852 .args(["-c", "sleep 10 & printf 'ready\\n'; wait"])
853 .process_group(0)
854 .stdout(Stdio::piped())
855 .spawn()
856 .unwrap();
857 let stdout = child.stdout.take().unwrap();
858 let starting = StartingWorker(Some(child));
859 let (send, receive) = mpsc::channel();
860 let reader = std::thread::spawn(move || {
861 let mut reader = io::BufReader::new(stdout);
862 let mut ready = String::new();
863 reader.read_line(&mut ready).unwrap();
864 send.send(ready).unwrap();
865 let mut remaining = Vec::new();
866 reader.read_to_end(&mut remaining).unwrap();
867 let _ = send.send("eof".into());
868 });
869 assert_eq!(
870 receive.recv_timeout(Duration::from_secs(2)).unwrap(),
871 "ready\n"
872 );
873 drop(starting);
875 assert_eq!(receive.recv_timeout(Duration::from_secs(2)).unwrap(), "eof");
876 reader.join().unwrap();
877 }
878
879 #[cfg(unix)]
880 #[test]
881 fn failed_io_thread_startup_reaps_the_spawned_worker() {
882 let root = tempfile::tempdir().unwrap();
883 let executable = root.path().join("worker");
884 let pid_file = root.path().join("worker.pid");
885 std::fs::write(
886 &executable,
887 b"#!/bin/sh\nprintf '%s' \"$$\" > \"$0.pid\"\nexec sleep 10\n",
888 )
889 .unwrap();
890 let mut pid = None;
891 let error = Worker::spawn_command(script_command(&executable), |_task| {
892 let deadline = std::time::Instant::now() + Duration::from_secs(2);
893 while std::time::Instant::now() < deadline {
894 pid = std::fs::read_to_string(&pid_file)
895 .ok()
896 .and_then(|text| text.parse::<libc::pid_t>().ok());
897 if pid.is_some() {
898 break;
899 }
900 std::thread::sleep(Duration::from_millis(5));
901 }
902 assert!(pid.is_some(), "worker did not announce startup");
903 Err(io::Error::other("injected I/O thread startup failure"))
904 })
905 .err()
906 .unwrap();
907 assert!(
908 error
909 .to_string()
910 .contains("injected I/O thread startup failure")
911 );
912 let mut status = 0;
913 assert_eq!(
915 unsafe { libc::waitpid(pid.unwrap(), &mut status, libc::WNOHANG) },
916 -1
917 );
918 assert_eq!(
919 io::Error::last_os_error().raw_os_error(),
920 Some(libc::ECHILD)
921 );
922 }
923
924 #[cfg(unix)]
925 #[tokio::test]
926 async fn cancellation_terminates_the_owned_request_but_not_another_call() {
927 use std::future::Future;
928 use std::task::Poll;
929 let root = tempfile::tempdir().unwrap();
930 let executable = root.path().join("worker");
931 std::fs::write(&executable, b"#!/bin/sh\nexec sleep 10\n").unwrap();
932 let worker = spawn_script(&executable).unwrap();
933 let context = ProcessContext {
934 worker: Arc::clone(&worker),
935 timeout: 5000,
936 };
937 let active_cancel = CancellationToken::new();
938 let mut active =
939 Box::pin(context.request(Request::Execute(Vec::new()), active_cancel.clone()));
940 std::future::poll_fn(|cx| {
941 assert!(active.as_mut().poll(cx).is_pending());
942 Poll::Ready(())
943 })
944 .await;
945 let other_cancel = CancellationToken::new();
946 other_cancel.cancel();
947 assert!(matches!(
948 context
949 .request(Request::Execute(Vec::new()), other_cancel)
950 .await,
951 Err(PluginError::Cancelled { .. })
952 ));
953 assert!(worker.alive.load(Ordering::Acquire));
954 assert!(worker.busy.load(Ordering::Acquire));
955 active_cancel.cancel();
956 assert!(matches!(active.await, Err(PluginError::Cancelled { .. })));
957 assert!(!worker.alive.load(Ordering::Acquire));
958 assert!(!worker.busy.load(Ordering::Acquire));
959 assert!(worker.child.lock().unwrap().try_wait().unwrap().is_some());
960 }
961
962 #[cfg(unix)]
963 #[tokio::test]
964 async fn dropping_async_request_terminates_only_its_owned_worker_request() {
965 use std::future::Future;
966 use std::task::Poll;
967 let root = tempfile::tempdir().unwrap();
968 let executable = root.path().join("worker");
969 std::fs::write(&executable, b"#!/bin/sh\nexec sleep 10\n").unwrap();
970 let worker = spawn_script(&executable).unwrap();
971 let mut first = Box::pin(worker.request_async(Request::Shutdown, 5000));
972 std::future::poll_fn(|cx| {
973 assert!(first.as_mut().poll(cx).is_pending());
974 Poll::Ready(())
975 })
976 .await;
977 assert!(worker.busy.load(Ordering::Acquire));
978 let error = worker
980 .request_async(Request::Shutdown, 5000)
981 .await
982 .err()
983 .unwrap();
984 assert!(error.to_string().contains("active request"));
985 assert!(worker.alive.load(Ordering::Acquire));
986 drop(first);
987 assert!(!worker.alive.load(Ordering::Acquire));
988 assert!(!worker.busy.load(Ordering::Acquire));
989 assert!(worker.child.lock().unwrap().try_wait().unwrap().is_some());
990 let error = worker.request(Request::Shutdown, 5000).err().unwrap();
991 assert!(error.to_string().contains("worker terminated"));
992 }
993
994 #[cfg(unix)]
995 #[test]
996 fn async_deadline_includes_blocking_executor_queue_time() {
997 for abandon in [false, true] {
998 let root = tempfile::tempdir().unwrap();
999 let executable = root.path().join("worker");
1000 std::fs::write(&executable, b"#!/bin/sh\nexec sleep 30\n").unwrap();
1001 let worker = spawn_script(&executable).unwrap();
1002 let executor = tokio::runtime::Builder::new_current_thread()
1003 .enable_all()
1004 .max_blocking_threads(1)
1005 .build()
1006 .unwrap();
1007 let (release, wait) = mpsc::channel::<()>();
1008 let (ready, started) = mpsc::channel();
1009 let blocker = executor.spawn_blocking(move || {
1010 ready.send(()).unwrap();
1011 let _ = wait.recv();
1012 });
1013 started.recv_timeout(Duration::from_secs(5)).unwrap();
1014 let dispatched = Arc::new(AtomicBool::new(false));
1015 let observed = Arc::clone(&dispatched);
1016 executor.block_on(async {
1017 let future = worker.request_async_with_dispatch(
1018 Request::Shutdown,
1019 25,
1020 move |worker, request, timeout| {
1021 observed.store(true, Ordering::Release);
1022 worker.request_reserved(request, timeout)
1023 },
1024 );
1025 if abandon {
1026 use std::future::Future;
1027 let mut future = Box::pin(future);
1028 std::future::poll_fn(|cx| {
1029 assert!(future.as_mut().poll(cx).is_pending());
1030 std::task::Poll::Ready(())
1031 })
1032 .await;
1033 drop(future);
1034 } else {
1035 let result = tokio::time::timeout(Duration::from_secs(2), future)
1036 .await
1037 .expect("worker deadline must expire while blocking executor is occupied");
1038 assert!(matches!(
1039 result,
1040 Err(PluginError::Timeout { timeout_ms: 25, .. })
1041 ));
1042 }
1043 assert!(!worker.alive.load(Ordering::Acquire));
1044 assert!(worker.permit.lock().unwrap().is_none());
1045 assert!(worker.child.lock().unwrap().try_wait().unwrap().is_some());
1046 });
1047 drop(release);
1048 executor.block_on(blocker).unwrap();
1049 drop(executor);
1052 assert!(!worker.busy.load(Ordering::Acquire));
1053 assert!(
1054 !dispatched.load(Ordering::Acquire),
1055 "expired queued work must never dispatch"
1056 );
1057 }
1058 }
1059
1060 #[cfg(unix)]
1061 #[test]
1062 fn disconnected_io_channel_terminates_worker_and_releases_permit() {
1063 let root = tempfile::tempdir().unwrap();
1064 let executable = root.path().join("worker");
1065 std::fs::write(&executable, b"#!/bin/sh\nexec sleep 30\n").unwrap();
1066 let worker = Worker::spawn_command(script_command(&executable), |task| {
1069 drop(task);
1070 Ok(())
1071 })
1072 .unwrap();
1073 let error = worker.request(Request::Shutdown, 2000).err().unwrap();
1074 assert!(error.to_string().contains("channel closed"), "{error:?}");
1075 assert!(!worker.alive.load(Ordering::Acquire));
1076 assert!(worker.permit.lock().unwrap().is_none());
1077 assert!(worker.child.lock().unwrap().try_wait().unwrap().is_some());
1078 assert!(!worker.busy.load(Ordering::Acquire));
1079 let error = worker.request(Request::Shutdown, 2000).err().unwrap();
1080 assert!(error.to_string().contains("worker terminated"), "{error:?}");
1081 }
1082
1083 #[cfg(unix)]
1084 #[test]
1085 fn mismatched_reply_types_terminate_workers_before_reuse() {
1086 let cases = [
1087 (
1088 Request::Init(PluginConfig::default()),
1089 Response::Value(WireValue::Null),
1090 ),
1091 (Request::Get("example".into()), Response::Unit),
1092 (
1093 Request::Load {
1094 name: "example".into(),
1095 source: "example.lua".into(),
1096 },
1097 Response::Unit,
1098 ),
1099 (
1100 Request::Hook {
1101 handle: 0,
1102 hook: Hook::OnStartup,
1103 context: (&HookContext::new()).into(),
1104 },
1105 Response::Unit,
1106 ),
1107 (
1108 Request::Set {
1109 name: "example".into(),
1110 value: WireValue::Null,
1111 },
1112 Response::Value(WireValue::Null),
1113 ),
1114 ];
1115 for (request, response) in cases {
1116 let root = tempfile::tempdir().unwrap();
1117 let executable = root.path().join("worker");
1118 std::fs::write(
1119 &executable,
1120 b"#!/bin/sh\ncat \"$0.response\"\nexec sleep 10\n",
1121 )
1122 .unwrap();
1123 let bytes = encode(Reply {
1124 result: Ok(response),
1125 log: String::new(),
1126 log_truncated: false,
1127 })
1128 .unwrap();
1129 let mut file = std::fs::File::create(root.path().join("worker.response")).unwrap();
1130 write_frame(&mut file, &bytes).unwrap();
1131 let worker = spawn_script(&executable).unwrap();
1132 let error = worker.request(request, 2000).err().unwrap();
1133 assert!(
1134 error.to_string().contains("unexpected plugin response"),
1135 "expected response-kind rejection, received {error:?}"
1136 );
1137 assert!(!worker.alive.load(Ordering::Acquire));
1138 let error = worker.request(Request::Shutdown, 2000).err().unwrap();
1139 assert!(error.to_string().contains("worker terminated"));
1140 }
1141 }
1142
1143 #[cfg(unix)]
1144 #[test]
1145 fn oversized_log_reply_terminates_worker_and_rejects_reuse() {
1146 let root = tempfile::tempdir().unwrap();
1147 let executable = root.path().join("worker");
1148 std::fs::write(
1149 &executable,
1150 b"#!/bin/sh\ncat \"$0.response\"\nexec sleep 10\n",
1151 )
1152 .unwrap();
1153 let bytes = encode(Reply {
1154 result: Ok(Response::Unit),
1155 log: "x".repeat(MAX_LOG_BYTES + 1),
1156 log_truncated: false,
1157 })
1158 .unwrap();
1159 let mut response = std::fs::File::create(root.path().join("worker.response")).unwrap();
1160 write_frame(&mut response, &bytes).unwrap();
1161 let worker = spawn_script(&executable).unwrap();
1162 let error = worker.request(Request::Shutdown, 2000).err().unwrap();
1163 assert!(error.to_string().contains("log exceeds 64 KiB"));
1164 assert!(!worker.alive.load(Ordering::Acquire));
1165 let error = worker.request(Request::Shutdown, 2000).err().unwrap();
1166 assert!(error.to_string().contains("worker terminated"));
1167 }
1168}