feat(java): worker — dispatch bridge, handlers, events

crates/taskito-java/Cargo.toml

@@ -20,3 +20,9 @@ jni = "0.21"
 serde = { workspace = true }
 serde_json = { workspace = true }
 once_cell = "1"
+tokio = { workspace = true }
+async-trait = { workspace = true }
+crossbeam-channel = { workspace = true }
+uuid = { workspace = true }
+log = { workspace = true }
+gethostname = "1"

crates/taskito-java/src/convert.rs

@@ -150,6 +150,15 @@ impl<'a> From<&'a Job> for JobView<'a> {
     }
 }
 
+/// Options accepted by `NativeQueue.runWorker`.
+#[derive(Deserialize, Default)]
+#[serde(rename_all = "camelCase", default)]
+pub struct WorkerOptions {
+    pub queues: Option<Vec<String>>,
+    pub channel_capacity: Option<u32>,
+    pub batch_size: Option<u32>,
+}
+
 /// Filter accepted by `NativeQueue.listJobs`.
 #[derive(Deserialize, Default)]
 #[serde(rename_all = "camelCase", default)]

crates/taskito-java/src/dispatcher.rs

@@ -0,0 +1,207 @@
+//! `JavaDispatcher` — runs each job by calling back into Java.
+//!
+//! Implements the core [`WorkerDispatcher`] trait. The completion registry lives
+//! here in Rust (the inverse of the Node shell's JS-promise bridge): for each
+//! job a token + oneshot is registered, the job is submitted to Java via
+//! `onJob`, and the awaiting task parks on the oneshot — no worker thread blocks.
+//! Java runs the task (sync or async) and completes it through the native
+//! `NativeWorker.completeJob/failJob/cancelJob` entry points (see `worker.rs`),
+//! which resolve the oneshot.
+
+use std::collections::HashMap;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::sync::{Arc, Mutex};
+use std::time::{Duration, Instant};
+
+use crossbeam_channel::Sender;
+use jni::objects::{GlobalRef, JValue};
+use taskito_core::job::Job;
+use taskito_core::scheduler::JobResult;
+use taskito_core::worker::WorkerDispatcher;
+use taskito_core::{Storage, StorageBackend};
+use tokio::sync::oneshot;
+
+use crate::jvm;
+
+/// The outcome Java reports for a submitted job.
+pub enum TaskOutcome {
+    Success(Vec<u8>),
+    Failure(String),
+    Cancelled,
+}
+
+/// Pending-job registry shared between the dispatcher and the Java completion
+/// callbacks. Held alive by the worker handle so its pointer stays valid.
+#[derive(Default)]
+pub struct Registry {
+    pending: Mutex<HashMap<u64, oneshot::Sender<TaskOutcome>>>,
+    next_token: AtomicU64,
+}
+
+impl Registry {
+    fn register(&self) -> (u64, oneshot::Receiver<TaskOutcome>) {
+        let token = self.next_token.fetch_add(1, Ordering::Relaxed);
+        let (tx, rx) = oneshot::channel();
+        self.pending.lock().unwrap().insert(token, tx);
+        (token, rx)
+    }
+
+    fn forget(&self, token: u64) {
+        self.pending.lock().unwrap().remove(&token);
+    }
+
+    /// Resolve a pending job. A no-op if the token already completed or timed out.
+    pub fn complete(&self, token: u64, outcome: TaskOutcome) {
+        if let Some(tx) = self.pending.lock().unwrap().remove(&token) {
+            let _ = tx.send(outcome);
+        }
+    }
+}
+
+/// Executes jobs by dispatching them to a Java `WorkerBridge` callback object.
+pub struct JavaDispatcher {
+    callbacks: GlobalRef,
+    registry: Arc<Registry>,
+    storage: StorageBackend,
+}
+
+impl JavaDispatcher {
+    pub fn new(callbacks: GlobalRef, registry: Arc<Registry>, storage: StorageBackend) -> Self {
+        Self {
+            callbacks,
+            registry,
+            storage,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl WorkerDispatcher for JavaDispatcher {
+    async fn run(
+        &self,
+        mut job_rx: tokio::sync::mpsc::Receiver<Job>,
+        result_tx: Sender<JobResult>,
+    ) {
+        while let Some(job) = job_rx.recv().await {
+            let callbacks = self.callbacks.clone();
+            let registry = self.registry.clone();
+            let storage = self.storage.clone();
+            let result_tx = result_tx.clone();
+            tokio::spawn(async move {
+                let result = run_one(&callbacks, &registry, &storage, job).await;
+                let _ = result_tx.send(result);
+            });
+        }
+    }
+
+    fn shutdown(&self) {}
+}
+
+/// Submit one job to Java, await its completion, and translate to a [`JobResult`].
+async fn run_one(
+    callbacks: &GlobalRef,
+    registry: &Registry,
+    storage: &StorageBackend,
+    job: Job,
+) -> JobResult {
+    let started = Instant::now();
+    let (token, rx) = registry.register();
+
+    if let Err(err) = submit_to_java(callbacks, token, &job) {
+        registry.forget(token);
+        return failure(job, err, started.elapsed().as_nanos() as i64, false);
+    }
+
+    // `timeout_ms <= 0` means no limit. On timeout we drop the pending entry; a
+    // late completion then finds no sender and is harmlessly ignored.
+    let outcome = if job.timeout_ms > 0 {
+        match tokio::time::timeout(Duration::from_millis(job.timeout_ms as u64), rx).await {
+            Ok(result) => result,
+            Err(_) => {
+                registry.forget(token);
+                let wall = started.elapsed().as_nanos() as i64;
+                return failure(job, "task timed out".to_string(), wall, true);
+            }
+        }
+    } else {
+        rx.await
+    };
+
+    let wall = started.elapsed().as_nanos() as i64;
+    match outcome {
+        Ok(TaskOutcome::Success(result)) => JobResult::Success {
+            job_id: job.id,
+            result: Some(result),
+            task_name: job.task_name,
+            wall_time_ns: wall,
+        },
+        Ok(TaskOutcome::Cancelled) => cancelled(job, wall),
+        Ok(TaskOutcome::Failure(error)) => {
+            // A failure on a cancel-requested job is a cancellation, not a fault.
+            if storage.is_cancel_requested(&job.id).unwrap_or(false) {
+                cancelled(job, wall)
+            } else {
+                failure(job, error, wall, false)
+            }
+        }
+        // The oneshot sender dropped without completing — the Java side died.
+        Err(_) => failure(job, "java task channel dropped".to_string(), wall, false),
+    }
+}
+
+/// Invoke `WorkerBridge.onJob` on an attached thread. Local refs are freed when
+/// the per-call attachment guard drops.
+fn submit_to_java(callbacks: &GlobalRef, token: u64, job: &Job) -> Result<(), String> {
+    let vm = jvm::vm();
+    let mut env = vm
+        .attach_current_thread()
+        .map_err(|e| format!("attach failed: {e}"))?;
+    let job_id = env.new_string(&job.id).map_err(|e| e.to_string())?;
+    let task_name = env.new_string(&job.task_name).map_err(|e| e.to_string())?;
+    let payload = env
+        .byte_array_from_slice(&job.payload)
+        .map_err(|e| e.to_string())?;
+    // An `Err(JavaException)` leaves the exception pending on the attached
+    // thread; clear it before returning so the thread isn't poisoned for the
+    // next call.
+    if let Err(e) = env.call_method(
+        callbacks,
+        "onJob",
+        "(JLjava/lang/String;Ljava/lang/String;[B)V",
+        &[
+            JValue::Long(token as i64),
+            JValue::Object(&job_id),
+            JValue::Object(&task_name),
+            JValue::Object(&payload),
+        ],
+    ) {
+        let _ = env.exception_clear();
+        return Err(e.to_string());
+    }
+    if env.exception_check().unwrap_or(false) {
+        let _ = env.exception_clear();
+        return Err("WorkerBridge.onJob threw".to_string());
+    }
+    Ok(())
+}
+
+fn cancelled(job: Job, wall_time_ns: i64) -> JobResult {
+    JobResult::Cancelled {
+        job_id: job.id,
+        task_name: job.task_name,
+        wall_time_ns,
+    }
+}
+
+fn failure(job: Job, error: String, wall_time_ns: i64, timed_out: bool) -> JobResult {
+    JobResult::Failure {
+        job_id: job.id,
+        error,
+        retry_count: job.retry_count,
+        max_retries: job.max_retries,
+        task_name: job.task_name,
+        wall_time_ns,
+        should_retry: true,
+        timed_out,
+    }
+}

crates/taskito-java/src/jvm.rs

@@ -12,7 +12,6 @@ pub fn init(vm: JavaVM) {
 
 /// The cached VM (present after `JNI_OnLoad`). Used by the worker bridge to
 /// attach worker threads that invoke Java callbacks.
-#[allow(dead_code)]
 pub fn vm() -> &'static JavaVM {
     JAVA_VM.get().expect("JNI_OnLoad has not run")
 }

crates/taskito-java/src/lib.rs

@@ -10,11 +10,13 @@
 
 mod backend;
 mod convert;
+mod dispatcher;
 mod error;
 mod ffi;
 mod handle;
 mod jvm;
 mod queue;
+mod worker;
 
 use std::ffi::c_void;