Skip to main content

little_durable_objects/actor/
executor_connection.rs

1use std::{
2    collections::{HashMap, HashSet},
3    error::Error,
4    fmt::{Display, Formatter},
5    os::unix::fs::FileTypeExt,
6    path::{Path, PathBuf},
7    sync::{
8        Arc, Mutex,
9        atomic::{AtomicU64, Ordering},
10    },
11};
12
13use anyhow::{Context, Result, ensure};
14use async_trait::async_trait;
15use serde::{Deserialize, Serialize};
16use serde_json::Value;
17use tokio::{
18    io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
19    net::{UnixListener, unix::OwnedWriteHalf},
20    sync::{Mutex as AsyncMutex, oneshot},
21    task::JoinHandle,
22};
23use tokio_util::sync::CancellationToken;
24use tracing::{debug, info};
25
26use super::{ActorInvocationFailure, ActorKey};
27
28const ACTOR_EXECUTOR_PROTOCOL_VERSION: u32 = 11;
29pub(crate) const MAX_ACTOR_EXECUTOR_MESSAGE_BYTES: usize = 16 * 1024 * 1024;
30
31#[derive(Debug, Serialize)]
32pub struct ActorMethodInvocation {
33    pub request_id: String,
34    pub actor: ActorKey,
35    pub method: String,
36    pub args: Vec<Value>,
37    pub state: Option<Value>,
38}
39
40#[derive(Debug, Serialize)]
41pub struct ActorMethodEviction {
42    pub actor: ActorKey,
43}
44
45#[derive(Debug, PartialEq)]
46pub enum ActorMethodOutcome {
47    Completed { result: Value, state: Value },
48    Failed(ActorInvocationFailure),
49}
50
51#[async_trait]
52pub trait ActorExecutor: Send + Sync {
53    fn supports(&self, actor_type: &str) -> bool;
54
55    async fn invoke(&self, invocation: ActorMethodInvocation) -> Result<ActorMethodOutcome>;
56
57    async fn evict(&self, _eviction: ActorMethodEviction) -> Result<()> {
58        Ok(())
59    }
60}
61
62pub(crate) struct ActorExecutorListener {
63    listener: UnixListener,
64    socket_path: PathBuf,
65}
66
67impl ActorExecutorListener {
68    pub(crate) async fn bind(socket_path: impl Into<PathBuf>) -> Result<Self> {
69        let socket_path = socket_path.into();
70        prepare_socket_path(&socket_path).await?;
71        if let Some(parent) = socket_path.parent() {
72            tokio::fs::create_dir_all(parent)
73                .await
74                .with_context(|| format!("create actor executor directory {}", parent.display()))?;
75        }
76        let listener = UnixListener::bind(&socket_path)
77            .with_context(|| format!("bind actor executor socket {}", socket_path.display()))?;
78        Ok(Self {
79            listener,
80            socket_path,
81        })
82    }
83
84    pub(crate) async fn accept(self) -> Result<ActorExecutorConnection> {
85        let result = self.accept_connection().await;
86        let cleanup = remove_socket(&self.socket_path).await;
87        match (result, cleanup) {
88            (Ok(connection), Ok(())) => Ok(connection),
89            (Err(error), _) => Err(error),
90            (Ok(_), Err(error)) => Err(error),
91        }
92    }
93
94    async fn accept_connection(&self) -> Result<ActorExecutorConnection> {
95        let (stream, _) =
96            self.listener.accept().await.with_context(|| {
97                format!("accept actor executor at {}", self.socket_path.display())
98            })?;
99        let (reader, writer) = stream.into_split();
100        let mut reader = BufReader::new(reader);
101        let attach = match read_client_message(&mut reader).await? {
102            Some(ActorExecutorClientMessage::Attach {
103                protocol,
104                actor_types,
105            }) => {
106                ensure!(
107                    protocol == ACTOR_EXECUTOR_PROTOCOL_VERSION,
108                    "customer actor executor uses unsupported protocol version {protocol}"
109                );
110                ensure!(
111                    !actor_types.is_empty(),
112                    "customer actor executor did not advertise any actor types"
113                );
114                actor_types
115            }
116            Some(_) => {
117                anyhow::bail!("first customer actor executor message must attach the process")
118            }
119            None => anyhow::bail!("customer actor executor disconnected before attaching"),
120        };
121
122        let executor = Arc::new(JsActorExecutor::new(writer, attach));
123        let task = tokio::spawn(read_executor_messages(reader, executor.clone()));
124        debug!(
125            socket = %self.socket_path.display(),
126            actor_types = ?executor.actor_types,
127            "customer JavaScript process connected to actor executor"
128        );
129        Ok(ActorExecutorConnection { executor, task })
130    }
131}
132
133pub(crate) struct ActorExecutorConnection {
134    executor: Arc<JsActorExecutor>,
135    task: JoinHandle<Result<()>>,
136}
137
138impl ActorExecutorConnection {
139    pub(crate) fn executor(&self) -> Arc<dyn ActorExecutor> {
140        self.executor.clone()
141    }
142
143    pub(crate) async fn mark_ready(&self) -> Result<()> {
144        self.executor
145            .send(&ActorExecutorServerMessage::Attached {
146                protocol: ACTOR_EXECUTOR_PROTOCOL_VERSION,
147            })
148            .await?;
149        info!(
150            actor_types = ?self.executor.actor_types,
151            "customer JavaScript process attached to actor executor"
152        );
153        Ok(())
154    }
155
156    pub(crate) async fn run(mut self, shutdown: CancellationToken) -> Result<()> {
157        tokio::select! {
158            result = &mut self.task => {
159                match result {
160                    Ok(result) => result,
161                    Err(error) => Err(error.into()),
162                }
163            }
164            _ = shutdown.cancelled() => {
165                self.executor.close().await;
166                self.task.abort();
167                let _ = (&mut self.task).await;
168                Ok(())
169            }
170        }
171    }
172}
173
174impl Drop for ActorExecutorConnection {
175    fn drop(&mut self) {
176        self.task.abort();
177    }
178}
179
180struct JsActorExecutor {
181    actor_types: HashSet<String>,
182    next_message_id: AtomicU64,
183    pending: Mutex<HashMap<u64, oneshot::Sender<ExecutorReply>>>,
184    writer: AsyncMutex<OwnedWriteHalf>,
185}
186
187#[async_trait]
188impl ActorExecutor for JsActorExecutor {
189    fn supports(&self, actor_type: &str) -> bool {
190        self.actor_types.contains(actor_type)
191    }
192
193    async fn invoke(&self, invocation: ActorMethodInvocation) -> Result<ActorMethodOutcome> {
194        match self.exchange(ExecutorCommand::Invoke(invocation)).await? {
195            ExecutorReply::Invoked { result, state } => {
196                Ok(ActorMethodOutcome::Completed { result, state })
197            }
198            ExecutorReply::Failed { code, message } => {
199                Ok(ActorMethodOutcome::Failed(ActorInvocationFailure {
200                    code,
201                    message,
202                }))
203            }
204            ExecutorReply::Evicted => {
205                anyhow::bail!("actor executor returned eviction reply to invocation")
206            }
207        }
208    }
209
210    async fn evict(&self, eviction: ActorMethodEviction) -> Result<()> {
211        match self.exchange(ExecutorCommand::Evict(eviction)).await? {
212            ExecutorReply::Evicted => Ok(()),
213            ExecutorReply::Failed { code, message } => {
214                anyhow::bail!("actor executor rejected eviction ({code}): {message}")
215            }
216            ExecutorReply::Invoked { .. } => {
217                anyhow::bail!("actor executor returned the wrong reply to eviction")
218            }
219        }
220    }
221}
222
223impl JsActorExecutor {
224    fn new(writer: OwnedWriteHalf, actor_types: Vec<String>) -> Self {
225        Self {
226            actor_types: actor_types.into_iter().collect(),
227            next_message_id: AtomicU64::new(1),
228            pending: Mutex::new(HashMap::new()),
229            writer: AsyncMutex::new(writer),
230        }
231    }
232
233    async fn exchange(&self, command: ExecutorCommand) -> Result<ExecutorReply> {
234        let message_id = self.next_message_id.fetch_add(1, Ordering::Relaxed);
235        let (reply_tx, reply_rx) = oneshot::channel();
236        self.pending
237            .lock()
238            .map_err(|_| anyhow::anyhow!("actor executor pending-reply lock poisoned"))?
239            .insert(message_id, reply_tx);
240
241        let write_result = self
242            .send(&ActorExecutorServerMessage::Command {
243                message_id,
244                command: Box::new(command),
245            })
246            .await;
247        if let Err(error) = write_result {
248            self.remove_pending(message_id)?;
249            if let Some(limit) = error.downcast_ref::<ActorExecutorMessageTooLarge>() {
250                return Ok(ExecutorReply::Failed {
251                    code: "resource_exhausted".into(),
252                    message: limit.to_string(),
253                });
254            }
255            return Err(error.context("send command to customer actor executor"));
256        }
257
258        reply_rx
259            .await
260            .context("customer actor executor disconnected before replying")
261    }
262
263    async fn send(&self, message: &ActorExecutorServerMessage) -> Result<()> {
264        write_server_message(&mut *self.writer.lock().await, message).await
265    }
266
267    fn deliver(&self, message_id: u64, reply: ExecutorReply) -> Result<()> {
268        let sender = self
269            .pending
270            .lock()
271            .map_err(|_| anyhow::anyhow!("actor executor pending-reply lock poisoned"))?
272            .remove(&message_id)
273            .with_context(|| format!("actor executor replied to unknown message {message_id}"))?;
274        let _ = sender.send(reply);
275        Ok(())
276    }
277
278    fn remove_pending(&self, message_id: u64) -> Result<()> {
279        self.pending
280            .lock()
281            .map_err(|_| anyhow::anyhow!("actor executor pending-reply lock poisoned"))?
282            .remove(&message_id);
283        Ok(())
284    }
285
286    async fn close(&self) {
287        let _ = self.writer.lock().await.shutdown().await;
288        self.disconnect();
289    }
290
291    fn disconnect(&self) {
292        if let Ok(mut pending) = self.pending.lock() {
293            pending.clear();
294        }
295    }
296}
297
298async fn read_executor_messages(
299    mut reader: BufReader<tokio::net::unix::OwnedReadHalf>,
300    executor: Arc<JsActorExecutor>,
301) -> Result<()> {
302    let result = async {
303        loop {
304            match read_client_message(&mut reader).await? {
305                None => {
306                    anyhow::bail!("customer JavaScript actor executor disconnected")
307                }
308                Some(ActorExecutorClientMessage::Reply { message_id, reply }) => {
309                    executor.deliver(message_id, reply)?;
310                }
311                Some(ActorExecutorClientMessage::Attach { .. }) => {
312                    anyhow::bail!("customer actor executor attached more than once")
313                }
314            }
315        }
316    }
317    .await;
318    executor.disconnect();
319    result
320}
321
322async fn read_client_message(
323    reader: &mut BufReader<tokio::net::unix::OwnedReadHalf>,
324) -> Result<Option<ActorExecutorClientMessage>> {
325    let mut line = String::new();
326    let bytes = reader.read_line(&mut line).await?;
327    if bytes == 0 {
328        return Ok(None);
329    }
330    ensure!(
331        bytes <= MAX_ACTOR_EXECUTOR_MESSAGE_BYTES,
332        "customer actor executor message exceeds {MAX_ACTOR_EXECUTOR_MESSAGE_BYTES} bytes"
333    );
334    serde_json::from_str(line.trim_end())
335        .map(Some)
336        .context("decode customer actor executor message")
337}
338
339async fn write_server_message(
340    writer: &mut OwnedWriteHalf,
341    message: &ActorExecutorServerMessage,
342) -> Result<()> {
343    let mut bytes = serde_json::to_vec(message)?;
344    bytes.push(b'\n');
345    if bytes.len() > MAX_ACTOR_EXECUTOR_MESSAGE_BYTES {
346        return Err(ActorExecutorMessageTooLarge.into());
347    }
348    writer.write_all(&bytes).await?;
349    Ok(())
350}
351
352#[derive(Debug)]
353struct ActorExecutorMessageTooLarge;
354
355impl Display for ActorExecutorMessageTooLarge {
356    fn fmt(&self, formatter: &mut Formatter<'_>) -> std::fmt::Result {
357        write!(
358            formatter,
359            "actor executor command exceeds {MAX_ACTOR_EXECUTOR_MESSAGE_BYTES} bytes"
360        )
361    }
362}
363
364impl Error for ActorExecutorMessageTooLarge {}
365
366async fn prepare_socket_path(path: &Path) -> Result<()> {
367    match tokio::fs::symlink_metadata(path).await {
368        Ok(metadata) => {
369            ensure!(
370                metadata.file_type().is_socket(),
371                "refusing to replace non-socket actor executor path {}",
372                path.display()
373            );
374            tokio::fs::remove_file(path).await?;
375        }
376        Err(error) if error.kind() == std::io::ErrorKind::NotFound => {}
377        Err(error) => return Err(error.into()),
378    }
379    Ok(())
380}
381
382async fn remove_socket(path: &Path) -> Result<()> {
383    match tokio::fs::remove_file(path).await {
384        Ok(()) => {
385            debug!(socket = %path.display(), "actor executor socket removed");
386            Ok(())
387        }
388        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(()),
389        Err(error) => Err(error.into()),
390    }
391}
392
393#[derive(Debug, Serialize)]
394#[serde(tag = "type", rename_all = "snake_case")]
395enum ActorExecutorServerMessage {
396    Attached {
397        protocol: u32,
398    },
399    Command {
400        message_id: u64,
401        command: Box<ExecutorCommand>,
402    },
403}
404
405#[derive(Debug, Deserialize)]
406#[serde(tag = "type", rename_all = "snake_case")]
407enum ActorExecutorClientMessage {
408    Attach {
409        protocol: u32,
410        actor_types: Vec<String>,
411    },
412    Reply {
413        message_id: u64,
414        reply: ExecutorReply,
415    },
416}
417
418#[derive(Debug, Serialize)]
419#[serde(tag = "type", rename_all = "snake_case")]
420enum ExecutorCommand {
421    Invoke(ActorMethodInvocation),
422    Evict(ActorMethodEviction),
423}
424
425#[derive(Debug, Deserialize)]
426#[serde(tag = "type", rename_all = "snake_case")]
427enum ExecutorReply {
428    Invoked { result: Value, state: Value },
429    Failed { code: String, message: String },
430    Evicted,
431}
432
433#[cfg(test)]
434mod tests {
435    use super::*;
436    use serde_json::json;
437    use tempfile::TempDir;
438    use tokio::{
439        io::{AsyncBufReadExt, AsyncWriteExt, BufReader},
440        net::UnixStream,
441    };
442
443    #[tokio::test]
444    async fn one_javascript_executor_runs_until_host_shutdown() -> Result<()> {
445        let root = TempDir::new_in("/tmp")?;
446        let socket = root.path().join("actor-executor.sock");
447        let host = ActorExecutorListener::bind(&socket).await?;
448        let customer = tokio::spawn(run_incrementing_customer(socket.clone()));
449        let connection = host.accept().await?;
450        let executor = connection.executor();
451        connection.mark_ready().await?;
452        assert!(executor.supports("counter"));
453
454        let shutdown = CancellationToken::new();
455        let connection_task = tokio::spawn(connection.run(shutdown.clone()));
456        let outcome = executor
457            .invoke(ActorMethodInvocation {
458                request_id: "request-1".into(),
459                actor: ActorKey {
460                    namespace_id: "namespace-1".into(),
461                    actor_type: "counter".into(),
462                    actor_id: "counter-1".into(),
463                },
464                method: "increment".into(),
465                args: vec![json!(2)],
466                state: None,
467            })
468            .await?;
469        assert_eq!(
470            outcome,
471            ActorMethodOutcome::Completed {
472                result: json!(2),
473                state: json!({ "count": 2 }),
474            }
475        );
476        shutdown.cancel();
477        connection_task.await??;
478        customer.await??;
479        Ok(())
480    }
481
482    #[tokio::test]
483    async fn oversized_commands_are_reported_as_resource_exhausted() -> Result<()> {
484        let root = TempDir::new_in("/tmp")?;
485        let socket = root.path().join("actor-executor.sock");
486        let host = ActorExecutorListener::bind(&socket).await?;
487        let customer = tokio::spawn(run_attached_customer(socket.clone()));
488        let connection = host.accept().await?;
489        let executor = connection.executor();
490        connection.mark_ready().await?;
491
492        let shutdown = CancellationToken::new();
493        let connection_task = tokio::spawn(connection.run(shutdown.clone()));
494        let outcome = executor
495            .invoke(ActorMethodInvocation {
496                request_id: "request-1".into(),
497                actor: ActorKey {
498                    namespace_id: "namespace-1".into(),
499                    actor_type: "counter".into(),
500                    actor_id: "counter-1".into(),
501                },
502                method: "accept".into(),
503                args: vec![json!("x".repeat(MAX_ACTOR_EXECUTOR_MESSAGE_BYTES))],
504                state: None,
505            })
506            .await?;
507
508        assert!(matches!(
509            outcome,
510            ActorMethodOutcome::Failed(ref failure) if failure.code == "resource_exhausted"
511        ));
512        shutdown.cancel();
513        connection_task.await??;
514        customer.await??;
515        Ok(())
516    }
517
518    async fn run_incrementing_customer(socket: PathBuf) -> Result<()> {
519        let stream = UnixStream::connect(socket).await?;
520        let (reader, mut writer) = stream.into_split();
521        let mut reader = BufReader::new(reader);
522        writer
523            .write_all(b"{\"type\":\"attach\",\"protocol\":11,\"actor_types\":[\"counter\"]}\n")
524            .await?;
525        ensure!(
526            read_json_line(&mut reader).await? == json!({ "type": "attached", "protocol": 11 })
527        );
528
529        let invocation = read_json_line(&mut reader).await?;
530        let invocation_id = invocation["message_id"]
531            .as_u64()
532            .context("invocation message ID")?;
533        ensure!(invocation["command"]["type"] == "invoke");
534        ensure!(invocation["command"].get("timeout_ms").is_none());
535        write_json_line(
536            &mut writer,
537            &json!({
538                "type": "reply",
539                "message_id": invocation_id,
540                "reply": {
541                    "type": "invoked",
542                    "result": 2,
543                    "state": { "count": 2 }
544                }
545            }),
546        )
547        .await?;
548
549        let mut trailing = String::new();
550        ensure!(
551            reader.read_line(&mut trailing).await? == 0,
552            "expected Rust host to close the actor executor"
553        );
554        Ok(())
555    }
556
557    async fn run_attached_customer(socket: PathBuf) -> Result<()> {
558        let stream = UnixStream::connect(socket).await?;
559        let (reader, mut writer) = stream.into_split();
560        let mut reader = BufReader::new(reader);
561        writer
562            .write_all(b"{\"type\":\"attach\",\"protocol\":11,\"actor_types\":[\"counter\"]}\n")
563            .await?;
564        ensure!(
565            read_json_line(&mut reader).await? == json!({ "type": "attached", "protocol": 11 })
566        );
567        let mut trailing = String::new();
568        ensure!(
569            reader.read_line(&mut trailing).await? == 0,
570            "oversized command reached the customer actor executor"
571        );
572        Ok(())
573    }
574
575    async fn read_json_line<R>(reader: &mut R) -> Result<Value>
576    where
577        R: tokio::io::AsyncBufRead + Unpin,
578    {
579        let mut line = String::new();
580        ensure!(reader.read_line(&mut line).await? > 0, "expected JSON line");
581        Ok(serde_json::from_str(line.trim_end())?)
582    }
583
584    async fn write_json_line<W>(writer: &mut W, value: &Value) -> Result<()>
585    where
586        W: tokio::io::AsyncWrite + Unpin,
587    {
588        writer
589            .write_all(serde_json::to_string(value)?.as_bytes())
590            .await?;
591        writer.write_all(b"\n").await?;
592        Ok(())
593    }
594}