1use std::collections::BTreeMap;
9use std::sync::Arc;
10use std::time::Duration;
11
12use beamr::scheduler::Scheduler;
13
14use crate::routing::FieldValue;
15use crate::routing::function::loader::{ContentHash, RoutingFunction};
16mod actor;
17
18#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
20pub struct ConsumerId(String);
21
22impl ConsumerId {
23 #[must_use]
25 pub fn new(id: impl Into<String>) -> Self {
26 Self(id.into())
27 }
28
29 #[must_use]
31 pub fn as_str(&self) -> &str {
32 self.0.as_str()
33 }
34}
35
36#[derive(Clone, Debug, PartialEq, Eq)]
38pub struct ConsumerStateView {
39 pub consumer: ConsumerId,
41 pub current_in_flight: u32,
43 pub max_in_flight: u32,
45 pub buffer_depth: u32,
47 pub affinity_tags: Vec<String>,
49}
50
51impl ConsumerStateView {
52 #[must_use]
54 pub const fn new(
55 consumer: ConsumerId,
56 current_in_flight: u32,
57 max_in_flight: u32,
58 buffer_depth: u32,
59 affinity_tags: Vec<String>,
60 ) -> Self {
61 Self {
62 consumer,
63 current_in_flight,
64 max_in_flight,
65 buffer_depth,
66 affinity_tags,
67 }
68 }
69
70 #[must_use]
72 pub const fn available_capacity(&self) -> u32 {
73 self.max_in_flight.saturating_sub(self.current_in_flight)
74 }
75
76 #[must_use]
78 pub const fn has_capacity(&self) -> bool {
79 self.available_capacity() > 0
80 }
81
82 #[must_use]
84 pub fn has_affinity(&self, tag: &str) -> bool {
85 self.affinity_tags
86 .iter()
87 .any(|advertised| advertised == tag)
88 }
89}
90
91#[derive(Clone, Debug, Default, PartialEq, Eq)]
93pub struct RoutingDecision {
94 selected: Option<ConsumerId>,
95}
96
97impl RoutingDecision {
98 #[must_use]
100 pub const fn select(consumer: ConsumerId) -> Self {
101 Self {
102 selected: Some(consumer),
103 }
104 }
105
106 #[must_use]
108 pub const fn none() -> Self {
109 Self { selected: None }
110 }
111
112 #[must_use]
114 pub const fn selected(&self) -> Option<&ConsumerId> {
115 self.selected.as_ref()
116 }
117
118 #[must_use]
120 pub const fn is_selected(&self) -> bool {
121 self.selected.is_some()
122 }
123}
124
125#[derive(Clone, Debug, Default, PartialEq)]
127pub struct RoutingMessage {
128 fields: BTreeMap<String, FieldValue>,
129}
130
131impl RoutingMessage {
132 #[must_use]
134 pub fn new() -> Self {
135 Self::default()
136 }
137
138 #[must_use]
140 pub fn with(mut self, field: impl Into<String>, value: FieldValue) -> Self {
141 self.fields.insert(field.into(), value);
142 self
143 }
144
145 #[must_use]
147 pub fn get(&self, field: &str) -> Option<&FieldValue> {
148 self.fields.get(field)
149 }
150
151 pub fn fields(&self) -> impl Iterator<Item = (&str, &FieldValue)> {
153 self.fields
154 .iter()
155 .map(|(name, value)| (name.as_str(), value))
156 }
157}
158
159#[derive(Clone, Debug, PartialEq, Eq, thiserror::Error)]
161pub enum FunctionError {
162 #[error("routing function '{0}' panicked during execution")]
164 Crashed(ContentHash),
165 #[error("routing function '{0}' exceeded the supervision timeout")]
167 TimedOut(ContentHash),
168 #[error("routing function execution process could not be started: {0}")]
170 SpawnFailed(String),
171}
172
173#[derive(Clone)]
180pub struct SupervisedExecutor {
181 scheduler: Arc<Scheduler>,
182 timeout: Duration,
183}
184
185impl SupervisedExecutor {
186 pub const DEFAULT_TIMEOUT: Duration = Duration::from_secs(5);
188
189 #[must_use]
191 pub const fn new(scheduler: Arc<Scheduler>, timeout: Duration) -> Self {
192 Self { scheduler, timeout }
193 }
194
195 #[must_use]
197 pub const fn with_default_timeout(scheduler: Arc<Scheduler>) -> Self {
198 Self::new(scheduler, Self::DEFAULT_TIMEOUT)
199 }
200
201 #[must_use]
203 pub fn scheduler(&self) -> Arc<Scheduler> {
204 Arc::clone(&self.scheduler)
205 }
206
207 pub fn execute(
218 &self,
219 function: &RoutingFunction,
220 message: RoutingMessage,
221 consumers: Vec<ConsumerStateView>,
222 ) -> Result<RoutingDecision, FunctionError> {
223 let invocation = actor::BeamrInvocation::new(Arc::clone(&self.scheduler), self.timeout);
224 let hash = function.content_hash();
225 match invocation.execute(function.clone(), message, consumers) {
226 Ok(decision) => Ok(decision),
227 Err(actor::InvocationError::Crashed) => Err(FunctionError::Crashed(hash)),
228 Err(actor::InvocationError::TimedOut(timed_out_hash)) => {
229 Err(FunctionError::TimedOut(timed_out_hash))
230 }
231 Err(actor::InvocationError::SpawnFailed(message)) => {
232 Err(FunctionError::SpawnFailed(message))
233 }
234 }
235 }
236}
237
238impl std::fmt::Debug for SupervisedExecutor {
239 fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
240 formatter
241 .debug_struct("SupervisedExecutor")
242 .field("timeout", &self.timeout)
243 .finish_non_exhaustive()
244 }
245}
246
247#[cfg(test)]
248mod tests {
249 use std::error::Error;
250 use std::sync::{Arc, Mutex, mpsc};
251 use std::thread;
252 use std::time::Duration;
253
254 use super::{
255 ConsumerId, ConsumerStateView, FunctionError, RoutingDecision, RoutingMessage,
256 SupervisedExecutor,
257 };
258 use crate::conversation::ConversationSupervisor;
259 use crate::routing::FieldValue;
260 use crate::routing::function::loader::{ModuleLoader, RoutingModule, RoutingSlot};
261
262 fn consumer(id: &str, current: u32, max: u32, tags: &[&str]) -> ConsumerStateView {
263 ConsumerStateView::new(
264 ConsumerId::new(id),
265 current,
266 max,
267 0,
268 tags.iter().map(|tag| (*tag).to_owned()).collect(),
269 )
270 }
271
272 fn select_first_with_capacity_module(bytecode: &[u8]) -> RoutingModule {
273 RoutingModule::new(bytecode, |_message, consumers| {
274 consumers
275 .iter()
276 .find(|state| state.has_capacity())
277 .map_or_else(RoutingDecision::none, |state| {
278 RoutingDecision::select(state.consumer.clone())
279 })
280 })
281 }
282
283 fn selected_name(decision: &RoutingDecision) -> Option<&str> {
284 decision.selected().map(ConsumerId::as_str)
285 }
286
287 fn supervised_executor() -> Result<(ConversationSupervisor, SupervisedExecutor), Box<dyn Error>>
288 {
289 let supervisor = ConversationSupervisor::new()?;
290 let executor = SupervisedExecutor::with_default_timeout(supervisor.scheduler());
291 Ok((supervisor, executor))
292 }
293
294 #[test]
295 fn execution_returns_decision_using_consumer_state_view() -> Result<(), Box<dyn Error>> {
296 let loader = ModuleLoader::new();
297 let function = loader.load(select_first_with_capacity_module(b"v1"));
298 let (_supervisor, executor) = supervised_executor()?;
299 let consumers = vec![
300 consumer("saturated", 5, 5, &["fast"]),
301 consumer("ready", 1, 4, &["fast"]),
302 ];
303
304 let decision = executor.execute(&function, RoutingMessage::new(), consumers);
305
306 assert!(matches!(decision, Ok(ref outcome) if selected_name(outcome) == Some("ready")));
307 Ok(())
308 }
309
310 #[test]
311 fn message_fields_are_visible_to_routing_function() -> Result<(), Box<dyn Error>> {
312 let loader = ModuleLoader::new();
313 let function = loader.load(RoutingModule::new(
314 b"amount-router",
315 |message, consumers| {
316 let high_value = matches!(
317 message.get("amount"),
318 Some(FieldValue::Integer(amount)) if *amount > 1_000
319 );
320 if high_value {
321 consumers
322 .first()
323 .map_or_else(RoutingDecision::none, |state| {
324 RoutingDecision::select(state.consumer.clone())
325 })
326 } else {
327 RoutingDecision::none()
328 }
329 },
330 ));
331 let (_supervisor, executor) = supervised_executor()?;
332 let message = RoutingMessage::new().with("amount", FieldValue::Integer(5_000));
333
334 let decision = executor.execute(&function, message, vec![consumer("priority", 0, 1, &[])]);
335
336 assert!(matches!(decision, Ok(ref outcome) if selected_name(outcome) == Some("priority")));
337 Ok(())
338 }
339
340 #[test]
341 fn panic_in_function_is_contained_and_other_channels_proceed() -> Result<(), Box<dyn Error>> {
342 let loader = ModuleLoader::new();
343 let crashing = loader.load(RoutingModule::new(b"channel-a", |_message, _consumers| {
344 std::panic::resume_unwind(Box::new(
345 "intentional crash for fault-isolation test".to_owned(),
346 ))
347 }));
348 let healthy = loader.load(select_first_with_capacity_module(b"channel-b"));
349 let (_supervisor, executor) = supervised_executor()?;
350
351 let crashed = executor.execute(&crashing, RoutingMessage::new(), Vec::new());
352 assert_eq!(
353 crashed,
354 Err(FunctionError::Crashed(crashing.content_hash()))
355 );
356
357 let recovered = executor.execute(
358 &healthy,
359 RoutingMessage::new(),
360 vec![consumer("ready", 0, 1, &[])],
361 );
362 assert!(matches!(recovered, Ok(ref outcome) if selected_name(outcome) == Some("ready")));
363 Ok(())
364 }
365
366 #[test]
367 fn repeated_panics_do_not_poison_the_shared_supervisor() -> Result<(), Box<dyn Error>> {
368 let loader = ModuleLoader::new();
376 let crashing = loader.load(RoutingModule::new(b"flaky", |_message, _consumers| {
377 std::panic::resume_unwind(Box::new("repeated intentional crash".to_owned()))
378 }));
379 let healthy = loader.load(select_first_with_capacity_module(b"steady"));
380 let (_supervisor, executor) = supervised_executor()?;
381
382 for _ in 0..16 {
383 let crashed = executor.execute(&crashing, RoutingMessage::new(), Vec::new());
384 assert_eq!(
385 crashed,
386 Err(FunctionError::Crashed(crashing.content_hash()))
387 );
388
389 let served = executor.execute(
390 &healthy,
391 RoutingMessage::new(),
392 vec![consumer("ready", 0, 1, &[])],
393 );
394 assert!(
395 matches!(served, Ok(ref outcome) if selected_name(outcome) == Some("ready")),
396 "scheduler must keep serving healthy invocations after a contained panic"
397 );
398 }
399 Ok(())
400 }
401
402 #[test]
403 fn function_exceeding_timeout_is_terminated_with_error() -> Result<(), Box<dyn Error>> {
404 let supervisor = ConversationSupervisor::new()?;
405 let loader = ModuleLoader::new();
406 let (release_tx, release_rx) = mpsc::channel::<()>();
407 let release = Mutex::new(release_rx);
408 let slow = loader.load(RoutingModule::new(b"slow", move |_message, _consumers| {
409 let released = release.lock().map(|receiver| receiver.recv());
410 assert!(
411 matches!(released, Ok(Err(mpsc::RecvError))),
412 "the blocked function is released only by dropping its sender"
413 );
414 drop(released);
415 RoutingDecision::none()
416 }));
417 let executor = SupervisedExecutor::new(supervisor.scheduler(), Duration::from_millis(20));
418
419 let result = executor.execute(&slow, RoutingMessage::new(), Vec::new());
420
421 assert_eq!(result, Err(FunctionError::TimedOut(slow.content_hash())));
422 drop(release_tx);
423 supervisor.shutdown();
424 Ok(())
425 }
426
427 #[test]
428 fn hot_deploy_does_not_interrupt_in_flight_and_swaps_next_version() -> Result<(), Box<dyn Error>>
429 {
430 let (_supervisor, executor) = supervised_executor()?;
431 let loader = ModuleLoader::new();
432 let (entered_tx, entered_rx) = mpsc::channel();
433 let (release_tx, release_rx) = mpsc::channel();
434 let release = Mutex::new(release_rx);
435
436 let old = loader.load(RoutingModule::new(b"v1", move |_message, _consumers| {
437 assert!(
438 entered_tx.send(()).is_ok(),
439 "the in-flight function must acknowledge entry"
440 );
441 let released = release.lock().map(|receiver| receiver.recv());
442 assert!(
443 matches!(released, Ok(Ok(()))),
444 "the in-flight function must receive its release"
445 );
446 drop(released);
447 RoutingDecision::select(ConsumerId::new("old"))
448 }));
449 let new = loader.load(RoutingModule::new(b"v2", |_message, _consumers| {
450 RoutingDecision::select(ConsumerId::new("new"))
451 }));
452 let old_hash = old.content_hash();
453 let new_hash = new.content_hash();
454
455 let slot = Arc::new(RoutingSlot::new(old));
456 let slot_for_thread = Arc::clone(&slot);
457 let executor_for_thread = executor.clone();
458
459 let in_flight = thread::spawn(move || {
460 let function = slot_for_thread.current();
461 executor_for_thread.execute(&function, RoutingMessage::new(), Vec::new())
462 });
463
464 entered_rx.recv()?;
465
466 slot.deploy(new);
467 assert_eq!(slot.active_hash(), new_hash);
468 assert!(
469 loader.is_loaded(old_hash),
470 "old module must remain loaded while in flight"
471 );
472 assert_eq!(loader.loaded_count(), 2);
473
474 release_tx.send(())?;
475
476 let in_flight_result = in_flight.join();
477 assert!(matches!(
478 in_flight_result,
479 Ok(Ok(ref outcome)) if selected_name(outcome) == Some("old")
480 ));
481
482 let next = executor.execute(&slot.current(), RoutingMessage::new(), Vec::new());
483 assert!(matches!(next, Ok(ref outcome) if selected_name(outcome) == Some("new")));
484 Ok(())
485 }
486}