1use fission_core::{
2 ActionEnvelope, BoxFuture, CapabilityCtx, CapabilityType, DataStreamRegistry, JobCtx, JobRef,
3 JobSpec, OperationCapability, ResourceExecutionContext, ServiceCtx, ServiceRunner, ServiceSpec,
4 ServiceType,
5};
6use std::collections::HashMap;
7use std::future::Future;
8use std::sync::{mpsc, Arc};
9use tokio::runtime::{Builder as TokioRuntimeBuilder, Runtime as TokioRuntime};
10
11pub type WakeFn = Arc<dyn Fn() + Send + Sync>;
12
13#[derive(Clone, Debug)]
14pub enum AsyncMessage {
15 JobOk {
16 job_name: String,
17 req_id: u64,
18 payload: Vec<u8>,
19 on_ok: Option<ActionEnvelope>,
20 resource: Option<ResourceExecutionContext>,
21 },
22 JobErr {
23 job_name: String,
24 req_id: u64,
25 payload: Option<Vec<u8>>,
26 on_err: Option<ActionEnvelope>,
27 message: Option<String>,
28 resource: Option<ResourceExecutionContext>,
29 },
30 ServiceStarted {
31 service_name: String,
32 slot_key: String,
33 instance_id: u64,
34 resource: Option<ResourceExecutionContext>,
35 },
36 ServiceStartFailed {
37 service_name: String,
38 slot_key: String,
39 instance_id: u64,
40 payload: Option<Vec<u8>>,
41 message: Option<String>,
42 resource: Option<ResourceExecutionContext>,
43 },
44 ServiceEvent {
45 service_name: String,
46 slot_key: String,
47 instance_id: u64,
48 payload: Vec<u8>,
49 resource: Option<ResourceExecutionContext>,
50 },
51 ServiceStopped {
52 service_name: String,
53 slot_key: String,
54 instance_id: u64,
55 resource: Option<ResourceExecutionContext>,
56 },
57 ServiceCommandOk {
58 service_name: String,
59 slot_key: String,
60 instance_id: u64,
61 req_id: u64,
62 payload: Option<Vec<u8>>,
63 on_ok: Option<ActionEnvelope>,
64 resource: Option<ResourceExecutionContext>,
65 },
66 ServiceCommandErr {
67 service_name: String,
68 slot_key: String,
69 instance_id: u64,
70 req_id: u64,
71 payload: Option<Vec<u8>>,
72 on_err: Option<ActionEnvelope>,
73 message: Option<String>,
74 resource: Option<ResourceExecutionContext>,
75 },
76 CapabilityOk {
77 capability_name: String,
78 req_id: u64,
79 payload: Vec<u8>,
80 on_ok: Option<ActionEnvelope>,
81 resource: Option<ResourceExecutionContext>,
82 },
83 CapabilityErr {
84 capability_name: String,
85 req_id: u64,
86 payload: Option<Vec<u8>>,
87 on_err: Option<ActionEnvelope>,
88 message: Option<String>,
89 resource: Option<ResourceExecutionContext>,
90 },
91}
92
93#[derive(Clone)]
94pub enum ServiceControlMessage {
95 Command {
96 req_id: u64,
97 payload: Vec<u8>,
98 on_ok: Option<ActionEnvelope>,
99 on_err: Option<ActionEnvelope>,
100 },
101 Stop,
102}
103
104#[derive(Clone)]
105pub struct RunningServiceHandle {
106 pub instance_id: u64,
107 pub control_tx: mpsc::Sender<ServiceControlMessage>,
108}
109
110#[derive(Clone)]
111struct JobLaunch {
112 req_id: u64,
113 payload: Vec<u8>,
114 on_ok: Option<ActionEnvelope>,
115 on_err: Option<ActionEnvelope>,
116 resource: Option<ResourceExecutionContext>,
117 tx: mpsc::Sender<AsyncMessage>,
118 wake: WakeFn,
119}
120
121#[derive(Clone)]
122struct CapabilityLaunch {
123 req_id: u64,
124 payload: Vec<u8>,
125 on_ok: Option<ActionEnvelope>,
126 on_err: Option<ActionEnvelope>,
127 resource: Option<ResourceExecutionContext>,
128 tx: mpsc::Sender<AsyncMessage>,
129 wake: WakeFn,
130}
131
132#[derive(Clone)]
133struct ServiceLaunch {
134 service_name: String,
135 slot_key: String,
136 instance_id: u64,
137 config: Vec<u8>,
138 resource: Option<ResourceExecutionContext>,
139 tx: mpsc::Sender<AsyncMessage>,
140 wake: WakeFn,
141}
142
143type JobHandler = dyn Fn(JobLaunch) + Send + Sync;
144type ServiceSpawner = dyn Fn(ServiceLaunch) -> RunningServiceHandle + Send + Sync;
145type CapabilitySpawner = dyn Fn(CapabilityLaunch) + Send + Sync;
146
147pub struct AsyncRegistry {
148 jobs: HashMap<String, Arc<JobHandler>>,
149 services: HashMap<String, Arc<ServiceSpawner>>,
150 operations: HashMap<String, Arc<CapabilitySpawner>>,
151 data_streams: DataStreamRegistry,
152}
153
154impl Default for AsyncRegistry {
155 fn default() -> Self {
156 Self {
157 jobs: HashMap::new(),
158 services: HashMap::new(),
159 operations: HashMap::new(),
160 data_streams: DataStreamRegistry::new(),
161 }
162 }
163}
164
165impl AsyncRegistry {
166 pub fn new() -> Self {
167 Self::default()
168 }
169
170 pub fn register_operation_capability<C, F, Fut>(
171 &mut self,
172 capability: CapabilityType<C>,
173 handler: F,
174 ) where
175 C: OperationCapability,
176 F: Fn(C::Request, CapabilityCtx) -> Fut + Send + Sync + 'static,
177 Fut: Future<Output = Result<C::Ok, C::Err>> + Send + 'static,
178 {
179 let handler = Arc::new(handler);
180 let data_streams = self.data_streams.clone();
181 self.operations.insert(
182 capability.name.to_string(),
183 Arc::new(move |launch: CapabilityLaunch| {
184 let handler = handler.clone();
185 let name = capability.name.to_string();
186 let data_streams = data_streams.clone();
187 std::thread::spawn(move || {
188 let runtime = match new_job_runtime() {
189 Ok(runtime) => runtime,
190 Err(err) => {
191 let _ = launch.tx.send(AsyncMessage::CapabilityErr {
192 capability_name: name,
193 req_id: launch.req_id,
194 payload: None,
195 on_err: launch.on_err,
196 message: Some(err),
197 resource: launch.resource,
198 });
199 (launch.wake)();
200 return;
201 }
202 };
203
204 let request = match serde_json::from_slice::<C::Request>(&launch.payload) {
205 Ok(request) => request,
206 Err(err) => {
207 let _ = launch.tx.send(AsyncMessage::CapabilityErr {
208 capability_name: name,
209 req_id: launch.req_id,
210 payload: None,
211 on_err: launch.on_err,
212 message: Some(err.to_string()),
213 resource: launch.resource,
214 });
215 (launch.wake)();
216 return;
217 }
218 };
219
220 match runtime.block_on(handler(
221 request,
222 CapabilityCtx::new_runtime(launch.req_id, data_streams),
223 )) {
224 Ok(ok) => match serde_json::to_vec(&ok) {
225 Ok(payload) => {
226 let _ = launch.tx.send(AsyncMessage::CapabilityOk {
227 capability_name: name,
228 req_id: launch.req_id,
229 payload,
230 on_ok: launch.on_ok,
231 resource: launch.resource,
232 });
233 }
234 Err(err) => {
235 let _ = launch.tx.send(AsyncMessage::CapabilityErr {
236 capability_name: name,
237 req_id: launch.req_id,
238 payload: None,
239 on_err: launch.on_err,
240 message: Some(err.to_string()),
241 resource: launch.resource,
242 });
243 }
244 },
245 Err(err) => {
246 let (payload, message) = serde_json::to_vec(&err)
247 .ok()
248 .map(|payload| (Some(payload), None))
249 .unwrap_or_else(|| {
250 (None, Some("capability error serialization failed".into()))
251 });
252 let _ = launch.tx.send(AsyncMessage::CapabilityErr {
253 capability_name: name,
254 req_id: launch.req_id,
255 payload,
256 on_err: launch.on_err,
257 message,
258 resource: launch.resource,
259 });
260 }
261 }
262
263 (launch.wake)();
264 });
265 }),
266 );
267 }
268
269 pub fn register_job<J, F, Fut>(&mut self, job: JobRef<J>, handler: F)
270 where
271 J: JobSpec,
272 F: Fn(J::Request, JobCtx) -> Fut + Send + Sync + 'static,
273 Fut: Future<Output = Result<J::Ok, J::Err>> + Send + 'static,
274 {
275 let handler = Arc::new(handler);
276 let data_streams = self.data_streams.clone();
277 self.jobs.insert(
278 job.name.to_string(),
279 Arc::new(move |launch: JobLaunch| {
280 let handler = handler.clone();
281 let data_streams = data_streams.clone();
282 std::thread::spawn(move || {
283 let runtime = match new_job_runtime() {
284 Ok(runtime) => runtime,
285 Err(err) => {
286 let _ = launch.tx.send(AsyncMessage::JobErr {
287 job_name: J::NAME.to_string(),
288 req_id: launch.req_id,
289 payload: None,
290 on_err: launch.on_err,
291 message: Some(err),
292 resource: launch.resource,
293 });
294 (launch.wake)();
295 return;
296 }
297 };
298 let request = match serde_json::from_slice::<J::Request>(&launch.payload) {
299 Ok(request) => request,
300 Err(err) => {
301 let _ = launch.tx.send(AsyncMessage::JobErr {
302 job_name: J::NAME.to_string(),
303 req_id: launch.req_id,
304 payload: None,
305 on_err: launch.on_err,
306 message: Some(err.to_string()),
307 resource: launch.resource,
308 });
309 (launch.wake)();
310 return;
311 }
312 };
313
314 match runtime.block_on(handler(
315 request,
316 JobCtx::new_runtime(launch.req_id, data_streams),
317 )) {
318 Ok(ok) => match serde_json::to_vec(&ok) {
319 Ok(payload) => {
320 let _ = launch.tx.send(AsyncMessage::JobOk {
321 job_name: J::NAME.to_string(),
322 req_id: launch.req_id,
323 payload,
324 on_ok: launch.on_ok,
325 resource: launch.resource,
326 });
327 }
328 Err(err) => {
329 let _ = launch.tx.send(AsyncMessage::JobErr {
330 job_name: J::NAME.to_string(),
331 req_id: launch.req_id,
332 payload: None,
333 on_err: launch.on_err,
334 message: Some(err.to_string()),
335 resource: launch.resource,
336 });
337 }
338 },
339 Err(err) => {
340 let (payload, message) = serde_json::to_vec(&err)
341 .ok()
342 .map(|payload| (Some(payload), None))
343 .unwrap_or_else(|| {
344 (None, Some("job error serialization failed".into()))
345 });
346 let _ = launch.tx.send(AsyncMessage::JobErr {
347 job_name: J::NAME.to_string(),
348 req_id: launch.req_id,
349 payload,
350 on_err: launch.on_err,
351 message,
352 resource: launch.resource,
353 });
354 }
355 }
356
357 (launch.wake)();
358 });
359 }),
360 );
361 }
362
363 pub fn register_service<S, F, Fut>(&mut self, service: ServiceType<S>, starter: F)
364 where
365 S: ServiceSpec + 'static,
366 F: Fn(S::Config, ServiceCtx<S>) -> Fut + Send + Sync + 'static,
367 Fut: Future<Output = Result<Box<dyn ServiceRunner<S>>, S::StartErr>> + Send + 'static,
368 {
369 let starter = Arc::new(starter);
370 let data_streams = self.data_streams.clone();
371 self.services.insert(
372 service.name.to_string(),
373 Arc::new(move |launch: ServiceLaunch| {
374 let (control_tx, control_rx) = mpsc::channel();
375 let starter = starter.clone();
376 let data_streams = data_streams.clone();
377 let tx = launch.tx.clone();
378 let wake = launch.wake.clone();
379 let service_name = launch.service_name.clone();
380 let slot_key = launch.slot_key.clone();
381 let resource = launch.resource.clone();
382 let instance_id = launch.instance_id;
383 let config_bytes = launch.config.clone();
384
385 std::thread::spawn(move || {
386 let runtime = match new_service_runtime() {
387 Ok(runtime) => runtime,
388 Err(err) => {
389 let _ = tx.send(AsyncMessage::ServiceStartFailed {
390 service_name,
391 slot_key,
392 instance_id,
393 payload: None,
394 message: Some(err),
395 resource,
396 });
397 wake();
398 return;
399 }
400 };
401 let tx_for_emit = tx.clone();
402 let wake_for_emit = wake.clone();
403 let service_name_for_emit = service_name.clone();
404 let slot_key_for_emit = slot_key.clone();
405 let resource_for_emit = resource.clone();
406 let emit = Arc::new(move |payload: Vec<u8>| -> BoxFuture<Result<(), String>> {
407 let tx = tx_for_emit.clone();
408 let wake = wake_for_emit.clone();
409 let service_name = service_name_for_emit.clone();
410 let slot_key = slot_key_for_emit.clone();
411 let resource = resource_for_emit.clone();
412 Box::pin(async move {
413 tx.send(AsyncMessage::ServiceEvent {
414 service_name,
415 slot_key,
416 instance_id,
417 payload,
418 resource,
419 })
420 .map_err(|err| err.to_string())?;
421 wake();
422 Ok(())
423 })
424 });
425
426 let ctx = ServiceCtx::<S>::new_runtime(
427 service_name.clone(),
428 slot_key.clone(),
429 instance_id,
430 data_streams,
431 emit,
432 );
433
434 let config = match serde_json::from_slice::<S::Config>(&config_bytes) {
435 Ok(config) => config,
436 Err(err) => {
437 let _ = tx.send(AsyncMessage::ServiceStartFailed {
438 service_name,
439 slot_key,
440 instance_id,
441 payload: None,
442 message: Some(err.to_string()),
443 resource,
444 });
445 wake();
446 return;
447 }
448 };
449
450 let mut runner = match runtime.block_on(starter(config, ctx.clone())) {
451 Ok(runner) => {
452 let _ = tx.send(AsyncMessage::ServiceStarted {
453 service_name: service_name.clone(),
454 slot_key: slot_key.clone(),
455 instance_id,
456 resource: resource.clone(),
457 });
458 wake();
459 runner
460 }
461 Err(err) => {
462 let payload = serde_json::to_vec(&err).ok();
463 let _ = tx.send(AsyncMessage::ServiceStartFailed {
464 service_name,
465 slot_key,
466 instance_id,
467 payload,
468 message: None,
469 resource,
470 });
471 wake();
472 return;
473 }
474 };
475
476 while let Ok(message) = control_rx.recv() {
477 match message {
478 ServiceControlMessage::Command {
479 req_id,
480 payload,
481 on_ok,
482 on_err,
483 } => {
484 let command = match serde_json::from_slice::<S::Command>(&payload) {
485 Ok(command) => command,
486 Err(err) => {
487 let _ = tx.send(AsyncMessage::ServiceCommandErr {
488 service_name: service_name.clone(),
489 slot_key: slot_key.clone(),
490 instance_id,
491 req_id,
492 payload: None,
493 on_err,
494 message: Some(err.to_string()),
495 resource: resource.clone(),
496 });
497 wake();
498 continue;
499 }
500 };
501
502 match runtime.block_on(runner.on_command(command, ctx.clone())) {
503 Ok(ok) => {
504 let payload = serde_json::to_vec(&ok).ok();
505 let _ = tx.send(AsyncMessage::ServiceCommandOk {
506 service_name: service_name.clone(),
507 slot_key: slot_key.clone(),
508 instance_id,
509 req_id,
510 payload,
511 on_ok,
512 resource: resource.clone(),
513 });
514 }
515 Err(err) => {
516 let (payload, message) = serde_json::to_vec(&err)
517 .ok()
518 .map(|payload| (Some(payload), None))
519 .unwrap_or_else(|| {
520 (
521 None,
522 Some(
523 "service command error serialization failed"
524 .into(),
525 ),
526 )
527 });
528 let _ = tx.send(AsyncMessage::ServiceCommandErr {
529 service_name: service_name.clone(),
530 slot_key: slot_key.clone(),
531 instance_id,
532 req_id,
533 payload,
534 on_err,
535 message,
536 resource: resource.clone(),
537 });
538 }
539 }
540 wake();
541 }
542 ServiceControlMessage::Stop => {
543 runtime.block_on(runner.on_stop(ctx.clone()));
544 let _ = tx.send(AsyncMessage::ServiceStopped {
545 service_name: service_name.clone(),
546 slot_key: slot_key.clone(),
547 instance_id,
548 resource: resource.clone(),
549 });
550 wake();
551 return;
552 }
553 }
554 }
555
556 runtime.block_on(runner.on_stop(ctx));
557 let _ = tx.send(AsyncMessage::ServiceStopped {
558 service_name,
559 slot_key,
560 instance_id,
561 resource,
562 });
563 wake();
564 });
565
566 RunningServiceHandle {
567 instance_id: launch.instance_id,
568 control_tx,
569 }
570 }),
571 );
572 }
573
574 pub fn spawn_job(
575 &self,
576 job_name: &str,
577 req_id: u64,
578 payload: Vec<u8>,
579 on_ok: Option<ActionEnvelope>,
580 on_err: Option<ActionEnvelope>,
581 resource: Option<ResourceExecutionContext>,
582 tx: &mpsc::Sender<AsyncMessage>,
583 wake: WakeFn,
584 ) -> bool {
585 let Some(handler) = self.jobs.get(job_name) else {
586 return false;
587 };
588 handler(JobLaunch {
589 req_id,
590 payload,
591 on_ok,
592 on_err,
593 resource,
594 tx: tx.clone(),
595 wake,
596 });
597 true
598 }
599
600 pub fn spawn_capability(
601 &self,
602 capability_name: &str,
603 req_id: u64,
604 payload: Vec<u8>,
605 on_ok: Option<ActionEnvelope>,
606 on_err: Option<ActionEnvelope>,
607 resource: Option<ResourceExecutionContext>,
608 tx: &mpsc::Sender<AsyncMessage>,
609 wake: WakeFn,
610 ) -> bool {
611 let Some(handler) = self.operations.get(capability_name) else {
612 return false;
613 };
614 handler(CapabilityLaunch {
615 req_id,
616 payload,
617 on_ok,
618 on_err,
619 resource,
620 tx: tx.clone(),
621 wake,
622 });
623 true
624 }
625
626 pub fn spawn_service(
627 &self,
628 service_name: &str,
629 slot_key: &str,
630 instance_id: u64,
631 config: Vec<u8>,
632 resource: Option<ResourceExecutionContext>,
633 tx: &mpsc::Sender<AsyncMessage>,
634 wake: WakeFn,
635 ) -> Option<RunningServiceHandle> {
636 let spawner = self.services.get(service_name)?;
637 Some(spawner(ServiceLaunch {
638 service_name: service_name.to_string(),
639 slot_key: slot_key.to_string(),
640 instance_id,
641 config,
642 resource,
643 tx: tx.clone(),
644 wake,
645 }))
646 }
647}
648
649fn new_job_runtime() -> Result<TokioRuntime, String> {
650 TokioRuntimeBuilder::new_current_thread()
651 .enable_all()
652 .build()
653 .map_err(|err| err.to_string())
654}
655
656fn new_service_runtime() -> Result<TokioRuntime, String> {
657 TokioRuntimeBuilder::new_multi_thread()
658 .worker_threads(1)
659 .enable_all()
660 .build()
661 .map_err(|err| err.to_string())
662}