capsule_runtime/guest_worker/
worker.rs

1//! The worker side: holds the engine, the compiled guests and every store, and answers the
2//! supervisor's requests. Every host import that reaches past WASI, clocks, random and the
3//! execution limits becomes a host call to the supervisor through a [`RemoteLink`].
4
5use std::collections::HashMap;
6use std::os::unix::net::UnixStream;
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::{Arc, Mutex};
9use std::time::Duration;
10
11use bytes::Bytes;
12use tokio::sync::{mpsc, oneshot};
13use wasmtime::component::Component;
14
15use super::wire::{
16    read_frame, write_frame, FromWorker, HostReply, HostRequest, HttpError, HttpRequestHead,
17    HttpResponseHead, Reply, StreamEvent, TaskIoCall, TaskIoReply, ToWorker,
18};
19use super::{WorkerKind, PROTOCOL_VERSION};
20use crate::bindings::hook::murmur::conversation::read::MessagePage;
21use crate::bindings::hook::murmur::runtime::inference::{InferenceRequest, InferenceResponse};
22use crate::bindings::hook::murmur::task_io::read::IoError;
23use crate::inference_import::HookInferencePurpose;
24use crate::limits::CallDeadlines;
25use crate::process_driver::{DriverAnswer, DriverError, ProcessDriver};
26use crate::runtime::{execute_tool_call, CapsuleCall, CapsuleReply, ToolHostSide};
27use crate::session_process::{ProcessCall, ProcessReply};
28
29type BodySender = mpsc::UnboundedSender<Result<Bytes, HttpError>>;
30
31/// The worker's connection: the writer every task shares, and what is waiting on the supervisor.
32/// Every pending host call and response body is keyed by its host id and remembers the call it
33/// belongs to, so a call that ends or is cancelled leaves nothing behind.
34pub(crate) struct WorkerShared {
35    writer: Mutex<UnixStream>,
36    next_host: AtomicU64,
37    pending: Mutex<HashMap<u64, (u64, oneshot::Sender<HostReply>)>>,
38    bodies: Mutex<HashMap<u64, (u64, BodySender)>>,
39}
40
41impl WorkerShared {
42    /// Forgets every host call and response body of call `call`.
43    fn forget_call(&self, call: u64) {
44        self.pending
45            .lock()
46            .unwrap_or_else(std::sync::PoisonError::into_inner)
47            .retain(|_, (owner, _)| *owner != call);
48        self.bodies
49            .lock()
50            .unwrap_or_else(std::sync::PoisonError::into_inner)
51            .retain(|_, (owner, _)| *owner != call);
52    }
53
54    fn send(&self, message: &FromWorker) {
55        let mut writer = self
56            .writer
57            .lock()
58            .unwrap_or_else(std::sync::PoisonError::into_inner);
59        // A failed write means the supervisor is gone; the reader sees the same and ends the
60        // worker.
61        let _ = write_frame(&mut *writer, message);
62    }
63}
64
65/// One guest store's way to the supervisor: every host call it makes names the call id the
66/// supervisor issued for the call running on it, and is answered only while that call is in
67/// flight. A session-long hook instance's link is pointed at each new call before it runs.
68pub struct RemoteLink {
69    call: AtomicU64,
70    shared: Arc<WorkerShared>,
71}
72
73impl RemoteLink {
74    fn new(call: u64, shared: &Arc<WorkerShared>) -> Arc<Self> {
75        Arc::new(Self {
76            call: AtomicU64::new(call),
77            shared: Arc::clone(shared),
78        })
79    }
80
81    fn call(&self) -> u64 {
82        self.call.load(Ordering::SeqCst)
83    }
84
85    fn set_call(&self, call: u64) {
86        self.call.store(call, Ordering::SeqCst);
87    }
88
89    async fn host(&self, request: HostRequest) -> HostReply {
90        let host = self.shared.next_host.fetch_add(1, Ordering::Relaxed);
91        let call = self.call();
92        let (tx, rx) = oneshot::channel();
93        self.shared
94            .pending
95            .lock()
96            .unwrap_or_else(std::sync::PoisonError::into_inner)
97            .insert(host, (call, tx));
98        self.shared.send(&FromWorker::HostCall {
99            host,
100            call,
101            request,
102        });
103        rx.await.unwrap_or_else(|_| {
104            HostReply::Refused("the supervisor's connection closed".to_string())
105        })
106    }
107
108    pub(crate) fn stream_event(&self, event: StreamEvent) {
109        self.shared.send(&FromWorker::Event {
110            call: self.call(),
111            event,
112        });
113    }
114
115    pub(crate) fn run_started(&self, deadline: Duration) {
116        self.shared.send(&FromWorker::RunStarted {
117            call: self.call(),
118            deadline,
119        });
120    }
121
122    pub(crate) async fn inference(
123        &self,
124        purpose: HookInferencePurpose,
125        request: InferenceRequest,
126    ) -> Result<InferenceResponse, String> {
127        match self.host(HostRequest::Inference { purpose, request }).await {
128            HostReply::Inference(result) => result,
129            HostReply::Refused(reason) => Err(reason),
130            other => Err(format!("unexpected reply {other:?}")),
131        }
132    }
133
134    pub(crate) async fn task_io(&self, call: TaskIoCall) -> TaskIoReply {
135        let reads_len = matches!(call, TaskIoCall::InputLen(_) | TaskIoCall::OutputLen);
136        match self.host(HostRequest::TaskIo(call)).await {
137            HostReply::TaskIo(reply) => reply,
138            // A refused read is the same `not-granted` value an ungranted hook branches on.
139            _ if reads_len => TaskIoReply::Len(Err(IoError::NotGranted)),
140            _ => TaskIoReply::Text(Err(IoError::NotGranted)),
141        }
142    }
143
144    pub(crate) async fn conversation(
145        &self,
146        cursor: Option<String>,
147        limit: u32,
148    ) -> Result<MessagePage, String> {
149        match self.host(HostRequest::Conversation { cursor, limit }).await {
150            HostReply::Conversation(result) => result,
151            HostReply::Refused(reason) => Err(reason),
152            other => Err(format!("unexpected reply {other:?}")),
153        }
154    }
155
156    pub(crate) async fn capsule(&self, call: CapsuleCall) -> CapsuleReply {
157        match self.host(HostRequest::Capsule(Box::new(call))).await {
158            HostReply::Capsule(reply) => *reply,
159            HostReply::Refused(reason) => CapsuleReply::Unanswered(reason),
160            other => CapsuleReply::Unanswered(format!("unexpected reply {other:?}")),
161        }
162    }
163
164    pub(crate) async fn request_input(&self, prompt: String) -> wasmtime::Result<String> {
165        match self.host(HostRequest::RequestInput { prompt }).await {
166            HostReply::Input(Ok(answer)) => Ok(answer),
167            HostReply::Input(Err(error)) | HostReply::Refused(error) => {
168                Err(wasmtime::Error::msg(error))
169            }
170            other => Err(wasmtime::Error::msg(format!("unexpected reply {other:?}"))),
171        }
172    }
173
174    pub(crate) async fn process(&self, call: ProcessCall) -> ProcessReply {
175        match self.host(HostRequest::Process(call)).await {
176            HostReply::Process(reply) => reply,
177            HostReply::Refused(reason) => ProcessReply::Unanswered(reason),
178            other => ProcessReply::Unanswered(format!("unexpected reply {other:?}")),
179        }
180    }
181
182    /// Sends `head` and returns the response head and a receiver for its body.
183    pub(crate) async fn http(
184        &self,
185        head: HttpRequestHead,
186    ) -> Result<
187        (
188            HttpResponseHead,
189            mpsc::UnboundedReceiver<Result<Bytes, HttpError>>,
190        ),
191        HttpError,
192    > {
193        let host = self.shared.next_host.fetch_add(1, Ordering::Relaxed);
194        let call = self.call();
195        let (body_tx, body_rx) = mpsc::unbounded_channel();
196        // Registered before the request is sent, so no body frame can arrive unrouted.
197        self.shared
198            .bodies
199            .lock()
200            .unwrap_or_else(std::sync::PoisonError::into_inner)
201            .insert(host, (call, body_tx));
202        let (tx, rx) = oneshot::channel();
203        self.shared
204            .pending
205            .lock()
206            .unwrap_or_else(std::sync::PoisonError::into_inner)
207            .insert(host, (call, tx));
208        self.shared.send(&FromWorker::HostCall {
209            host,
210            call,
211            request: HostRequest::Http(head),
212        });
213        let reply = rx.await.unwrap_or_else(|_| {
214            HostReply::Refused("the supervisor's connection closed".to_string())
215        });
216        let forget = || {
217            self.shared
218                .bodies
219                .lock()
220                .unwrap_or_else(std::sync::PoisonError::into_inner)
221                .remove(&host);
222        };
223        match reply {
224            HostReply::Http(Ok(head)) => Ok((head, body_rx)),
225            HostReply::Http(Err(error)) => {
226                forget();
227                Err(error)
228            }
229            HostReply::Refused(reason) => {
230                forget();
231                Err(HttpError::InternalError(Some(reason)))
232            }
233            other => {
234                forget();
235                Err(HttpError::InternalError(Some(format!(
236                    "unexpected reply {other:?}"
237                ))))
238            }
239        }
240    }
241}
242
243/// Serves the supervisor on `stream` until `Shutdown` or end of file, as a worker of `kind`.
244/// Returns the process exit code a worker process ends with.
245#[allow(unsafe_code)]
246pub(crate) fn serve(kind: WorkerKind, stream: UnixStream) -> i32 {
247    let Ok(writer) = stream.try_clone() else {
248        return 2;
249    };
250    let shared = Arc::new(WorkerShared {
251        writer: Mutex::new(writer),
252        next_host: AtomicU64::new(1),
253        pending: Mutex::new(HashMap::new()),
254        bodies: Mutex::new(HashMap::new()),
255    });
256    shared.send(&FromWorker::Hello {
257        protocol: PROTOCOL_VERSION,
258        version: env!("CARGO_PKG_VERSION").to_string(),
259        pid: std::process::id(),
260    });
261    let mut reader = stream;
262    match read_frame::<ToWorker>(&mut reader) {
263        Ok(ToWorker::Init { kind: told }) if told == kind => {}
264        _ => return 2,
265    }
266    let Ok(engine) = crate::runtime::build_engine() else {
267        return 2;
268    };
269    let _ticker = crate::limits::EpochTicker::spawn(&engine);
270    let Ok(runtime) = tokio::runtime::Builder::new_multi_thread()
271        .worker_threads(kind.runtime_threads())
272        .thread_name(format!("guest-worker-{}", kind.as_str()))
273        .enable_all()
274        .build()
275    else {
276        return 2;
277    };
278    let state = Arc::new(ServerState {
279        engine,
280        shared: Arc::clone(&shared),
281        guests: Mutex::new(HashMap::new()),
282        hooks: Mutex::new(HashMap::new()),
283        drivers: Mutex::new(HashMap::new()),
284    });
285    let mut running: HashMap<u64, tokio::task::AbortHandle> = HashMap::new();
286    let (done_tx, done_rx) = std::sync::mpsc::channel::<u64>();
287
288    while let Ok(message) = read_frame::<ToWorker>(&mut reader) {
289        while let Ok(done) = done_rx.try_recv() {
290            running.remove(&done);
291        }
292        let spawn_call =
293            |req: u64, body: std::pin::Pin<Box<dyn std::future::Future<Output = Reply> + Send>>| {
294                let shared = Arc::clone(&shared);
295                let done = done_tx.clone();
296                runtime
297                    .spawn(async move {
298                        let reply = body.await;
299                        shared.forget_call(req);
300                        shared.send(&FromWorker::Reply { req, reply });
301                        let _ = done.send(req);
302                    })
303                    .abort_handle()
304            };
305        match message {
306            ToWorker::Init { .. } => {}
307            ToWorker::Load {
308                req,
309                guest,
310                name,
311                form,
312            } => {
313                // SAFETY: `form` is a component this same binary serialized with an engine built by
314                // the same `build_engine`, handed over a connection only the supervisor holds.
315                let loaded =
316                    unsafe { Component::deserialize(&state.engine, &form) }.map_err(|error| {
317                        format!("failed to load '{name}' in the guest worker: {error}")
318                    });
319                let reply = match loaded {
320                    Ok(component) => {
321                        state
322                            .guests
323                            .lock()
324                            .unwrap_or_else(std::sync::PoisonError::into_inner)
325                            .insert(guest, component);
326                        Ok(())
327                    }
328                    Err(error) => Err(error),
329                };
330                shared.send(&FromWorker::Reply {
331                    req,
332                    reply: Reply::Loaded(reply),
333                });
334            }
335            ToWorker::RunTool {
336                req,
337                guest,
338                spec,
339                input,
340            } => {
341                let state = Arc::clone(&state);
342                let task = spawn_call(
343                    req,
344                    Box::pin(async move {
345                        let deadlines = CallDeadlines::default();
346                        let result = match state.guest(guest) {
347                            None => Err(format!(
348                                "tool '{}' is not loaded in the guest worker",
349                                spec.name
350                            )),
351                            Some(component) => {
352                                let link = RemoteLink::new(req, &state.shared);
353                                Box::pin(execute_tool_call(
354                                    &state.engine,
355                                    &spec,
356                                    ToolHostSide::Remote(link),
357                                    &deadlines,
358                                    &component,
359                                    input,
360                                ))
361                                .await
362                            }
363                        };
364                        #[cfg(all(feature = "escape-probe", target_os = "linux"))]
365                        let result = result.map(|mut result| {
366                            super::escape_probe::answer(&spec, &mut result);
367                            result
368                        });
369                        Reply::Tool {
370                            result,
371                            deadlines: deadlines.take(),
372                        }
373                    }),
374                );
375                running.insert(req, task);
376            }
377            ToWorker::HookOpen {
378                req,
379                guest,
380                instance,
381                spec,
382            } => {
383                let state = Arc::clone(&state);
384                let task = spawn_call(
385                    req,
386                    Box::pin(async move {
387                        let deadlines = CallDeadlines::default();
388                        let result = match state.guest(guest) {
389                            None => Err(format!(
390                                "hook '{}' is not loaded in the guest worker",
391                                spec.name
392                            )),
393                            Some(component) => {
394                                let link = RemoteLink::new(req, &state.shared);
395                                crate::hooks::open_remote_instance(
396                                    &state.engine,
397                                    &component,
398                                    &spec,
399                                    Arc::clone(&link),
400                                    deadlines.clone(),
401                                )
402                                .await
403                                .map(|hook| {
404                                    let fns = hook.exported_fns();
405                                    state
406                                        .hooks
407                                        .lock()
408                                        .unwrap_or_else(std::sync::PoisonError::into_inner)
409                                        .insert(
410                                            instance,
411                                            Arc::new(tokio::sync::Mutex::new(WorkerHook {
412                                                hook,
413                                                link,
414                                                deadlines: deadlines.clone(),
415                                                in_call: None,
416                                            })),
417                                        );
418                                    fns
419                                })
420                            }
421                        };
422                        Reply::HookOpened {
423                            result,
424                            deadlines: deadlines.take(),
425                        }
426                    }),
427                );
428                running.insert(req, task);
429            }
430            ToWorker::HookCall {
431                req,
432                instance,
433                fn_name,
434                arg,
435                purpose,
436            } => {
437                let state = Arc::clone(&state);
438                let task = spawn_call(
439                    req,
440                    Box::pin(async move {
441                        let hook = state
442                            .hooks
443                            .lock()
444                            .unwrap_or_else(std::sync::PoisonError::into_inner)
445                            .get(&instance)
446                            .cloned();
447                        let Some(hook) = hook else {
448                            return Reply::Hook {
449                                result: Err(format!(
450                                    "hook instance {instance} is not held by the guest worker"
451                                )),
452                                deadlines: Vec::new(),
453                            };
454                        };
455                        let mut hook = hook.lock().await;
456                        hook.link.set_call(req);
457                        let WorkerHook {
458                            hook: instance,
459                            deadlines,
460                            in_call,
461                            ..
462                        } = &mut *hook;
463                        // A call dropped mid-guest by `Cancel` left the instance where its future
464                        // stopped: it is cut, as an instance the wall-clock bound stopped is.
465                        if let Some(stopped) = in_call.take() {
466                            crate::hooks::mark_cut(instance, &stopped);
467                        }
468                        *in_call = Some(fn_name.clone());
469                        let result =
470                            crate::hooks::call_instance(instance, &fn_name, *arg, purpose).await;
471                        *in_call = None;
472                        Reply::Hook {
473                            result,
474                            deadlines: deadlines.take(),
475                        }
476                    }),
477                );
478                running.insert(req, task);
479            }
480            ToWorker::HookStage {
481                req,
482                guest,
483                spec,
484                event,
485            } => {
486                let state = Arc::clone(&state);
487                let task = spawn_call(
488                    req,
489                    Box::pin(async move {
490                        let deadlines = CallDeadlines::default();
491                        let result = match state.guest(guest) {
492                            None => Err(format!(
493                                "hook '{}' is not loaded in the guest worker",
494                                spec.name
495                            )),
496                            Some(component) => {
497                                let link = RemoteLink::new(req, &state.shared);
498                                crate::hooks::stage_remote_once(
499                                    &state.engine,
500                                    &component,
501                                    &spec,
502                                    link,
503                                    deadlines.clone(),
504                                    event,
505                                )
506                                .await
507                                .map(Some)
508                            }
509                        };
510                        Reply::Hook {
511                            result,
512                            deadlines: deadlines.take(),
513                        }
514                    }),
515                );
516                running.insert(req, task);
517            }
518            ToWorker::HookClose { instance } => {
519                state
520                    .hooks
521                    .lock()
522                    .unwrap_or_else(std::sync::PoisonError::into_inner)
523                    .remove(&instance);
524            }
525            ToWorker::DriverOpen {
526                req,
527                guest,
528                instance,
529                name,
530                version,
531            } => {
532                let state = Arc::clone(&state);
533                let task = spawn_call(
534                    req,
535                    Box::pin(async move {
536                        let deadlines = CallDeadlines::default();
537                        let result = match state.guest(guest) {
538                            None => Err(DriverError::Other(format!(
539                                "process driver '{name}' is not loaded in the guest worker"
540                            ))),
541                            Some(component) => ProcessDriver::instantiate(
542                                &state.engine,
543                                &component,
544                                &name,
545                                &version,
546                                deadlines.clone(),
547                            )
548                            .await
549                            .map(|driver| {
550                                state
551                                    .drivers
552                                    .lock()
553                                    .unwrap_or_else(std::sync::PoisonError::into_inner)
554                                    .insert(
555                                        instance,
556                                        Arc::new(tokio::sync::Mutex::new((
557                                            driver,
558                                            deadlines.clone(),
559                                        ))),
560                                    );
561                                DriverAnswer::Opened
562                            })
563                            .map_err(DriverError::from_runtime),
564                        };
565                        Reply::Driver {
566                            result,
567                            deadlines: deadlines.take(),
568                        }
569                    }),
570                );
571                running.insert(req, task);
572            }
573            ToWorker::DriverCall {
574                req,
575                instance,
576                call,
577            } => {
578                let state = Arc::clone(&state);
579                let task = spawn_call(
580                    req,
581                    Box::pin(async move {
582                        let driver = state
583                            .drivers
584                            .lock()
585                            .unwrap_or_else(std::sync::PoisonError::into_inner)
586                            .get(&instance)
587                            .cloned();
588                        let Some(driver) = driver else {
589                            return Reply::Driver {
590                                result: Err(DriverError::Other(format!(
591                                    "process driver instance {instance} is not held by the guest worker"
592                                ))),
593                                deadlines: Vec::new(),
594                            };
595                        };
596                        let mut driver = driver.lock().await;
597                        let (driver, deadlines) = &mut *driver;
598                        let result = driver.call(*call).await.map_err(DriverError::from_runtime);
599                        Reply::Driver {
600                            result,
601                            deadlines: deadlines.take(),
602                        }
603                    }),
604                );
605                running.insert(req, task);
606            }
607            ToWorker::DriverClose { instance } => {
608                state
609                    .drivers
610                    .lock()
611                    .unwrap_or_else(std::sync::PoisonError::into_inner)
612                    .remove(&instance);
613            }
614            ToWorker::RunCapsule { req, guest, spec } => {
615                let state = Arc::clone(&state);
616                let task = spawn_call(
617                    req,
618                    Box::pin(async move {
619                        let deadlines = CallDeadlines::default();
620                        let result = match state.guest(guest) {
621                            None => Err(crate::runtime::CapsuleRunError::Runtime(
622                                "the capsule component is not loaded in the guest worker"
623                                    .to_string(),
624                            )),
625                            Some(component) => {
626                                let link = RemoteLink::new(req, &state.shared);
627                                crate::runtime::run_capsule_guest(
628                                    &state.engine,
629                                    &component,
630                                    &spec,
631                                    link,
632                                    &deadlines,
633                                )
634                                .await
635                            }
636                        };
637                        Reply::Capsule {
638                            result,
639                            deadlines: deadlines.take(),
640                        }
641                    }),
642                );
643                running.insert(req, task);
644            }
645            ToWorker::Cancel { req } => {
646                if let Some(task) = running.remove(&req) {
647                    task.abort();
648                }
649                shared.forget_call(req);
650            }
651            ToWorker::HostReply { host, reply } => {
652                let waiter = shared
653                    .pending
654                    .lock()
655                    .unwrap_or_else(std::sync::PoisonError::into_inner)
656                    .remove(&host);
657                if let Some((_, waiter)) = waiter {
658                    let _ = waiter.send(reply);
659                }
660            }
661            ToWorker::BodyChunk { stream, bytes } => {
662                let bodies = shared
663                    .bodies
664                    .lock()
665                    .unwrap_or_else(std::sync::PoisonError::into_inner);
666                if let Some((_, body)) = bodies.get(&stream) {
667                    let _ = body.send(Ok(Bytes::from(bytes)));
668                }
669            }
670            ToWorker::BodyEnd { stream, error } => {
671                let body = shared
672                    .bodies
673                    .lock()
674                    .unwrap_or_else(std::sync::PoisonError::into_inner)
675                    .remove(&stream);
676                if let (Some((_, body)), Some(error)) = (body, error) {
677                    let _ = body.send(Err(error));
678                }
679            }
680            ToWorker::Shutdown => break,
681        }
682    }
683    for (_, task) in running {
684        task.abort();
685    }
686    runtime.shutdown_background();
687    0
688}
689
690/// A link on a connection no supervisor reads, for tests that need a store's host side but never
691/// reach it. The other end is returned so the connection stays open.
692#[cfg(test)]
693pub(crate) fn detached_link() -> (Arc<RemoteLink>, UnixStream) {
694    let (ours, theirs) = UnixStream::pair().expect("a socket pair");
695    let shared = Arc::new(WorkerShared {
696        writer: Mutex::new(ours),
697        next_host: AtomicU64::new(1),
698        pending: Mutex::new(HashMap::new()),
699        bodies: Mutex::new(HashMap::new()),
700    });
701    (RemoteLink::new(0, &shared), theirs)
702}
703
704/// A session-long hook instance in the hooks worker, with the link its host imports use and the
705/// latch its wall-clock bounds record into, drained after each call.
706struct WorkerHook {
707    hook: crate::hooks::HookInstance,
708    link: Arc<RemoteLink>,
709    deadlines: CallDeadlines,
710    /// The lifecycle function a call is running, while one is.
711    in_call: Option<String>,
712}
713
714/// What the worker holds: its engine, the guests handed to it, and the session-long instances.
715struct ServerState {
716    engine: wasmtime::Engine,
717    shared: Arc<WorkerShared>,
718    guests: Mutex<HashMap<u64, Component>>,
719    hooks: Mutex<HashMap<u64, Arc<tokio::sync::Mutex<WorkerHook>>>>,
720    drivers: Mutex<HashMap<u64, Arc<tokio::sync::Mutex<WorkerDriver>>>>,
721}
722
723/// A process driver instance in the driver worker, with the latch its wall-clock bounds record
724/// into, drained after each call.
725type WorkerDriver = (ProcessDriver, CallDeadlines);
726
727impl ServerState {
728    fn guest(&self, guest: u64) -> Option<Component> {
729        self.guests
730            .lock()
731            .unwrap_or_else(std::sync::PoisonError::into_inner)
732            .get(&guest)
733            .cloned()
734    }
735}
capsule_runtime/guest_worker/worker.rs

capsule_runtime/guest_worker/
worker.rs
