mocra_core/common/model/
message.rs

1use crate::common::interface::StoreTrait;
2use crate::common::interface::storage::{BlobStorage, Offloadable};
3use crate::common::model::data::DataEvent;
4use crate::common::model::{ExecutionMark, Response};
5use async_trait::async_trait;
6use serde::{Deserialize, Serialize};
7use serde_json::Value;
8use std::sync::Arc;
9use uuid::Uuid;
10
11#[derive(Debug, Clone, Copy)]
12pub enum TopicType {
13    /// Task queue.
14    Task,
15    /// Request queue.
16    Request,
17    /// Response queue (raw task output).
18    Response,
19    /// Parser task queue (processed results).
20    ParserTask,
21    /// Error queue (error details).
22    Error,
23}
24
25impl TopicType {
26    /// Returns topic suffix.
27    pub(crate) fn suffix(&self) -> &'static str {
28        match self {
29            TopicType::Response => "response",
30            TopicType::ParserTask => "parser_task",
31            TopicType::Error => "error",
32            TopicType::Task => "task",
33            TopicType::Request => "request",
34        }
35    }
36    pub fn get_name(&self, name: &str) -> String {
37        format!("crawler-{}-{}", name, self.suffix())
38    }
39}
40
41/// Parser task message model.
42///
43/// Used to create downstream tasks after parsing, or move to the next processing stage.
44/// Contains task identity, metadata, execution context, and predecessor request reference.
45#[derive(Debug, Clone, Serialize, Deserialize)]
46pub struct TaskParserEvent {
47    /// Unique identifier.
48    pub id: Uuid,
49    /// Associated account task information.
50    pub account_task: TaskEvent,
51    /// Timestamp.
52    pub timestamp: u64,
53    /// Metadata (`TaskParserEvent.meta => Task.metadata => Context.meta.task_meta => Module.generate`).
54    pub metadata: serde_json::Map<String, Value>,
55    /// Execution context (`ExecutionMark`).
56    pub context: ExecutionMark,
57    /// Run identifier (Run ID).
58    #[serde(default = "default_run_id")]
59    pub run_id: Uuid,
60    /// Predecessor request identifier.
61    pub prefix_request: Uuid,
62}
63
64#[async_trait]
65impl Offloadable for TaskParserEvent {
66    fn should_offload(&self, _threshold: usize) -> bool {
67        false
68    }
69    async fn offload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
70        Ok(())
71    }
72    async fn reload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
73        Ok(())
74    }
75}
76
77impl TaskParserEvent {
78    /// Sets explicit execution context (`ExecutionMark`) for parser-chain execution.
79    ///
80    /// Typical uses:
81    /// - Precisely control which step index target module starts from.
82    /// - Combine with `stay_current_step` to prevent auto-advancing.
83    pub fn with_context(mut self, ctx: ExecutionMark) -> Self {
84        self.context = ctx;
85        self
86    }
87    /// Marks this task to stay on the current parser node and avoid auto-advance.
88    ///
89    /// Behavior:
90    /// - When parser returns `ParserTaskModel` and this is set, the chain loops on this step.
91    /// - Otherwise, the loop is considered complete and processing continues.
92    pub fn stay_current_step(mut self) -> Self {
93        // Set `ExecutionMark.stay_current_step`; executor locks progress to current response step.
94        self.context.stay_current_step = true;
95        self
96    }
97    pub fn get_context(&self) -> &ExecutionMark {
98        &self.context
99    }
100    pub fn with_meta<T>(mut self, meta: T) -> Self
101    where
102        T: serde::Serialize,
103    {
104        if let Ok(Value::Object(map)) = serde_json::to_value(meta) {
105            self.metadata = map;
106        }
107        self
108    }
109    pub fn add_meta<T>(mut self, key: impl AsRef<str>, value: T) -> Self
110    where
111        T: serde::Serialize,
112    {
113        if let Ok(val) = serde_json::to_value(value) {
114            self.metadata.insert(key.as_ref().to_string(), val);
115        }
116        self
117    }
118    /// Overrides chain backtracking pointer (points to predecessor `Request.id`).
119    ///
120    /// By default this pointer is inherited from `Response.prefix_request`.
121    /// For cross-module jumps, you can set it explicitly to change first-failure fallback target.
122    pub fn with_prefix_request(mut self, prefix: Uuid) -> Self {
123        self.prefix_request = prefix;
124        self
125    }
126    // Creates a new `ParserTaskModel` for target module within same account + platform.
127    pub fn start_other_module(response: &Response, module_name: impl AsRef<str>) -> Self {
128        debug_assert!(
129            !module_name.as_ref().is_empty(),
130            "module_name must not be empty"
131        );
132        TaskParserEvent {
133            id: Uuid::now_v7(),
134            account_task: TaskEvent {
135                account: response.account.clone(),
136                platform: response.platform.clone(),
137                module: Some(vec![module_name.as_ref().to_string()]),
138                priority: response.priority,
139                run_id: Uuid::now_v7(),
140            },
141            timestamp: chrono::Utc::now().timestamp() as u64,
142            metadata: serde_json::Map::new(),
143            // Reset context to ensure target module starts from the beginning.
144            context: ExecutionMark::default(),
145            run_id: Uuid::now_v7(),
146            prefix_request: response.prefix_request,
147        }
148    }
149    /// Creates `ParserTaskModel` for target module with explicit context.
150    ///
151    /// Typical usage:
152    /// - Cross-jump to step 0: pass `ExecutionMark::default().with_step_idx(0)`.
153    /// - Stay on current step during retry: pass `ExecutionMark` with `stay_current_step=true`.
154    pub fn start_other_module_with_ctx(
155        response: &Response,
156        module_name: impl AsRef<str>,
157        ctx: ExecutionMark,
158    ) -> Self {
159        debug_assert!(
160            !module_name.as_ref().is_empty(),
161            "module_name must not be empty"
162        );
163        TaskParserEvent {
164            id: Uuid::now_v7(),
165            account_task: TaskEvent {
166                account: response.account.clone(),
167                platform: response.platform.clone(),
168                module: Some(vec![module_name.as_ref().to_string()]),
169                priority: response.priority,
170                run_id: Uuid::now_v7(),
171            },
172            timestamp: chrono::Utc::now().timestamp() as u64,
173            metadata: serde_json::Map::new(),
174            context: ctx,
175            run_id: Uuid::now_v7(),
176            prefix_request: response.prefix_request,
177        }
178    }
179}
180
181/// Error task message model.
182///
183/// Records processing-time errors, including error details and execution context.
184#[derive(Debug, Clone, Serialize, Deserialize)]
185pub struct TaskErrorEvent {
186    /// Unique identifier.
187    pub id: Uuid,
188    /// Associated account task information.
189    pub account_task: TaskEvent,
190    /// Error message.
191    pub error_msg: String,
192    /// Timestamp.
193    pub timestamp: u64,
194    /// Metadata.
195    pub metadata: serde_json::Map<String, Value>,
196    /// Execution context.
197    pub context: ExecutionMark,
198    /// Run identifier.
199    #[serde(default = "default_run_id")]
200    pub run_id: Uuid,
201    /// Predecessor request identifier (points to predecessor `Request.id`).
202    pub prefix_request: Uuid,
203}
204
205#[async_trait]
206impl Offloadable for TaskErrorEvent {
207    fn should_offload(&self, _threshold: usize) -> bool {
208        false
209    }
210    async fn offload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
211        Ok(())
212    }
213    async fn reload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
214        Ok(())
215    }
216}
217
218impl TaskErrorEvent {
219    pub fn get_context(&self) -> &ExecutionMark {
220        &self.context
221    }
222    pub fn with_meta<T>(mut self, meta: T) -> Self
223    where
224        T: serde::Serialize,
225    {
226        if let Ok(Value::Object(map)) = serde_json::to_value(meta) {
227            self.metadata = map;
228        }
229        self
230    }
231    pub fn add_meta<T>(mut self, key: impl AsRef<str>, value: T) -> Self
232    where
233        T: serde::Serialize,
234    {
235        if let Ok(val) = serde_json::to_value(value) {
236            self.metadata.insert(key.as_ref().to_string(), val);
237        }
238        self
239    }
240    pub fn with_context(mut self, ctx: ExecutionMark) -> Self {
241        self.context = ctx;
242        self
243    }
244}
245/// Base task model.
246///
247/// Defines minimal task identity with account, platform, and module information.
248#[derive(Debug, Clone, Serialize, Deserialize)]
249pub struct TaskEvent {
250    /// Account identifier.
251    pub account: String,
252    /// Platform identifier.
253    pub platform: String,
254    /// Module list (optional; empty means all modules).
255    pub module: Option<Vec<String>>,
256    /// Priority.
257    #[serde(default)]
258    pub priority: crate::common::model::Priority,
259    /// Run identifier.
260    #[serde(default = "default_run_id")]
261    pub run_id: Uuid,
262}
263
264#[derive(Debug, Clone, Serialize, Deserialize)]
265#[serde(tag = "type", content = "payload", rename_all = "snake_case")]
266pub enum UnifiedTaskInput {
267    Task(TaskEvent),
268    ParserTask(TaskParserEvent),
269    ErrorTask(TaskErrorEvent),
270}
271
272impl UnifiedTaskInput {
273    pub fn run_id(&self) -> Uuid {
274        match self {
275            UnifiedTaskInput::Task(value) => value.run_id,
276            UnifiedTaskInput::ParserTask(value) => value.run_id,
277            UnifiedTaskInput::ErrorTask(value) => value.run_id,
278        }
279    }
280}
281
282impl From<TaskEvent> for UnifiedTaskInput {
283    fn from(value: TaskEvent) -> Self {
284        UnifiedTaskInput::Task(value)
285    }
286}
287
288impl From<TaskParserEvent> for UnifiedTaskInput {
289    fn from(value: TaskParserEvent) -> Self {
290        UnifiedTaskInput::ParserTask(value)
291    }
292}
293
294impl From<TaskErrorEvent> for UnifiedTaskInput {
295    fn from(value: TaskErrorEvent) -> Self {
296        UnifiedTaskInput::ErrorTask(value)
297    }
298}
299
300#[async_trait]
301impl Offloadable for TaskEvent {
302    fn should_offload(&self, _threshold: usize) -> bool {
303        false
304    }
305    async fn offload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
306        Ok(())
307    }
308    async fn reload(&mut self, _storage: &Arc<dyn BlobStorage>) -> crate::errors::Result<()> {
309        Ok(())
310    }
311}
312
313impl crate::common::model::priority::Prioritizable for TaskEvent {
314    fn get_priority(&self) -> crate::common::model::Priority {
315        self.priority
316    }
317}
318
319impl crate::common::model::priority::Prioritizable for TaskParserEvent {
320    fn get_priority(&self) -> crate::common::model::Priority {
321        self.account_task.priority
322    }
323}
324
325impl crate::common::model::priority::Prioritizable for TaskErrorEvent {
326    fn get_priority(&self) -> crate::common::model::Priority {
327        self.account_task.priority
328    }
329}
330
331impl From<&Response> for TaskParserEvent {
332    fn from(value: &Response) -> Self {
333        // Forward task metadata from the response so downstream nodes receive
334        // the same params that were used to generate the request.
335        let metadata = value.metadata.task.as_object().cloned().unwrap_or_default();
336
337        TaskParserEvent {
338            id: Uuid::now_v7(),
339            account_task: TaskEvent {
340                account: value.account.clone(),
341                platform: value.platform.clone(),
342                module: Some(vec![value.module.clone()]),
343                priority: value.priority,
344                run_id: value.run_id,
345            },
346            timestamp: chrono::Utc::now().timestamp() as u64,
347            metadata,
348            context: value.context.clone(),
349            run_id: value.run_id,
350            prefix_request: value.prefix_request,
351        }
352    }
353}
354
355impl From<&Response> for TaskErrorEvent {
356    fn from(value: &Response) -> Self {
357        // Forward only the task metadata slot as a flat map, consistent with
358        // TaskParserEvent so downstream generate() receives the same shape.
359        let metadata = value.metadata.task.as_object().cloned().unwrap_or_default();
360
361        TaskErrorEvent {
362            id: Uuid::now_v7(),
363            account_task: TaskEvent {
364                account: value.account.clone(),
365                platform: value.platform.clone(),
366                module: Some(vec![value.module.clone()]),
367                priority: value.priority,
368                run_id: value.run_id,
369            },
370            error_msg: "".to_string(),
371            timestamp: chrono::Utc::now().timestamp() as u64,
372            metadata,
373            context: value.context.clone(),
374            run_id: value.run_id,
375            prefix_request: value.prefix_request,
376        }
377    }
378}
379
380/// Serde default function to auto-generate a run_id using UUID v7
381fn default_run_id() -> Uuid {
382    Uuid::now_v7()
383}
384
385/// Parser output envelope.
386///
387/// Contains parsed data, next task, error task, and control flags.
388#[derive(Debug, Default, Clone)]
389pub struct TaskOutputEvent {
390    /// Parsed data list.
391    pub data: Vec<DataEvent>,
392    /// Generated next parser tasks.
393    pub parser_task: Vec<TaskParserEvent>,
394    /// Generated error task (optional).
395    pub error_task: Option<TaskErrorEvent>,
396    /// Stop flag (optional).
397    pub stop: Option<bool>,
398}
399impl TaskOutputEvent {
400    pub fn with_data(mut self, data: Vec<impl StoreTrait>) -> Self {
401        self.data = data.into_iter().map(|d| d.build()).collect();
402        self
403    }
404    pub fn with_task(mut self, task: TaskParserEvent) -> Self {
405        self.parser_task.push(task);
406        self
407    }
408    pub fn with_tasks(mut self, tasks: Vec<TaskParserEvent>) -> Self {
409        self.parser_task.extend(tasks);
410        self
411    }
412    pub fn with_error(mut self, error: TaskErrorEvent) -> Self {
413        self.error_task = Some(error);
414        self
415    }
416
417    /// Module-level stop flag; when `true`, subsequent requests are no longer processed.
418    /// For WSS flows this can be used to signal upper layers to close connection.
419    pub fn with_stop(mut self, stop: bool) -> Self {
420        self.stop = Some(stop);
421        self
422    }
423}
424
425#[cfg(test)]
426mod tests {
427    use super::*;
428    use crate::common::model::ExecutionMark;
429
430    #[test]
431    fn test_topic_type() {
432        assert_eq!(TopicType::Task.suffix(), "task");
433        assert_eq!(TopicType::Request.suffix(), "request");
434        assert_eq!(TopicType::Response.suffix(), "response");
435
436        let name = TopicType::Task.get_name("test_mod");
437        assert_eq!(name, "crawler-test_mod-task");
438    }
439
440    #[test]
441    fn test_parser_task_model_builder() {
442        let id = Uuid::now_v7();
443        let run_id = Uuid::now_v7();
444        let prefix_req = Uuid::now_v7();
445
446        let task = TaskParserEvent {
447            id,
448            account_task: TaskEvent {
449                account: "acc".into(),
450                platform: "plat".into(),
451                module: None,
452                priority: Default::default(),
453                run_id,
454            },
455            timestamp: 123456,
456            metadata: serde_json::Map::new(),
457            context: ExecutionMark::default(),
458            run_id,
459            prefix_request: prefix_req,
460        };
461
462        let ctx = ExecutionMark::default().with_step_idx(5);
463        let task = task.with_context(ctx.clone());
464        assert_eq!(task.context.step_idx, Some(5));
465
466        let task = task.stay_current_step();
467        assert!(task.context.stay_current_step);
468
469        let new_prefix = Uuid::now_v7();
470        let task = task.with_prefix_request(new_prefix);
471        assert_eq!(task.prefix_request, new_prefix);
472
473        let task = task.add_meta("foo", "bar");
474        assert_eq!(task.metadata["foo"], "bar");
475    }
476
477    #[test]
478    fn test_response_conversion() {
479        let run_id = Uuid::now_v7();
480        let prefix_request = Uuid::now_v7();
481        let response = Response {
482            id: Uuid::now_v7(),
483            platform: "plat".into(),
484            account: "acc".into(),
485            module: "mod".into(),
486            status_code: 200,
487            cookies: Default::default(),
488            content: vec![],
489            storage_path: None,
490            headers: vec![],
491            task_retry_times: 0,
492            metadata: crate::common::model::meta::MetaData::default(),
493            download_middleware: vec![],
494            data_middleware: vec![],
495            task_finished: false,
496            context: ExecutionMark::default(),
497            run_id,
498            prefix_request,
499            request_hash: None,
500            priority: Default::default(),
501        };
502
503        let parser_task: TaskParserEvent = (&response).into();
504        assert_eq!(parser_task.account_task.account, "acc");
505        assert_eq!(parser_task.account_task.platform, "plat");
506        assert_eq!(parser_task.run_id, run_id);
507        assert_eq!(parser_task.prefix_request, prefix_request);
508
509        let error_task: TaskErrorEvent = (&response).into();
510        assert_eq!(error_task.account_task.account, "acc");
511        assert_eq!(error_task.run_id, run_id);
512    }
513}
mocra_core/common/model/message.rs

mocra_core/common/model/
message.rs
