feat(java): add mesh scheduling worker

.github/workflows/ci-java.yml

@@ -26,7 +26,7 @@ jobs:
         uses: ./.github/actions/setup-rust
 
       - name: Build the JNI cdylib
-        run: cargo build --release --features postgres,redis,workflows -p taskito-java
+        run: cargo build --release --features postgres,redis,workflows,mesh -p taskito-java
 
       - name: Upload the cdylib
         uses: actions/upload-artifact@v7

.github/workflows/publish-java.yml

@@ -95,7 +95,7 @@ jobs:
         uses: ./.github/actions/setup-rust
 
       - name: Build cdylib
-        run: cargo build -p taskito-java --release --features postgres,redis,workflows
+        run: cargo build -p taskito-java --release --features postgres,redis,workflows,mesh
 
       - name: Stage native binary under its classifier
         shell: bash

crates/taskito-java/Cargo.toml

@@ -14,10 +14,12 @@ default = []
 postgres = ["taskito-core/postgres", "taskito-workflows?/postgres"]
 redis = ["taskito-core/redis", "taskito-workflows?/redis"]
 workflows = ["dep:taskito-workflows"]
+mesh = ["dep:taskito-mesh"]
 
 [dependencies]
 taskito-core = { path = "../taskito-core" }
 taskito-workflows = { path = "../taskito-workflows", optional = true }
+taskito-mesh = { path = "../taskito-mesh", optional = true }
 jni = "0.21"
 serde = { workspace = true }
 serde_json = { workspace = true }

crates/taskito-java/src/convert.rs

@@ -165,6 +165,11 @@ pub struct WorkerOptions {
     /// Per-task retry-backoff policies, registered with the scheduler at start.
     /// The core owns the retry engine; this only feeds it the backoff curve.
     pub task_configs: Option<Vec<TaskRetryConfig>>,
+    /// Raw `taskito_mesh::MeshConfig` JSON. When present (and the `mesh` feature
+    /// is built), the worker interposes the mesh bridge between scheduler and
+    /// dispatcher; otherwise this is ignored.
+    #[cfg_attr(not(feature = "mesh"), allow(dead_code))]
+    pub mesh_config: Option<String>,
 }
 
 /// A task's retry-backoff curve. Fields left unset fall back to the core's

crates/taskito-java/src/lib.rs

@@ -16,6 +16,8 @@ mod ffi;
 mod ffi_c;
 mod handle;
 mod jvm;
+#[cfg(feature = "mesh")]
+mod mesh;
 mod queue;
 mod worker;
 #[cfg(feature = "workflows")]

crates/taskito-java/src/mesh.rs

@@ -0,0 +1,120 @@
+//! Mesh scheduling bridge (feature `mesh`).
+//!
+//! When a worker is started with a `meshConfig`, [`spawn_mesh`] interposes a
+//! [`taskito_mesh::MeshNode`] between the scheduler and the dispatcher: the
+//! scheduler emits ready jobs into an intermediate channel, the bridge sorts
+//! them into the node's affinity-aware local deque (and steals from busy peers),
+//! then drains the deque to the real dispatcher channel. The DB stays the source
+//! of truth, so the scheduler and dispatcher remain mesh-unaware — identical to
+//! the Python and Node shells.
+
+use std::sync::Arc;
+
+use taskito_core::job::Job;
+use taskito_core::scheduler::Scheduler;
+use taskito_mesh::{MeshConfig, MeshNode};
+use tokio::sync::mpsc::Sender;
+
+/// Build the mesh node, start its gossip + steal-server tasks and the bridge
+/// loop on `runtime`, and return the node so the worker can request shutdown and
+/// read cluster info. `threads` is advertised to peers as this worker's capacity.
+pub fn spawn_mesh(
+    runtime: &tokio::runtime::Runtime,
+    scheduler: Arc<Scheduler>,
+    job_tx: Sender<Job>,
+    config_json: &str,
+    worker_id: String,
+    queues: Vec<String>,
+    threads: u16,
+) -> Arc<MeshNode> {
+    let config: MeshConfig = serde_json::from_str(config_json).unwrap_or_else(|e| {
+        log::error!("[taskito-java] invalid mesh config, falling back to defaults: {e}");
+        MeshConfig::default()
+    });
+    log::info!(
+        "[taskito-java] mesh scheduling enabled (gossip_port={}, local_buffer={})",
+        config.gossip_port,
+        config.local_buffer_capacity,
+    );
+    let node = Arc::new(MeshNode::new(worker_id, config));
+
+    // `spawn_gossip`/`spawn_steal_server` call the free `tokio::spawn` internally,
+    // so they must run inside the runtime context (we are on the JNI thread here).
+    let (gossip, steal_server) = {
+        let _guard = runtime.enter();
+        (
+            node.spawn_gossip(queues, threads),
+            node.spawn_steal_server(),
+        )
+    };
+
+    let bridge_node = node.clone();
+    runtime.spawn(async move {
+        run_mesh_bridge(scheduler, bridge_node, job_tx).await;
+    });
+    // Keep the gossip + steal-server tasks alive for the worker's lifetime; they
+    // stop when the node's shutdown is signalled (see `request_shutdown`).
+    runtime.spawn(async move {
+        let (gossip_res, steal_res) = tokio::join!(gossip, steal_server);
+        if let Err(e) = gossip_res {
+            log::error!("[taskito-java] mesh gossip task failed: {e}");
+        }
+        if let Err(e) = steal_res {
+            log::error!("[taskito-java] mesh steal-server task failed: {e}");
+        }
+    });
+
+    node
+}
+
+/// Pump jobs from the scheduler through the mesh local deque to the dispatcher.
+async fn run_mesh_bridge(scheduler: Arc<Scheduler>, mesh_node: Arc<MeshNode>, job_tx: Sender<Job>) {
+    let (mesh_tx, mut mesh_rx) = tokio::sync::mpsc::channel::<Job>(64);
+
+    let sched = scheduler.clone();
+    let sched_task = tokio::spawn(async move {
+        sched.run(mesh_tx).await;
+    });
+
+    loop {
+        // Drain the local deque to the dispatcher first.
+        while let Some(job) = mesh_node.pop_local() {
+            if job_tx.send(job).await.is_err() {
+                let _ = sched_task.await;
+                return;
+            }
+        }
+
+        // Steal from a busier peer when our deque is low.
+        if mesh_node.should_steal() {
+            mesh_node.try_steal().await;
+            while let Some(job) = mesh_node.pop_local() {
+                if job_tx.send(job).await.is_err() {
+                    let _ = sched_task.await;
+                    return;
+                }
+            }
+        }
+
+        // Wait for the scheduler to produce jobs, then prefetch the batch.
+        match mesh_rx.recv().await {
+            Some(job) => {
+                let mut batch = vec![job];
+                while let Ok(extra) = mesh_rx.try_recv() {
+                    batch.push(extra);
+                }
+                mesh_node.prefetch(batch);
+            }
+            None => break,
+        }
+    }
+
+    // Scheduler done: drain whatever remains in the deque.
+    while let Some(job) = mesh_node.pop_local() {
+        if job_tx.send(job).await.is_err() {
+            break;
+        }
+    }
+
+    let _ = sched_task.await;
+}

crates/taskito-java/src/worker.rs

@@ -38,6 +38,10 @@ pub struct WorkerHandle {
     registry: Arc<Registry>,
     shutdown: Arc<Notify>,
     heartbeat_stop: Arc<Notify>,
+    /// The mesh node, when mesh scheduling is enabled. Held so `stop` can signal
+    /// its gossip + steal-server tasks and `meshClusterInfo` can read its state.
+    #[cfg(feature = "mesh")]
+    mesh: Option<Arc<taskito_mesh::MeshNode>>,
 }
 
 /// Build and start a worker over `storage`, calling back into the Java
@@ -56,6 +60,9 @@ fn start_worker(
     let queues = options
         .queues
         .unwrap_or_else(|| vec![DEFAULT_QUEUE.to_string()]);
+    // Mesh gossip advertises the served queues; capture before `queues` moves.
+    #[cfg(feature = "mesh")]
+    let mesh_queues = queues.clone();
     let capacity = options
         .channel_capacity
         .map(|c| (c as usize).max(1))
@@ -70,6 +77,9 @@ fn start_worker(
     let lifecycle_storage = storage.clone();
     let queues_csv = queues.join(",");
     let worker_id = format!("java-{}", uuid::Uuid::now_v7());
+    // The mesh node id is this worker's id; capture before `worker_id` moves.
+    #[cfg(feature = "mesh")]
+    let mesh_worker_id = worker_id.clone();
 
     let mut scheduler = Scheduler::new(storage, queues, config, namespace);
     // Claim execution under this worker's id so dead-worker recovery can
@@ -87,8 +97,30 @@ fn start_worker(
     // so the result queue can't grow without bound.
     let (result_tx, result_rx) = crossbeam_channel::unbounded();
 
-    // Scheduler loop: poll storage and dispatch ready jobs.
+    // Scheduler loop: poll storage and dispatch ready jobs. With mesh enabled and
+    // a config supplied, route the scheduler's output through the mesh bridge
+    // (affinity-sorted local deque + work-stealing) instead of straight to the
+    // dispatcher; the DB stays the source of truth, so the plain path is identical.
     let scheduler_run = scheduler.clone();
+    #[cfg(feature = "mesh")]
+    let mesh = match options.mesh_config.as_deref() {
+        Some(config_json) => Some(crate::mesh::spawn_mesh(
+            &runtime,
+            scheduler_run,
+            job_tx,
+            config_json,
+            mesh_worker_id,
+            mesh_queues,
+            capacity.min(u16::MAX as usize) as u16,
+        )),
+        None => {
+            runtime.spawn(async move {
+                scheduler_run.run(job_tx).await;
+            });
+            None
+        }
+    };
+    #[cfg(not(feature = "mesh"))]
     runtime.spawn(async move {
         scheduler_run.run(job_tx).await;
     });
@@ -121,6 +153,8 @@ fn start_worker(
         registry,
         shutdown,
         heartbeat_stop,
+        #[cfg(feature = "mesh")]
+        mesh,
     })
 }
 
@@ -412,10 +446,38 @@ pub extern "system" fn Java_org_byteveda_taskito_internal_NativeWorker_stop(
         let worker = unsafe { borrow_worker(handle) };
         worker.shutdown.notify_one();
         worker.heartbeat_stop.notify_one();
+        // Signal the mesh node so its gossip + steal-server tasks exit instead of
+        // lingering until the runtime drops.
+        #[cfg(feature = "mesh")]
+        if let Some(mesh) = &worker.mesh {
+            mesh.request_shutdown();
+        }
         Ok(())
     })
 }
 
+/// `String meshClusterInfo(long workerHandle)` — a JSON `ClusterInfo` snapshot,
+/// or `null` when this worker is not mesh-enabled.
+#[no_mangle]
+pub extern "system" fn Java_org_byteveda_taskito_internal_NativeWorker_meshClusterInfo<'local>(
+    mut env: JNIEnv<'local>,
+    _class: JClass<'local>,
+    handle: jlong,
+) -> jni::sys::jstring {
+    guard(&mut env, std::ptr::null_mut(), |_env| {
+        let worker = unsafe { borrow_worker(handle) };
+        #[cfg(feature = "mesh")]
+        if let Some(mesh) = &worker.mesh {
+            let json = serde_json::to_string(&mesh.cluster_info()).map_err(|e| {
+                crate::error::BindingError::new(format!("failed to encode cluster info: {e}"))
+            })?;
+            return crate::ffi::new_string(_env, json);
+        }
+        let _ = worker;
+        Ok(std::ptr::null_mut())
+    })
+}
+
 /// `void close(long workerHandle)` — stop the runtime and reclaim the handle.
 /// The Java side drains its handler executor before calling this, so no
 /// `completeJob`/`failJob` can still borrow the handle being freed.

sdks/java/build.gradle.kts

@@ -117,7 +117,7 @@ val nativeStaging = layout.buildDirectory.dir("native")
 // Build the native library for the local platform.
 val cargoBuild = tasks.register<Exec>("cargoBuild") {
     workingDir = crateDir.asFile
-    commandLine("cargo", "build", "--release", "--features", "postgres,redis,workflows")
+    commandLine("cargo", "build", "--release", "--features", "postgres,redis,workflows,mesh")
 }
 
 // Stage the built library under its platform-classifier resource path.

sdks/java/src/main/java/org/byteveda/taskito/internal/JniWorkerControl.java

@@ -1,5 +1,6 @@
 package org.byteveda.taskito.internal;
 
+import java.util.Optional;
 import org.byteveda.taskito.spi.WorkerControl;
 
 /** JNI-backed {@link WorkerControl} over a native worker handle. */
@@ -31,6 +32,11 @@ public void stop() {
         NativeWorker.stop(handle);
     }
 
+    @Override
+    public Optional<String> meshClusterInfoJson() {
+        return Optional.ofNullable(NativeWorker.meshClusterInfo(handle));
+    }
+
     /** Idempotent: frees the native worker handle exactly once. */
     @Override
     public synchronized void close() {

sdks/java/src/main/java/org/byteveda/taskito/internal/NativeWorker.java

@@ -23,4 +23,7 @@ private NativeWorker() {}
     public static native void stop(long handle);
 
     public static native void close(long handle);
+
+    /** A JSON {@code ClusterInfo} snapshot, or {@code null} when not mesh-enabled. */
+    public static native String meshClusterInfo(long handle);
 }

sdks/java/src/main/java/org/byteveda/taskito/spi/WorkerControl.java

@@ -1,5 +1,7 @@
 package org.byteveda.taskito.spi;
 
+import java.util.Optional;
+
 /** Controls a running worker and completes its in-flight jobs. */
 public interface WorkerControl extends AutoCloseable {
     void completeJob(long token, byte[] result);
@@ -11,6 +13,11 @@ public interface WorkerControl extends AutoCloseable {
     /** Stop the scheduler and heartbeat loops; in-flight jobs drain. */
     void stop();
 
+    /** A JSON {@code ClusterInfo} snapshot when mesh is enabled, else empty. */
+    default Optional<String> meshClusterInfoJson() {
+        return Optional.empty();
+    }
+
     @Override
     void close();
 }

sdks/java/src/main/java/org/byteveda/taskito/worker/MeshClusterInfo.java

@@ -0,0 +1,23 @@
+package org.byteveda.taskito.worker;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+/**
+ * A snapshot of a mesh worker's view of its cluster, from
+ * {@link Worker#meshClusterInfo()}. Capacities and loads are summed across alive
+ * peers (excluding this node); buffered counts are job-deque lengths.
+ *
+ * @param peerCount alive peers discovered via gossip (excludes self)
+ * @param totalCapacity summed advertised capacity of those peers
+ * @param totalLoad summed in-flight jobs across those peers
+ * @param totalBuffered summed local-deque length across those peers
+ * @param localBufferLen this worker's local-deque length
+ * @param adaptivePrefetch this worker's current prefetch budget
+ */
+public record MeshClusterInfo(
+        @JsonProperty("peer_count") int peerCount,
+        @JsonProperty("total_capacity") int totalCapacity,
+        @JsonProperty("total_load") int totalLoad,
+        @JsonProperty("total_buffered") int totalBuffered,
+        @JsonProperty("local_buffer_len") int localBufferLen,
+        @JsonProperty("adaptive_prefetch") int adaptivePrefetch) {}