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