a3s_flow/worker/
routing.rs

1use async_trait::async_trait;
2use std::collections::BTreeMap;
3use std::fmt;
4use std::sync::Arc;
5
6use crate::engine::FlowEngine;
7use crate::error::{FlowError, Result};
8use crate::runtime_build::RuntimeBuildId;
9
10use super::{FlowTask, FlowTaskDispatcher};
11
12/// Routes pinned Flow tasks to dispatchers serving exact runtime builds.
13///
14/// A route can point at an A3S Boot manager, a compatibility queue, or another
15/// host dispatcher. Register the same dispatcher under every build it can
16/// execute. Unpinned histories require a separate explicit route.
17#[derive(Clone, Default)]
18pub struct RuntimeBuildTaskRouter {
19    routes: BTreeMap<RuntimeBuildId, Arc<dyn FlowTaskDispatcher>>,
20    unpinned_route: Option<Arc<dyn FlowTaskDispatcher>>,
21}
22
23impl fmt::Debug for RuntimeBuildTaskRouter {
24    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
25        formatter
26            .debug_struct("RuntimeBuildTaskRouter")
27            .field("runtime_build_ids", &self.routes.keys().collect::<Vec<_>>())
28            .field("has_unpinned_route", &self.unpinned_route.is_some())
29            .finish()
30    }
31}
32
33impl RuntimeBuildTaskRouter {
34    /// Create a router with no pinned or legacy routes.
35    pub fn new() -> Self {
36        Self::default()
37    }
38
39    /// Register the only dispatcher used for an exact runtime build.
40    pub fn with_route(
41        mut self,
42        runtime_build_id: RuntimeBuildId,
43        dispatcher: Arc<dyn FlowTaskDispatcher>,
44    ) -> Result<Self> {
45        if self.routes.contains_key(&runtime_build_id) {
46            return Err(FlowError::InvalidWorkerConfiguration(format!(
47                "runtime build route {runtime_build_id} is already registered"
48            )));
49        }
50        self.routes.insert(runtime_build_id, dispatcher);
51        Ok(self)
52    }
53
54    /// Register the explicit dispatcher for legacy unpinned histories.
55    pub fn with_unpinned_route(mut self, dispatcher: Arc<dyn FlowTaskDispatcher>) -> Result<Self> {
56        if self.unpinned_route.is_some() {
57            return Err(FlowError::InvalidWorkerConfiguration(
58                "an unpinned runtime build route is already registered".to_string(),
59            ));
60        }
61        self.unpinned_route = Some(dispatcher);
62        Ok(self)
63    }
64
65    /// Iterate over the exact pinned build routes.
66    pub fn runtime_build_ids(&self) -> impl Iterator<Item = &RuntimeBuildId> {
67        self.routes.keys()
68    }
69
70    /// Return whether a legacy unpinned route is registered.
71    pub fn has_unpinned_route(&self) -> bool {
72        self.unpinned_route.is_some()
73    }
74
75    /// Resolve the persisted build for `run_id`, verify the task target, and
76    /// dispatch it through the matching route.
77    pub async fn dispatch_for_run(
78        &self,
79        engine: &FlowEngine,
80        run_id: &str,
81        task: FlowTask,
82    ) -> Result<()> {
83        let task_run_id = task.target_run_id().ok_or_else(|| {
84            FlowError::InvalidTransition(
85                "runtime build routing requires a Flow task with an explicit run id".to_string(),
86            )
87        })?;
88        if task_run_id != run_id {
89            return Err(FlowError::InvalidTransition(format!(
90                "Flow task targets run {task_run_id} but build routing requested {run_id}"
91            )));
92        }
93        let required_build_id = engine.runtime_build_id(run_id).await?;
94        self.dispatch_for_runtime_build(required_build_id.as_ref(), task)
95            .await
96    }
97
98    fn route(
99        &self,
100        required_build_id: Option<&RuntimeBuildId>,
101    ) -> Result<&Arc<dyn FlowTaskDispatcher>> {
102        match required_build_id {
103            Some(build_id) => self.routes.get(build_id),
104            None => self.unpinned_route.as_ref(),
105        }
106        .ok_or_else(|| FlowError::RuntimeBuildRouteNotFound {
107            required_build_id: required_build_id.cloned(),
108        })
109    }
110}
111
112#[async_trait]
113impl FlowTaskDispatcher for RuntimeBuildTaskRouter {
114    async fn dispatch(&self, task: FlowTask) -> Result<()> {
115        self.route(None)?.dispatch(task).await
116    }
117
118    fn has_runtime_build_route(&self, required_build_id: Option<&RuntimeBuildId>) -> bool {
119        match required_build_id {
120            Some(build_id) => self.routes.contains_key(build_id),
121            None => self.unpinned_route.is_some(),
122        }
123    }
124
125    async fn dispatch_for_runtime_build(
126        &self,
127        required_build_id: Option<&RuntimeBuildId>,
128        task: FlowTask,
129    ) -> Result<()> {
130        self.route(required_build_id)?.dispatch(task).await
131    }
132}
a3s_flow/worker/routing.rs

a3s_flow/worker/
routing.rs
