a3s_flow/worker/
boot.rs

1use std::fmt;
2use std::sync::Arc;
3use std::time::Duration;
4
5use a3s_boot::{BootError, Queue, QueueJob, QueueJobOptions, QueueJobReceipt, QueueRetryPolicy};
6use async_trait::async_trait;
7use sha2::{Digest, Sha256};
8
9use crate::engine::FlowEngine;
10use crate::error::{FlowError, Result};
11
12use super::runner::handle_flow_task;
13use super::{FlowTask, FlowTaskDispatcher};
14
15const DEFAULT_FLOW_JOB_NAME: &str = "a3s.flow.task";
16
17/// How a Boot queue coalesces duplicate Flow task targets.
18///
19/// The derived ID excludes scan timestamps and hook payloads. It identifies the
20/// logical Flow target instead: a run, wait, hook, callback token, targeted
21/// scheduled run, or compatibility-wide due scan. IDs are SHA-256 digests, so
22/// callback tokens are not exposed in queue metadata.
23#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
24#[non_exhaustive]
25pub enum BootFlowTaskDeduplication {
26    /// Submit every dispatch as a distinct Boot job.
27    #[default]
28    Disabled,
29    /// Coalesce matching jobs until the current owner becomes terminal.
30    UntilTerminal,
31    /// Coalesce matching jobs until the owner becomes terminal or the TTL
32    /// expires.
33    UntilTerminalOrTtl(Duration),
34}
35
36/// Typed Boot queue policy applied to every task dispatched by one manager.
37///
38/// Caller-assigned job IDs remain per-submission values and are therefore set
39/// through [`BootFlowTaskManager::enqueue_with_options`]. This policy owns the
40/// settings that are safe to share across scheduler dispatches: retry,
41/// execution timeout, stalled-job tolerance, terminal record cleanup, and
42/// logical-target deduplication.
43#[derive(Debug, Clone, PartialEq)]
44pub struct BootFlowTaskPolicy {
45    retry_policy: QueueRetryPolicy,
46    timeout: Option<Duration>,
47    max_stalled_count: u32,
48    remove_on_complete: bool,
49    remove_on_fail: bool,
50    deduplication: BootFlowTaskDeduplication,
51}
52
53impl Default for BootFlowTaskPolicy {
54    fn default() -> Self {
55        Self {
56            retry_policy: QueueRetryPolicy::none(),
57            timeout: None,
58            max_stalled_count: 1,
59            remove_on_complete: false,
60            remove_on_fail: false,
61            deduplication: BootFlowTaskDeduplication::Disabled,
62        }
63    }
64}
65
66impl BootFlowTaskPolicy {
67    /// Creates the default no-retry, no-timeout task policy.
68    pub fn new() -> Self {
69        Self::default()
70    }
71
72    /// Sets the Boot job retry policy.
73    pub fn with_retry_policy(mut self, retry_policy: QueueRetryPolicy) -> Self {
74        self.retry_policy = retry_policy;
75        self
76    }
77
78    /// Sets the maximum execution duration of each Boot job attempt.
79    pub fn with_timeout(mut self, timeout: Duration) -> Self {
80        self.timeout = Some(timeout);
81        self
82    }
83
84    /// Sets how many stalled recoveries a Boot job may survive.
85    pub fn with_max_stalled_count(mut self, max_stalled_count: u32) -> Self {
86        self.max_stalled_count = max_stalled_count;
87        self
88    }
89
90    /// Configures removal of completed Boot job records.
91    pub fn remove_on_complete(mut self, remove: bool) -> Self {
92        self.remove_on_complete = remove;
93        self
94    }
95
96    /// Configures removal of failed Boot job records.
97    pub fn remove_on_fail(mut self, remove: bool) -> Self {
98        self.remove_on_fail = remove;
99        self
100    }
101
102    /// Sets logical-target deduplication behavior.
103    pub fn with_deduplication(mut self, deduplication: BootFlowTaskDeduplication) -> Self {
104        self.deduplication = deduplication;
105        self
106    }
107
108    /// Returns the Boot job retry policy.
109    pub fn retry_policy(&self) -> &QueueRetryPolicy {
110        &self.retry_policy
111    }
112
113    /// Returns the optional per-attempt execution timeout.
114    pub fn timeout(&self) -> Option<Duration> {
115        self.timeout
116    }
117
118    /// Returns the maximum permitted stalled recovery count.
119    pub fn max_stalled_count(&self) -> u32 {
120        self.max_stalled_count
121    }
122
123    /// Returns whether completed Boot job records are removed.
124    pub fn removes_completed_jobs(&self) -> bool {
125        self.remove_on_complete
126    }
127
128    /// Returns whether failed Boot job records are removed.
129    pub fn removes_failed_jobs(&self) -> bool {
130        self.remove_on_fail
131    }
132
133    /// Returns logical-target deduplication behavior.
134    pub fn deduplication(&self) -> BootFlowTaskDeduplication {
135        self.deduplication
136    }
137
138    fn validate(&self) -> Result<()> {
139        if matches!(
140            self.deduplication,
141            BootFlowTaskDeduplication::UntilTerminalOrTtl(ttl) if ttl.is_zero()
142        ) {
143            return Err(FlowError::InvalidWorkerConfiguration(
144                "Boot Flow task deduplication TTL must be greater than zero".to_string(),
145            ));
146        }
147        Ok(())
148    }
149
150    fn job_options_for(&self, job_name: &str, task: &FlowTask) -> QueueJobOptions {
151        let mut options = QueueJobOptions::new()
152            .with_retry_policy(self.retry_policy.clone())
153            .with_max_stalled_count(self.max_stalled_count)
154            .remove_on_complete(self.remove_on_complete)
155            .remove_on_fail(self.remove_on_fail);
156        if let Some(timeout) = self.timeout {
157            options = options.with_timeout(timeout);
158        }
159
160        let ttl = match self.deduplication {
161            BootFlowTaskDeduplication::Disabled => return options,
162            BootFlowTaskDeduplication::UntilTerminal => None,
163            BootFlowTaskDeduplication::UntilTerminalOrTtl(ttl) => Some(ttl),
164        };
165        options = options.with_deduplication_id(flow_task_deduplication_id(job_name, task));
166        if let Some(deduplication) = options.deduplication.as_mut() {
167            deduplication.ttl = ttl;
168            deduplication.keep_last_if_active = flow_task_needs_active_successor(task);
169        }
170        options
171    }
172}
173
174/// A3S Boot-backed task manager for Flow scheduler and callback dispatch.
175///
176/// Boot owns queue processors, worker lifecycle, leasing, job state, and
177/// shutdown. Flow owns only task serialization and engine handling semantics.
178#[derive(Clone)]
179pub struct BootFlowTaskManager {
180    engine: FlowEngine,
181    queue: Arc<Queue>,
182    job_name: String,
183    task_policy: BootFlowTaskPolicy,
184}
185
186impl fmt::Debug for BootFlowTaskManager {
187    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
188        formatter
189            .debug_struct("BootFlowTaskManager")
190            .field("queue", &self.queue.name())
191            .field("job_name", &self.job_name)
192            .field("task_policy", &self.task_policy)
193            .finish_non_exhaustive()
194    }
195}
196
197impl BootFlowTaskManager {
198    /// Creates a manager for one engine and Boot queue.
199    pub fn new(engine: FlowEngine, queue: Arc<Queue>) -> Self {
200        Self {
201            engine,
202            queue,
203            job_name: DEFAULT_FLOW_JOB_NAME.to_string(),
204            task_policy: BootFlowTaskPolicy::new(),
205        }
206    }
207
208    /// Replaces the Boot processor job name.
209    pub fn with_job_name(mut self, job_name: impl Into<String>) -> Result<Self> {
210        let job_name = job_name.into().trim().to_string();
211        if job_name.is_empty() {
212            return Err(FlowError::InvalidWorkerConfiguration(
213                "Boot Flow job name cannot be empty".to_string(),
214            ));
215        }
216        self.job_name = job_name;
217        Ok(self)
218    }
219
220    /// Replaces and validates the shared task policy.
221    pub fn with_task_policy(mut self, task_policy: BootFlowTaskPolicy) -> Result<Self> {
222        task_policy.validate()?;
223        self.task_policy = task_policy;
224        Ok(self)
225    }
226
227    /// Returns the Flow engine used by the registered processor.
228    pub fn engine(&self) -> &FlowEngine {
229        &self.engine
230    }
231
232    /// Returns the backing Boot queue.
233    pub fn queue(&self) -> Arc<Queue> {
234        Arc::clone(&self.queue)
235    }
236
237    /// Returns the registered Boot job name.
238    pub fn job_name(&self) -> &str {
239        &self.job_name
240    }
241
242    /// Returns the shared Flow task policy.
243    pub fn task_policy(&self) -> &BootFlowTaskPolicy {
244        &self.task_policy
245    }
246
247    /// Build the concrete Boot options that this manager will use for `task`.
248    ///
249    /// Hosts can add a caller-assigned job ID or other one-off Boot option and
250    /// pass the result to [`Self::enqueue_with_options`].
251    pub fn job_options_for(&self, task: &FlowTask) -> QueueJobOptions {
252        self.task_policy.job_options_for(&self.job_name, task)
253    }
254
255    /// Register the Flow task processor with the Boot queue.
256    ///
257    /// The host still starts and stops the queue through `QueueModule` or the
258    /// corresponding `Queue::start` and `Queue::shutdown` lifecycle calls.
259    pub fn register(&self) -> Result<()> {
260        let engine = self.engine.clone();
261        self.queue
262            .process(self.job_name.clone(), move |job: QueueJob, _context| {
263                let engine = engine.clone();
264                async move {
265                    let task = job.data_as::<FlowTask>()?;
266                    handle_flow_task(&engine, task).await.map_err(|error| {
267                        BootError::Internal(format!("A3S Flow task handling failed: {error}"))
268                    })?;
269                    Ok(())
270                }
271            })
272            .map_err(boot_error)
273    }
274
275    /// Enqueues one task and returns its Boot job receipt.
276    pub async fn enqueue_with_receipt(&self, task: FlowTask) -> Result<QueueJobReceipt> {
277        let options = self.job_options_for(&task);
278        self.enqueue_with_options(task, options).await
279    }
280
281    /// Enqueue one task with explicit typed A3S Boot job options.
282    ///
283    /// This per-submission entrypoint supports caller-assigned job IDs and the
284    /// complete `QueueJobOptions` surface. Scheduler dispatch through
285    /// [`FlowTaskDispatcher`] uses this manager's [`BootFlowTaskPolicy`].
286    pub async fn enqueue_with_options(
287        &self,
288        task: FlowTask,
289        options: QueueJobOptions,
290    ) -> Result<QueueJobReceipt> {
291        self.queue
292            .enqueue_with_options(self.job_name.clone(), &task, options)
293            .await
294            .map_err(boot_error)
295    }
296}
297
298#[async_trait]
299impl FlowTaskDispatcher for BootFlowTaskManager {
300    async fn dispatch(&self, task: FlowTask) -> Result<()> {
301        self.enqueue_with_receipt(task).await.map(|_| ())
302    }
303
304    fn has_runtime_build_route(&self, required_build_id: Option<&crate::RuntimeBuildId>) -> bool {
305        self.engine.supports_runtime_build(required_build_id)
306    }
307}
308
309fn boot_error(error: BootError) -> FlowError {
310    FlowError::TaskManagement(format!("A3S Boot queue error: {error}"))
311}
312
313fn flow_task_deduplication_id(job_name: &str, task: &FlowTask) -> String {
314    let mut hasher = Sha256::new();
315    hash_deduplication_field(&mut hasher, job_name);
316    let kind = match task {
317        FlowTask::DriveRun { run_id } => {
318            hash_deduplication_field(&mut hasher, run_id);
319            "drive_run"
320        }
321        FlowTask::ResumeWait { run_id, wait_id } => {
322            hash_deduplication_field(&mut hasher, run_id);
323            hash_deduplication_field(&mut hasher, wait_id);
324            "resume_wait"
325        }
326        FlowTask::ResumeHook {
327            run_id, hook_id, ..
328        } => {
329            hash_deduplication_field(&mut hasher, run_id);
330            hash_deduplication_field(&mut hasher, hook_id);
331            "resume_hook"
332        }
333        FlowTask::ResumeHookByToken { token, .. } => {
334            hash_deduplication_field(&mut hasher, token);
335            "resume_hook_by_token"
336        }
337        FlowTask::SendSignal { run_id, signal } => {
338            hash_deduplication_field(&mut hasher, run_id);
339            hash_deduplication_field(&mut hasher, &signal.signal_id);
340            "send_signal"
341        }
342        FlowTask::DisposeHook { run_id, hook_id } => {
343            hash_deduplication_field(&mut hasher, run_id);
344            hash_deduplication_field(&mut hasher, hook_id);
345            "dispose_hook"
346        }
347        FlowTask::DisposeHookByToken { token } => {
348            hash_deduplication_field(&mut hasher, token);
349            "dispose_hook_by_token"
350        }
351        FlowTask::ResumeScheduledRun { run_id, .. } => {
352            hash_deduplication_field(&mut hasher, run_id);
353            "resume_scheduled_run"
354        }
355        FlowTask::ResumeDueWaits { .. } => "resume_due_waits",
356        FlowTask::ResumeDueRetries { .. } => "resume_due_retries",
357    };
358    hash_deduplication_field(&mut hasher, kind);
359    format!("a3s-flow:{kind}:{:x}", hasher.finalize())
360}
361
362fn hash_deduplication_field(hasher: &mut Sha256, value: &str) {
363    let length = u64::try_from(value.len()).unwrap_or(u64::MAX);
364    hasher.update(length.to_be_bytes());
365    hasher.update(value.as_bytes());
366}
367
368fn flow_task_needs_active_successor(task: &FlowTask) -> bool {
369    matches!(
370        task,
371        FlowTask::DriveRun { .. }
372            | FlowTask::SendSignal { .. }
373            | FlowTask::ResumeScheduledRun { .. }
374            | FlowTask::ResumeDueWaits { .. }
375            | FlowTask::ResumeDueRetries { .. }
376    )
377}
a3s_flow/worker/boot.rs

a3s_flow/worker/
boot.rs
