cortex/
worker.rs

1// Copyright 2015-2025 Deyan Ginev. See the LICENSE
2// file at the top-level directory of this distribution.
3//
4// Licensed under the MIT license <LICENSE-MIT or http://opensource.org/licenses/MIT>.
5// This file may not be copied, modified, or distributed
6// except according to those terms.
7
8//! Worker for performing corpus imports, when served as "init" tasks by the `CorTeX` dispatcher
9use rand::seq::IndexedRandom;
10use std::borrow::Cow;
11use std::error::Error;
12use std::fs::File;
13use std::path::Path;
14use std::thread;
15use std::time::Duration;
16use zmq::{Context, Message, SNDMORE};
17
18use crate::backend;
19use crate::backend::default_db_address;
20use crate::importer::Importer;
21use crate::models::{Corpus, NewCorpus, Task};
22use pericortex::worker::Worker;
23
24/// Resolves the ZMQ socket identity for a worker: the **operator-configured** `identity` when set
25/// (a stable, operator-controlled metadata key — its `worker_metadata` row then accumulates across
26/// restarts instead of fragmenting under a fresh random name each start, KNOWN_ISSUES W-3),
27/// otherwise a random ephemeral handle so an unconfigured worker still gets a unique name. The
28/// operator is responsible for keeping configured identities unique per worker (two DEALER sockets
29/// sharing an identity would break the dispatcher's ROUTER addressing).
30fn resolve_worker_identity(configured: &str, rng: &mut impl rand::Rng) -> String {
31  if !configured.is_empty() {
32    return configured.to_string();
33  }
34  let letters: Vec<char> = "abcdefghijklmonpqrstuvwxyz".chars().collect();
35  // 19 random letters — preserves the historical ephemeral-identity length/charset.
36  (1..20).map(|_| *letters.choose(rng).unwrap()).collect()
37}
38
39/// `Worker` for initializing/importing a new corpus into `CorTeX`
40#[derive(Debug, Clone)]
41pub struct InitWorker {
42  /// name of the service ("init")
43  pub service: String,
44  /// version, as usual
45  pub version: f32,
46  /// message size, as usual
47  pub message_size: usize,
48  /// full URL (including port) to task source/dispatcher
49  pub source: String,
50  /// full URL (including port) to task sink/receiver
51  pub sink: String,
52  /// address to the Task store backend
53  /// (special case, only for the init service, third-party workers can't access the Task store
54  /// directly)
55  pub backend_address: String,
56  /// thread-local unique identifier
57  pub identity: String,
58}
59impl Default for InitWorker {
60  fn default() -> InitWorker {
61    InitWorker {
62      service: "init".to_string(),
63      version: 0.1,
64      message_size: 100_000,
65      source: "tcp://localhost:51695".to_string(),
66      sink: "tcp://localhost:51696".to_string(),
67      backend_address: default_db_address().to_string(),
68      identity: String::new(),
69    }
70  }
71}
72impl Worker for InitWorker {
73  fn get_service(&self) -> &str { &self.service }
74  fn get_source_address(&self) -> Cow<'_, str> { Cow::Borrowed(&self.source) }
75  fn get_sink_address(&self) -> Cow<'_, str> { Cow::Borrowed(&self.sink) }
76  fn get_identity(&self) -> &str { &self.identity }
77  fn set_identity(&mut self, identity: String) { self.identity = identity; }
78  fn message_size(&self) -> usize { self.message_size }
79
80  fn convert(&self, path_opt: &Path) -> Result<File, Box<dyn Error>> {
81    let path = path_opt.to_str().unwrap().to_string();
82    let name = path
83      .rsplit_once('/')
84      .map(|x| x.1)
85      .unwrap_or(&path)
86      .to_lowercase(); // TODO: this is Unix path only
87    let mut backend = backend::from_address(&self.backend_address);
88    let corpus = NewCorpus {
89      name,
90      path,
91      complex: true,
92      description: String::new(),
93    };
94    // Add the new corpus.
95    backend.add(&corpus).expect("Failed to create new corpus.");
96    let registered_corpus = Corpus::find_by_name(&corpus.name, &mut backend.connection)
97      .expect("Failed to create new corpus.");
98
99    // Create an importer for the corpus, and then process all entries to populate CorTeX tasks
100    let mut importer = Importer {
101      corpus: registered_corpus,
102      backend: backend::from_address(&self.backend_address),
103      ..Importer::default()
104    };
105
106    importer.process()?;
107    // TODO: Stopgap, we should do the error-reporting well
108    Err(From::from("init worker does not return a file handle."))
109  }
110
111  fn start(&mut self, limit: Option<usize>) -> Result<(), Box<dyn Error>> {
112    let mut work_counter = 0;
113    let mut rng = rand::rng();
114    // Connect to a task ventilator
115    let context_source = Context::new();
116    let source = context_source.socket(zmq::DEALER).unwrap();
117    let identity = resolve_worker_identity(&self.identity, &mut rng);
118    source.set_identity(identity.as_bytes()).unwrap();
119
120    source.connect(&self.get_source_address()).unwrap();
121    // Connect to a task sink
122    let context_sink = Context::new();
123    let sink = context_sink.socket(zmq::PUSH).unwrap();
124    sink.connect(&self.get_sink_address()).unwrap();
125    let mut backend = backend::from_address(&self.backend_address);
126    // Work in perpetuity
127    loop {
128      let mut taskid_msg = Message::new();
129      let mut recv_msg = Message::new();
130
131      source.send(self.get_service(), 0).unwrap();
132      source.recv(&mut taskid_msg, 0).unwrap();
133      let taskid = taskid_msg.as_str().unwrap();
134
135      // Terminating with an empty message in place of a payload (INIT is special)
136      source.recv(&mut recv_msg, 0).unwrap();
137
138      let task_result = Task::find(taskid.parse::<i64>().unwrap(), &mut backend.connection);
139      let task = match task_result {
140        Ok(t) => t,
141        _ => {
142          // If there was nothing to do, retry a minute later
143          thread::sleep(Duration::new(60, 0));
144          continue;
145        },
146      };
147      // ignore error for now, complete the task.
148      let _ = self.convert(Path::new(&task.entry));
149      sink.send(&identity, SNDMORE).unwrap();
150      sink.send(self.get_service(), SNDMORE).unwrap();
151      sink.send(taskid, SNDMORE).unwrap();
152      sink.send(Vec::new(), 0).unwrap();
153
154      work_counter += 1;
155      if let Some(upper_bound) = limit
156        && work_counter >= upper_bound
157      {
158        // Give enough time to complete the Final job.
159        thread::sleep(Duration::from_millis(500));
160        break;
161      };
162    }
163    Ok(())
164  }
165}
166
167#[cfg(test)]
168mod tests {
169  use super::resolve_worker_identity;
170
171  #[test]
172  fn configured_identity_is_honored_verbatim() {
173    // An operator-set identity is used as-is, giving a stable worker_metadata key (W-3).
174    let mut rng = rand::rng();
175    assert_eq!(
176      resolve_worker_identity("arxiv-host3:init:1", &mut rng),
177      "arxiv-host3:init:1"
178    );
179  }
180
181  #[test]
182  fn empty_identity_falls_back_to_a_random_handle() {
183    let mut rng = rand::rng();
184    let id = resolve_worker_identity("", &mut rng);
185    assert_eq!(id.len(), 19, "preserves the historical 19-char length");
186    assert!(
187      id.chars().all(|c| c.is_ascii_lowercase()),
188      "random handle is lowercase letters only"
189    );
190  }
191}
cortex/worker.rs

cortex/
worker.rs
