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}