feat(sdks): taskito executor subcommand for Python, Node and Java

crates/taskito-core/src/lib.rs

@@ -46,7 +46,8 @@ pub use storage::Storage;
 pub use storage::StorageBackend;
 pub use storage::{DeadJob, QueueStats, SubscriptionBacklogStats};
 pub use worker::{
-    AttachError, AttachedExecutor, Capacity, ExecutorMessage, NativeDispatcher, ProtocolError,
-    RemoteConfig, RemoteDispatcher, SchedulerMessage, Secret, TaskError, TaskRegistry, TaskResult,
-    Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION,
+    AttachAddress, AttachError, AttachedExecutor, Capacity, ExecutorClient, ExecutorConfig,
+    ExecutorError, ExecutorHandle, ExecutorMessage, ExecutorSession, NativeDispatcher,
+    ProtocolError, RemoteConfig, RemoteDispatcher, SchedulerMessage, Secret, TaskError,
+    TaskRegistry, TaskResult, Transport, Worker, WorkerDispatcher, WorkerHandle, PROTOCOL_VERSION,
 };

crates/taskito-core/src/worker/cancel.rs

@@ -0,0 +1,145 @@
+//! How a dispatcher learns that a running job was cancelled.
+//!
+//! An in-process worker reads the storage flag `Storage::request_cancel` sets.
+//! An attached executor has no storage — it is the whole point of #546 that it
+//! carries no database credentials — and learns instead from the `cancel` frame
+//! the scheduler sends, which arrives as
+//! [`WorkerDispatcher::notify_cancel`](super::WorkerDispatcher::notify_cancel).
+//!
+//! Both sources answer one question, so both dispatchers ask it here rather
+//! than growing two cancel paths each.
+
+use std::collections::HashSet;
+use std::sync::{Mutex, PoisonError};
+
+use crate::storage::{Storage, StorageBackend};
+
+/// The cancel sources available to one dispatcher.
+pub struct CancelSignals {
+    /// Present only for a dispatcher running inside a worker.
+    storage: Option<StorageBackend>,
+    /// Ids delivered out of band. Kept until the job reports, so a cancel that
+    /// races a job's start still fires rather than being missed.
+    signalled: Mutex<HashSet<String>>,
+}
+
+impl CancelSignals {
+    /// Read cancels from storage, and from `notify_cancel` when it is called.
+    pub fn from_storage(storage: StorageBackend) -> Self {
+        Self {
+            storage: Some(storage),
+            signalled: Mutex::new(HashSet::new()),
+        }
+    }
+
+    /// Read cancels only from `notify_cancel` — the attached-executor case.
+    pub fn detached() -> Self {
+        Self {
+            storage: None,
+            signalled: Mutex::new(HashSet::new()),
+        }
+    }
+
+    /// Record a cancel request for `job_id`.
+    pub fn signal(&self, job_id: &str) {
+        self.lock().insert(job_id.to_string());
+    }
+
+    /// Whether `job_id` has been cancelled.
+    ///
+    /// The out-of-band set is checked first: it needs no I/O, and for a
+    /// detached executor it is the only answer there is.
+    pub fn is_cancelled(&self, job_id: &str) -> bool {
+        if self.lock().contains(job_id) {
+            return true;
+        }
+        self.storage
+            .as_ref()
+            .is_some_and(|storage| storage.is_cancel_requested(job_id).unwrap_or(false))
+    }
+
+    /// Drop the record for a finished job, so the set cannot grow for the life
+    /// of the process.
+    pub fn forget(&self, job_id: &str) {
+        self.lock().remove(job_id);
+    }
+
+    fn lock(&self) -> std::sync::MutexGuard<'_, HashSet<String>> {
+        // The state behind the lock is a plain set, so reading it stays safe
+        // even if a holder panicked.
+        self.signalled
+            .lock()
+            .unwrap_or_else(PoisonError::into_inner)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn a_detached_signal_is_observed_without_storage() {
+        let signals = CancelSignals::detached();
+        assert!(!signals.is_cancelled("job-1"));
+
+        signals.signal("job-1");
+        assert!(signals.is_cancelled("job-1"));
+    }
+
+    #[test]
+    fn forgetting_a_job_clears_its_signal() {
+        // Ids are held until the job reports, so they have to be released or
+        // the set grows for the life of the process.
+        let signals = CancelSignals::detached();
+        signals.signal("job-1");
+        signals.forget("job-1");
+        assert!(!signals.is_cancelled("job-1"));
+    }
+
+    #[test]
+    fn an_unknown_job_is_not_cancelled() {
+        let signals = CancelSignals::detached();
+        signals.signal("job-1");
+        assert!(!signals.is_cancelled("job-2"));
+    }
+
+    #[test]
+    fn a_storage_flag_is_honoured_too() {
+        use crate::job::{now_millis, NewJob};
+        use crate::storage::sqlite::SqliteStorage;
+
+        let storage = StorageBackend::Sqlite(SqliteStorage::in_memory().expect("in-memory"));
+        let job = storage
+            .enqueue(NewJob {
+                queue: "default".to_string(),
+                task_name: "resize".to_string(),
+                payload: Vec::new(),
+                priority: 0,
+                scheduled_at: now_millis(),
+                max_retries: 0,
+                timeout_ms: 0,
+                unique_key: None,
+                metadata: None,
+                notes: None,
+                depends_on: vec![],
+                expires_at: None,
+                result_ttl_ms: None,
+                namespace: None,
+            })
+            .expect("enqueue");
+
+        // `request_cancel` only flags a *running* job — a pending one is
+        // cancelled outright — so the job has to be dequeued first.
+        let running = storage
+            .dequeue("default", now_millis() + 1_000, None)
+            .expect("dequeue")
+            .expect("the enqueued job");
+        assert_eq!(running.id, job.id);
+
+        let signals = CancelSignals::from_storage(storage.clone());
+        assert!(!signals.is_cancelled(&job.id));
+
+        assert!(storage.request_cancel(&job.id).expect("request cancel"));
+        assert!(signals.is_cancelled(&job.id));
+    }
+}

crates/taskito-core/src/worker/dial.rs

@@ -0,0 +1,219 @@
+//! Dial the address an executor was pointed at.
+//!
+//! The listener parses the same grammar on the bind side
+//! (`taskito-server`'s `config::listen`), so the two stay readable against each
+//! other: whatever `TASKITO_LISTEN` accepts, `TASKITO_ATTACH` dials. Every SDK
+//! shares this rather than reimplementing the grammar in its own language,
+//! where `unix:` support would inevitably drift.
+
+use std::io;
+use std::net::{TcpStream, ToSocketAddrs};
+#[cfg(unix)]
+use std::os::unix::net::UnixStream;
+use std::time::Duration;
+
+#[cfg(unix)]
+use super::transport::UnixTransport;
+use super::transport::{TcpTransport, Transport};
+
+/// Where an executor attaches.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum AttachAddress {
+    /// TCP, for a scheduler in another container or host.
+    Tcp(String),
+    /// Unix domain socket, the same-pod sidecar case.
+    #[cfg(unix)]
+    Unix(std::path::PathBuf),
+}
+
+impl std::fmt::Display for AttachAddress {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::Tcp(target) => write!(f, "tcp://{target}"),
+            #[cfg(unix)]
+            Self::Unix(path) => write!(f, "unix:{}", path.display()),
+        }
+    }
+}
+
+impl AttachAddress {
+    /// Parse one attach spec: `unix:/path`, `tcp://host:port`, `host:port`, or
+    /// `:port`.
+    ///
+    /// A bare `:port` means loopback, matching the listener's reading of the
+    /// same ambiguous value.
+    pub fn parse(spec: &str) -> io::Result<Self> {
+        let spec = spec.trim();
+        if spec.is_empty() {
+            return Err(invalid("an attach address must not be empty"));
+        }
+
+        if let Some(path) = spec.strip_prefix("unix:") {
+            #[cfg(unix)]
+            {
+                if path.is_empty() {
+                    return Err(invalid(
+                        "a unix attach address needs a socket path, e.g. unix:/run/taskito.sock",
+                    ));
+                }
+                return Ok(Self::Unix(std::path::PathBuf::from(path)));
+            }
+            #[cfg(not(unix))]
+            {
+                let _ = path;
+                return Err(invalid(
+                    "unix socket attach addresses are not supported on this platform",
+                ));
+            }
+        }
+
+        // The listener prints itself as `tcp://host:port`, so an operator who
+        // copies that line out of the logs must get a working address back.
+        let target = spec.strip_prefix("tcp://").unwrap_or(spec);
+        let target = match target.strip_prefix(':') {
+            Some(port) => format!("127.0.0.1:{port}"),
+            None => target.to_string(),
+        };
+        if !target.contains(':') {
+            return Err(invalid(format!(
+                "'{spec}' has no port — an attach address looks like host:port or \
+                 unix:/run/taskito.sock"
+            )));
+        }
+        Ok(Self::Tcp(target))
+    }
+
+    /// Open a connection to this address.
+    ///
+    /// `timeout` bounds the TCP connect so an unreachable scheduler fails
+    /// promptly instead of sitting in the platform's default retry window,
+    /// which can be minutes.
+    pub fn connect(&self, timeout: Duration) -> io::Result<Box<dyn Transport>> {
+        match self {
+            Self::Tcp(target) => {
+                let addresses = target.to_socket_addrs().map_err(|error| {
+                    invalid(format!("'{target}' is not a valid host:port: {error}"))
+                })?;
+                // Every resolved address is tried, not just the first: a
+                // dual-stack scheduler resolves to both an AAAA and an A
+                // record, and a host that cannot route one still reaches the
+                // other. The last failure is what gets reported.
+                let mut last_error = None;
+                for address in addresses {
+                    match TcpStream::connect_timeout(&address, timeout) {
+                        Ok(stream) => return Ok(Box::new(TcpTransport::new(stream)?)),
+                        Err(error) => last_error = Some(error),
+                    }
+                }
+                Err(last_error
+                    .unwrap_or_else(|| invalid(format!("'{target}' resolved to no address"))))
+            }
+            #[cfg(unix)]
+            Self::Unix(path) => {
+                // No connect timeout exists for a Unix socket, and none is
+                // needed: the peer is on this host, so a connect either
+                // succeeds or fails at once.
+                let stream = UnixStream::connect(path)?;
+                Ok(Box::new(UnixTransport::new(stream)))
+            }
+        }
+    }
+}
+
+fn invalid(message: impl Into<String>) -> io::Error {
+    io::Error::new(io::ErrorKind::InvalidInput, message.into())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn a_host_and_port_parses_as_tcp() {
+        assert_eq!(
+            AttachAddress::parse("scheduler:7749").expect("parse"),
+            AttachAddress::Tcp("scheduler:7749".to_string())
+        );
+    }
+
+    #[test]
+    fn the_tcp_scheme_the_listener_prints_is_accepted() {
+        // The listener logs `attach listener on tcp://127.0.0.1:7749`; pasting
+        // that back must work.
+        assert_eq!(
+            AttachAddress::parse("tcp://127.0.0.1:7749").expect("parse"),
+            AttachAddress::Tcp("127.0.0.1:7749".to_string())
+        );
+    }
+
+    #[test]
+    fn a_bare_port_means_loopback() {
+        assert_eq!(
+            AttachAddress::parse(":7749").expect("parse"),
+            AttachAddress::Tcp("127.0.0.1:7749".to_string())
+        );
+    }
+
+    #[test]
+    fn surrounding_whitespace_is_ignored() {
+        // Shell heredocs and Kubernetes manifests both leak trailing newlines.
+        assert_eq!(
+            AttachAddress::parse("  127.0.0.1:7749\n").expect("parse"),
+            AttachAddress::Tcp("127.0.0.1:7749".to_string())
+        );
+    }
+
+    #[test]
+    fn an_address_without_a_port_is_rejected_with_the_shape_it_wanted() {
+        let error = AttachAddress::parse("scheduler").expect_err("must be rejected");
+        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
+        assert!(
+            error.to_string().contains("host:port"),
+            "the message must show the expected shape: {error}"
+        );
+    }
+
+    #[test]
+    fn an_empty_address_is_rejected() {
+        assert!(AttachAddress::parse("").is_err());
+        assert!(AttachAddress::parse("   ").is_err());
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn a_unix_path_parses_and_prints_back() {
+        let address = AttachAddress::parse("unix:/run/taskito.sock").expect("parse");
+        assert_eq!(
+            address,
+            AttachAddress::Unix(std::path::PathBuf::from("/run/taskito.sock"))
+        );
+        assert_eq!(address.to_string(), "unix:/run/taskito.sock");
+    }
+
+    #[cfg(unix)]
+    #[test]
+    fn a_unix_scheme_without_a_path_is_rejected() {
+        let error = AttachAddress::parse("unix:").expect_err("must be rejected");
+        assert!(error.to_string().contains("socket path"), "{error}");
+    }
+
+    #[test]
+    fn connecting_to_a_closed_port_fails_rather_than_hanging() {
+        // Port 1 on loopback: reserved, and nothing listens there.
+        let address = AttachAddress::parse("127.0.0.1:1").expect("parse");
+        assert!(address.connect(Duration::from_millis(500)).is_err());
+    }
+
+    #[test]
+    fn a_dialed_address_round_trips_through_a_real_listener() {
+        let listener = std::net::TcpListener::bind("127.0.0.1:0").expect("bind");
+        let port = listener.local_addr().expect("addr").port();
+        let accepting = std::thread::spawn(move || listener.accept().expect("accept"));
+
+        let address = AttachAddress::parse(&format!(":{port}")).expect("parse");
+        let transport = address.connect(Duration::from_secs(5)).expect("connect");
+        assert!(transport.peer().starts_with("tcp:"));
+
+        let _ = accepting.join();
+    }
+}