feat(java): handler registry, taskito-test, middleware contexts, submit payloads

crates/taskito-java/src/convert.rs

@@ -127,6 +127,8 @@ pub struct JobView<'a> {
     pub error: Option<&'a str>,
     pub unique_key: Option<&'a str>,
     pub namespace: Option<&'a str>,
+    /// Opaque metadata blob (JSON the SDK sets, e.g. middleware-injected trace ids).
+    pub metadata: Option<&'a str>,
 }
 
 impl<'a> From<&'a Job> for JobView<'a> {
@@ -148,6 +150,7 @@ impl<'a> From<&'a Job> for JobView<'a> {
             error: j.error.as_deref(),
             unique_key: j.unique_key.as_deref(),
             namespace: j.namespace.as_deref(),
+            metadata: j.metadata.as_deref(),
         }
     }
 }

sdks/java/README.md

@@ -130,6 +130,21 @@ try (Worker worker = queue.worker()
 A failed step (after its retries) fails the run and skips downstream nodes;
 `run.cancel()` skips pending nodes.
 
+Payloads can also be supplied **at submit** instead of baked into each step —
+declare structural steps with `stepAfter(name, task, deps...)` and pass a map:
+
+```java
+Workflow etl = Workflow.named("etl")
+        .stepAfter("extract", extract)
+        .stepAfter("transform", transform, "extract")
+        .stepAfter("load", load, "transform");
+
+queue.submitWorkflow(etl, Map.of("extract", 5, "transform", 6, "load", 7));
+```
+
+A step's effective payload is `map.get(name)` when present, else the one baked
+into the step.
+
 **Fan-out / fan-in** map a step over a producer's result list and gather the
 results:
 

sdks/java/processor/src/main/java/org/byteveda/taskito/processor/TaskHandlerProcessor.java

@@ -83,6 +83,8 @@ private void generate(TypeElement owner, List<ExecutableElement> methods) {
         }
         out.append("import com.fasterxml.jackson.core.type.TypeReference;\n")
                 .append("import org.byteveda.taskito.task.Task;\n")
+                .append("import org.byteveda.taskito.worker.Handler;\n")
+                .append("import org.byteveda.taskito.worker.HandlerRegistry;\n")
                 .append("import org.byteveda.taskito.worker.Worker;\n\n")
                 .append("/** Generated task descriptors + worker binding for {@link ")
                 .append(ownerSimple)
@@ -131,7 +133,27 @@ private void generate(TypeElement owner, List<ExecutableElement> methods) {
                         .append(");\n");
             }
         }
-        out.append("        return builder;\n    }\n}\n");
+        out.append("        return builder;\n    }\n\n");
+
+        // handlers() — returns a HandlerRegistry for Worker.Builder.register(...).
+        out.append("    public static HandlerRegistry handlers(")
+                .append(ownerSimple)
+                .append(" impl) {\n        return HandlerRegistry.of(\n");
+        for (int i = 0; i < methods.size(); i++) {
+            ExecutableElement method = methods.get(i);
+            String constant = upperSnake(method.getSimpleName().toString());
+            String methodName = method.getSimpleName().toString();
+            out.append("                Handler.of(").append(constant).append(", ");
+            if (isVoid(method)) {
+                out.append("payload -> {\n                    impl.")
+                        .append(methodName)
+                        .append("(payload);\n                    return null;\n                })");
+            } else {
+                out.append("impl::").append(methodName).append(")");
+            }
+            out.append(i + 1 < methods.size() ? ",\n" : ");\n");
+        }
+        out.append("    }\n}\n");
 
         write(owner, qualified, out.toString());
     }

sdks/java/settings.gradle.kts

@@ -1,3 +1,4 @@
 rootProject.name = "taskito"
 
 include(":processor")
+include(":test-support")

sdks/java/src/main/java/org/byteveda/taskito/DefaultQueue.java

@@ -86,7 +86,13 @@ private String dispatchEnqueue(String taskName, Object payload, EnqueueOptions o
         for (Middleware m : middleware) {
             m.onEnqueue(context);
         }
-        return backend.enqueue(taskName, serializer.serialize(context.payload()), encode(context.options()));
+        EnqueueOptions finalOptions = context.options();
+        if (!context.metadata().isEmpty()) {
+            finalOptions = finalOptions.toBuilder()
+                    .metadata(encode(context.metadata()))
+                    .build();
+        }
+        return backend.enqueue(taskName, serializer.serialize(context.payload()), encode(finalOptions));
     }
 
     @Override
@@ -313,18 +319,31 @@ public long registerPeriodic(PeriodicTask task) {
 
     @Override
     public WorkflowRun submitWorkflow(Workflow workflow) {
+        return submitWorkflow(workflow, java.util.Collections.emptyMap());
+    }
+
+    @Override
+    public WorkflowRun submitWorkflow(Workflow workflow, Map<String, Object> suppliedPayloads) {
         List<Step> steps = workflow.steps();
         Set<String> deferred = deferredNodes(steps);
         List<Map<String, Object>> specs = new ArrayList<>(steps.size());
         // Deferred nodes (fan-out/fan-in + their downstream) have no job — and so
-        // no payload — at submit; the tracker enqueues them at runtime.
+        // no payload — at submit; the tracker enqueues them at runtime. A static
+        // node's payload is the one supplied at submit, else the one baked in.
         List<String> payloadNames = new ArrayList<>();
         List<byte[]> payloads = new ArrayList<>();
         for (Step step : steps) {
             specs.add(stepSpec(step));
             if (!deferred.contains(step.name)) {
+                Object payload = suppliedPayloads.getOrDefault(step.name, step.payload);
+                // A payloadless structural step (stepAfter) must get its payload at
+                // submit; fail fast rather than enqueue a null-payload job.
+                if (payload == null && !suppliedPayloads.containsKey(step.name)) {
+                    throw new TaskitoException("workflow step '" + step.name
+                            + "' has no payload; supply one via submitWorkflow(wf, payloads)");
+                }
                 payloadNames.add(step.name);
-                payloads.add(serializer.serialize(step.payload));
+                payloads.add(serializer.serialize(payload));
             }
         }
         String runId = backend.submitWorkflow(

sdks/java/src/main/java/org/byteveda/taskito/Queue.java

@@ -146,6 +146,14 @@ public interface Queue extends AutoCloseable {
     /** Submit a workflow DAG; returns a handle to the run. */
     WorkflowRun submitWorkflow(Workflow workflow);
 
+    /**
+     * Submit a workflow, supplying per-step payloads keyed by step name. A step's
+     * effective payload is {@code payloads.get(name)} when present, else the
+     * payload baked into the step. Pairs with the structural
+     * {@code Workflow.stepAfter(name, task, deps...)} form.
+     */
+    WorkflowRun submitWorkflow(Workflow workflow, Map<String, Object> payloads);
+
     /** Current status of a workflow run, or empty if it no longer exists. */
     Optional<WorkflowStatus> workflowStatus(String runId);
 

sdks/java/src/main/java/org/byteveda/taskito/middleware/EnqueueContext.java

@@ -1,23 +1,35 @@
 package org.byteveda.taskito.middleware;
 
+import java.util.HashMap;
+import java.util.Map;
 import org.byteveda.taskito.task.EnqueueOptions;
 
 /**
  * A job being enqueued, passed to {@link Middleware#onEnqueue} before
- * serialization. Replace the payload or options to rewrite the job; throw to
- * abort the enqueue.
+ * serialization. Replace the payload or options to rewrite the job, add to
+ * {@link #metadata()} to travel with the job, or throw to abort the enqueue.
  */
 public final class EnqueueContext {
     public final String taskName;
     private Object payload;
     private EnqueueOptions options;
+    private final Map<String, Object> metadata = new HashMap<>();
 
     public EnqueueContext(String taskName, Object payload, EnqueueOptions options) {
         this.taskName = taskName;
         this.payload = payload;
         this.options = options;
     }
 
+    /**
+     * Mutable metadata that travels with the job (readable at execution via
+     * {@code TaskContext.job().metadata()}). When non-empty it becomes the job's
+     * metadata blob, replacing any set on the options.
+     */
+    public Map<String, Object> metadata() {
+        return metadata;
+    }
+
     public Object payload() {
         return payload;
     }

sdks/java/src/main/java/org/byteveda/taskito/middleware/JobInfo.java

@@ -0,0 +1,38 @@
+package org.byteveda.taskito.middleware;
+
+import java.util.Map;
+import java.util.function.Supplier;
+
+/**
+ * The executing job, exposed to {@link Middleware} hooks. {@link #metadata()} is
+ * loaded lazily (a storage read) only when first accessed, so middleware that
+ * never reads it pays nothing.
+ */
+public final class JobInfo {
+    private final String id;
+    private final String taskName;
+    private final Supplier<Map<String, Object>> metadataLoader;
+    private Map<String, Object> metadata;
+
+    public JobInfo(String id, String taskName, Supplier<Map<String, Object>> metadataLoader) {
+        this.id = id;
+        this.taskName = taskName;
+        this.metadataLoader = metadataLoader;
+    }
+
+    public String id() {
+        return id;
+    }
+
+    public String taskName() {
+        return taskName;
+    }
+
+    /** The job's metadata map (e.g. trace ids injected at enqueue); loaded on first call. */
+    public Map<String, Object> metadata() {
+        if (metadata == null) {
+            metadata = metadataLoader.get();
+        }
+        return metadata;
+    }
+}

sdks/java/src/main/java/org/byteveda/taskito/middleware/TaskContext.java

@@ -1,12 +1,37 @@
 package org.byteveda.taskito.middleware;
 
-/** Identifies a task as it executes on a worker. */
+import java.util.HashMap;
+import java.util.Map;
+
+/**
+ * Identifies a task as it executes on a worker. {@link #attributes()} is a
+ * per-execution scratch map shared across {@code before}/{@code after}/
+ * {@code onError} for the same job; {@link #job()} exposes its metadata.
+ */
 public final class TaskContext {
     public final String jobId;
     public final String taskName;
 
-    public TaskContext(String jobId, String taskName) {
+    private final Map<String, Object> attributes = new HashMap<>();
+    private final JobInfo job;
+
+    public TaskContext(String jobId, String taskName, JobInfo job) {
         this.jobId = jobId;
         this.taskName = taskName;
+        this.job = job;
+    }
+
+    public TaskContext(String jobId, String taskName) {
+        this(jobId, taskName, new JobInfo(jobId, taskName, java.util.Collections::emptyMap));
+    }
+
+    /** Mutable per-execution scratch shared across this job's middleware hooks. */
+    public Map<String, Object> attributes() {
+        return attributes;
+    }
+
+    /** The executing job, including its (lazily loaded) metadata. */
+    public JobInfo job() {
+        return job;
     }
 }

sdks/java/src/main/java/org/byteveda/taskito/model/Job.java

@@ -23,6 +23,7 @@ public final class Job {
     public final String error;
     public final String uniqueKey;
     public final String namespace;
+    public final String metadata;
 
     @JsonCreator
     public Job(
@@ -41,7 +42,8 @@ public Job(
             @JsonProperty("progress") Integer progress,
             @JsonProperty("error") String error,
             @JsonProperty("uniqueKey") String uniqueKey,
-            @JsonProperty("namespace") String namespace) {
+            @JsonProperty("namespace") String namespace,
+            @JsonProperty("metadata") String metadata) {
         this.id = id;
         this.queue = queue;
         this.taskName = taskName;
@@ -58,5 +60,6 @@ public Job(
         this.error = error;
         this.uniqueKey = uniqueKey;
         this.namespace = namespace;
+        this.metadata = metadata;
     }
 }

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

@@ -0,0 +1,27 @@
+package org.byteveda.taskito.worker;
+
+import org.byteveda.taskito.task.Task;
+import org.byteveda.taskito.task.TaskFunction;
+
+/** A task descriptor paired with the function that handles it. */
+public final class Handler<T, R> {
+    private final Task<T> task;
+    private final TaskFunction<T, R> function;
+
+    private Handler(Task<T> task, TaskFunction<T, R> function) {
+        this.task = task;
+        this.function = function;
+    }
+
+    public static <T, R> Handler<T, R> of(Task<T> task, TaskFunction<T, R> function) {
+        return new Handler<>(task, function);
+    }
+
+    public Task<T> task() {
+        return task;
+    }
+
+    public TaskFunction<T, R> function() {
+        return function;
+    }
+}

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

@@ -0,0 +1,28 @@
+package org.byteveda.taskito.worker;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+
+/**
+ * An immutable bundle of {@link Handler}s registered together via
+ * {@code Worker.Builder.register}. Generated {@code <Class>Tasks.handlers(impl)}
+ * returns one of these.
+ */
+public final class HandlerRegistry {
+    private final List<Handler<?, ?>> handlers;
+
+    private HandlerRegistry(List<Handler<?, ?>> handlers) {
+        this.handlers = handlers;
+    }
+
+    public static HandlerRegistry of(Handler<?, ?>... handlers) {
+        // Copy the varargs array so later caller mutations can't leak in.
+        return new HandlerRegistry(Collections.unmodifiableList(new ArrayList<>(Arrays.asList(handlers))));
+    }
+
+    public List<Handler<?, ?>> handlers() {
+        return handlers;
+    }
+}

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

@@ -120,6 +120,19 @@ public Builder apply(Consumer<Builder> customizer) {
             return this;
         }
 
+        /** Register a single {@link Handler} (a task + its function). */
+        public Builder register(Handler<?, ?> handler) {
+            handlers.put(
+                    handler.task().name(), new RegisteredTask(handler.task().payloadType(), cast(handler.function())));
+            return this;
+        }
+
+        /** Register every handler in a {@link HandlerRegistry} (e.g. a generated {@code XxxTasks.handlers}). */
+        public Builder register(HandlerRegistry registry) {
+            registry.handlers().forEach(this::register);
+            return this;
+        }
+
         public Builder queues(String... queues) {
             this.queues = Arrays.asList(queues);
             return this;
@@ -159,7 +172,8 @@ public Worker start() {
                     concurrency > 0 ? Executors.newFixedThreadPool(concurrency) : Executors.newCachedThreadPool();
             Emitter emitter = new Emitter();
             listeners.forEach((name, bound) -> bound.forEach(listener -> emitter.on(name, listener)));
-            WorkerDispatchBridge bridge = new WorkerDispatchBridge(handlers, serializer, executor, emitter, middleware);
+            WorkerDispatchBridge bridge =
+                    new WorkerDispatchBridge(backend, handlers, serializer, executor, emitter, middleware);
             WorkerControl control = backend.startWorker(bridge, encodeOptions());
             bridge.bind(control);
             return new Worker(control, executor);