[samza][producer] Share partitioned Venice write execution

[kvargha]

Problem Statement

VeniceSystemProducer STREAM writes currently enter the core writer on the caller thread. Segment initialization, schema work, writer-hook throttling, or PubSub backpressure can block a stream-processing task even though the API returns a future. OnlineVeniceProducer already isolates core writes behind partition-striped workers, but its execution primitive was not reusable.

Solution

Extract the bounded partition-striped executor into venice-common and reuse it from both producer APIs. Lifecycle behavior remains API-specific: OnlineVeniceProducer keeps its existing settings and metrics, while VeniceSystemProducer adds coordination for submission futures, durable futures, flush fences, startup rollback, and bounded shutdown.

For STREAM writes, VeniceSystemProducer serializes immutable write commands before enqueueing and routes them with the partition selected by the constructed core writer. Legacy writers that cannot expose partition routing use one conservative stripe to preserve ordering. The protected write path remains override-compatible and returns the existing durable future after core submission completes. Samza-facing send, put, and delete therefore retain caller-visible VeniceWriterHook backpressure and immediate submission failures without waiting for the broker acknowledgment.

The shared executor provides bounded blocking admission, per-stripe FIFO ordering, worker and callback thread management, and forced-shutdown cleanup. Active work retains terminal ownership during forced shutdown, while queued work is rejected exactly once. A Venice-owned daemon executor guarantees off-thread completion progress without relying on the common pool.

The SystemProducer-specific coordinator owns sticky asynchronous failures, flush markers, lifecycle fencing, interrupt restoration, and lifecycle-safe future completion. Terminal writer cleanup runs as one retained daemon task, so each stop() call observes its deadline without racing another flush or close. Later retries resume the same cleanup. Partial startup is rolled back before a retry can replace live resources. BATCH and STREAM_REPROCESSING writes remain inline.

OnlineVeniceProducer retains its existing client.producer.* settings, defaults, thread names, and metrics. It also gains an optional VeniceWriterHook factory overload so quota or admission-control hooks can be wired without another nearline producer implementation. Worker-originated and reentrant close completions are handed off before running arbitrary user continuations. Reentrant close exempts only its current tracked completion depth while draining all other work, and subclass clients close after core cleanup even when a sticky deferred failure is rethrown.

BatchingVeniceWriter now stops its periodic checker before closing the internal VeniceWriter it owns. Checker-thread close is deferred until the checker exits, and external checker shutdown remains bounded.

New VeniceSystemProducer settings:

• venice.system.producer.worker.count: default 4; set to 0 for the previous inline behavior.
• venice.system.producer.worker.queue.capacity: default 100000 per worker.
• venice.system.producer.callback.thread.count: default 0.
• venice.system.producer.callback.queue.capacity: default 100000.

Partitions assigned to the same worker stripe can still experience head-of-line blocking. Queue saturation intentionally applies caller backpressure, and flush() continues to wait for all accepted writes before returning.

stop() bounds the caller's wait, not the underlying writer operation. A timed-out writer cleanup continues on a daemon executor and can be observed by a later retry.

Code changes

• Added new code behind a config. Config names and defaults are listed above.
• Introduced new log lines.
  • Confirmed if logs need to be rate limited to avoid excessive logging. Lifecycle messages occur once per producer, and routing fallback warns once per dispatcher.
• VALIDATION_OVERRIDE: the final production diff exceeds the 2,000-line aggregate check because the bounded, retryable lifecycle and deterministic completion ownership require explicit state. Multiple independent reviews found the shared executor design to be the simplest complete implementation; compacting these paths would reduce readability of concurrency-critical code.

Concurrency-Specific Checks

Both reviewer and PR author to verify

• Code has no race conditions or thread safety issues.
• Proper synchronization mechanisms (e.g., synchronized, RWLock) are used where needed.
• No blocking calls inside critical sections that could lead to deadlocks or performance degradation.
  • Public flush() remains synchronous by contract. stop() runs terminal writer cleanup off-thread and bounds each caller wait.
• Verified thread-safe collections are used (e.g., ConcurrentHashMap, CopyOnWriteArrayList).
• Validated proper exception handling in multi-threaded code to avoid silent thread termination.

How was this PR tested?

• New unit tests added.
• New integration tests added.
• Modified or extended existing tests.
• Verified backward compatibility (if applicable).
• Local code review completed.

Coverage includes partition ordering, independent stripe progress, bounded blocking backpressure, inline mode, immutable command capture, exact and legacy writer routing, submission and durable completion, flush fencing, startup rollback, sticky failures, interruption, active-versus-queued shutdown ownership, bounded retryable cleanup, lifecycle reentrancy, deferred completion draining, batching-writer ownership, configuration isolation, writer-hook forwarding, and Online Producer compatibility.

All tests added or modified by this PR were audited for scheduler, timeout, latch, executor, and cleanup nondeterminism. Concurrency tests use explicit fence, admission, queue, marker, and completion signals rather than sleeps, yields, common-pool assumptions, Mockito timeout polling, or thread-state checks.

• VeniceSystemProducerDispatchTest passed 15/15 in five independent JDK 11 processes with retries disabled (75 invocations). The same suite also passed on JDK 8 and 17.
• PartitionedVeniceWriteExecutorTest passed 19/19 on JDK 8, 11, and 17 (57 invocations).
• OnlineVeniceProducerTest passed 27/27 on JDK 17. The protected-send compatibility test passed all four result variants on JDK 8, 11, and 17.
• Batching-writer close passed 26 tests on JDK 8, 11, and 17; six concurrency cases passed 50 runs each per JDK.
• All six affected SpotBugs main/test tasks, spotlessCheck, and git diff --check passed.
• A multi-agent review covered whole-diff correctness, SystemProducer lifecycle, Online/Liminal compatibility, executor/DIV semantics, test determinism, and architectural simplicity. The final post-fix review reported no findings.

Does this PR introduce any user-facing or breaking changes?

• No. You can skip the rest of this section.
• Yes. Clearly explain the behavior change and its impact.

VeniceSystemProducer STREAM writes now use four partition-striped workers by default. Public APIs and existing settings remain compatible, but core write submission moves off the caller thread after synchronous validation and serialization. Set venice.system.producer.worker.count=0 to restore inline execution.

[Copilot]

Pull request overview

This PR extracts partition-striped write dispatch into venice-common and reuses it across Online Venice Producer and VeniceSystemProducer STREAM writes to avoid blocking the caller thread on core-writer work (schema/segment init, backpressure), while preserving existing public APIs and configuration namespaces.

Changes:

• Introduces PartitionedVeniceWriteExecutor and PartitionedVeniceWriteDispatcher to provide partition-striped FIFO execution, bounded backpressure, flush fencing, sticky failure propagation, and orderly drain/close behavior.
• Updates VeniceSystemProducer STREAM path to enqueue immutable serialized write commands onto the dispatcher (with opt-out via venice.system.producer.worker.count=0) and adds corresponding config keys + tests.
• Retains Online Producer compatibility by shimming the old PartitionedProducerExecutor name onto the shared executor and adding optional VeniceWriterHook plumbing + tests.

Reviewed changes

Copilot reviewed 21 out of 21 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
internal/venice-common/src/test/java/com/linkedin/venice/writer/VeniceWriterUnitTest.java	Adds assertion that captured partition matches writer partition routing for serialized key.
internal/venice-common/src/test/java/com/linkedin/venice/writer/PartitionedVeniceWriteExecutorTest.java	New unit tests for striping FIFO, backpressure blocking, shutdown rejection semantics, and inline mode.
internal/venice-common/src/test/java/com/linkedin/venice/writer/PartitionedVeniceWriteDispatcherTest.java	New unit tests for routing, fencing/flush ordering, sticky failures, interruption handling, and shutdown behavior.
internal/venice-common/src/test/java/com/linkedin/venice/writer/BatchingVeniceWriterTest.java	Adds test ensuring getPartitionId delegates through batching wrapper to internal writer.
internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java	Implements getPartitionId(byte[]) to expose exact partition routing.
internal/venice-common/src/main/java/com/linkedin/venice/writer/PartitionedVeniceWriteExecutor.java	Adds shared bounded striped executor (workers + optional callback executor) with blocking admission.
internal/venice-common/src/main/java/com/linkedin/venice/writer/PartitionedVeniceWriteDispatcher.java	Adds STREAM write coordinator with immutable command capture, fences, sticky failure propagation, and drain-before-close.
internal/venice-common/src/main/java/com/linkedin/venice/writer/BatchingVeniceWriter.java	Overrides getPartitionId to delegate to internal writer.
internal/venice-common/src/main/java/com/linkedin/venice/writer/AbstractVeniceWriter.java	Adds non-abstract compatibility extension getPartitionId(byte[]) with default UnsupportedOperationException.
internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java	Adds new venice.system.producer.* config keys for STREAM dispatcher workers/callbacks.
integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerTest.java	Updates existing verifications to tolerate async dispatch by using Mockito timeouts.
integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerDispatchTest.java	New tests validating enqueue-vs-submission semantics, routing source, flush/stop fencing, sticky failures, and inline rollback behavior.
integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerConfigTest.java	Tests defaults + namespace isolation between Online Producer and System Producer configs.
integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerConfig.java	Adds typed accessors/defaults for new STREAM dispatcher config keys.
integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java	Converts STREAM writes to dispatcher-based execution with drain-before-close lifecycle and submission waiting semantics.
clients/venice-producer/src/test/java/com/linkedin/venice/producer/online/OnlineVeniceProducerTest.java	Adds tests for writer hook forwarding and config namespace isolation from system-producer settings.
clients/venice-producer/src/test/java/com/linkedin/venice/producer/AbstractVeniceProducerAsyncBehaviorTest.java	Adds test ensuring callback executor remains alive through writer close for durable completion.
clients/venice-producer/src/main/java/com/linkedin/venice/producer/PartitionedProducerExecutor.java	Replaces implementation with a deprecated shim extending the shared executor for compatibility.
clients/venice-producer/src/main/java/com/linkedin/venice/producer/online/OnlineVeniceProducer.java	Adds optional writer hook plumbing and makes close idempotent/guarded against concurrent cleanup.
clients/venice-producer/src/main/java/com/linkedin/venice/producer/online/OnlineProducerFactory.java	Adds overload to supply an optional VeniceWriterHook to the producer.
clients/venice-producer/src/main/java/com/linkedin/venice/producer/AbstractVeniceProducer.java	Integrates writer hook into writer options, adds stronger close sequencing (workers drain before writer close; callbacks drain after), and rejection-safe callback completion.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 2 comments.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 2 comments.

Suppressed comments (1)

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:318

• A full stripe makes this loop busy-spin for the entire duration of the blocked write (and flush() has no deadline), consuming a CPU core while merely waiting for queue capacity. Use a bounded park/wait between retries; the loop can still re-check sticky failure, interruption, and the stop deadline on each wake-up.

        Thread.yield();

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (4)

Previously missed (1) — in code that hasn't changed since the last review.

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:318

• When a stripe queue is full, flush()/stop() busy-spin here until capacity appears (potentially indefinitely for flush()). Thread.yield() does not guarantee descheduling, so saturation can consume a CPU core and starve the worker needed to drain the queue. Use the existing poll interval for a short interruptible/parked backoff between nonblocking attempts.

        Thread.yield();

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:534

• The new executor validates these settings only here, after clients and veniceWriter have already been created. For example, the default worker count with venice.system.producer.worker.queue.capacity=0 throws from this constructor and start() has no failure cleanup, leaking the writer/clients; because isStarted remains false, retries can leak additional instances. Validate the executor settings before startup side effects or clean up all partially initialized resources on failure.

        this.streamWriteDispatcher = new VeniceSystemProducerWriteDispatcher(
            veniceWriter,
            workerCount,
            workerQueueCapacity,
            callbackThreadCount,
            callbackQueueCapacity,
            storeName);

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:137

• The shutdown deadline does not bound this synchronous writer.flush() call. If PubSub backpressure stalls the flush, stop() remains here indefinitely and never reaches the timeout/forced-worker-shutdown path, despite the configured 60-second shutdown window. The stop path needs to await flush only until deadlineNanos, then record the failure and continue safe forced cleanup.

            flushWriter();

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:327

• The forced-cleanup path also performs an unbounded flush after the shutdown deadline may already have expired. A stalled PubSub flush therefore prevents closeWriter() from running and defeats forced shutdown precisely on the failure path. Skip or deadline-bound this best-effort flush so writer cleanup can proceed.

      flushWriter();

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:471

• A failed STREAM shutdown leaves isStarted false while cleanupComplete is false and the old dispatcher/writer may still be active. This check therefore allows a subsequent start() to overwrite those references, leaking the old workers/writer and permitting two producer lifecycles to run concurrently. Refuse restart until the previous dispatcher reports complete cleanup (or finish that cleanup before reinitializing).

      if (this.isStarted) {
        return;
      }
      this.cleanupComplete = false;

clients/venice-producer/src/main/java/com/linkedin/venice/producer/AbstractVeniceProducer.java:350

• shutdownWorkersNow() invokes this rejection callback for active tasks, but interruption does not guarantee that such a task stops before registering or receiving its PubSub callback. The rejection path can therefore record a failed request, followed by the late PubSub callback unconditionally recording another success/failure and decrementing the pending-operation gauge a second time. Coordinate terminal ownership across both paths—for example, only record terminal metrics when durableWriteFuture.complete* wins.

      }, rejection -> completeRejectedWriteFutureExceptionally(durableWriteFuture, rejection));

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:327

• Thread.yield() does not block, so while a stripe queue remains full this loop continuously retries and can consume a CPU core for the entire flush (unbounded) or shutdown timeout. Use a bounded park/wait between attempts while retaining the existing interruption and sticky-failure checks.

        Thread.yield();

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:471

• A failed start() or a timed-out stop() leaves isStarted false while old clients/writer/dispatcher resources may still be live, but the next start() proceeds and overwrites those handles at lines 526-534. The old dispatcher can then no longer be retried or stopped, leaking its writer and worker threads. Track initialization/cleanup state and reject or clean up a restart unless the prior lifecycle completed successfully.

      if (this.isStarted) {
        return;
      }
      this.cleanupComplete = false;

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:297

• This isFenceHeld() check is a TOCTOU race. With callbacks inline, a callback can observe no fence, then a concurrent flush() can acquire the fence and enter writer.flush() before directCompletion runs. If completing the future invokes a continuation that calls flush() or stop(), that callback blocks on the fence while the writer flush may be waiting for the callback to return, causing a deadlock. The direct-completion decision needs coordination that is atomic with fence acquisition, or the completion must be handed off when that guarantee cannot be made.

    if (!executor.isCallbackExecutorEnabled() && lifecycle.isFenceHeld()) {
      fallbackHandoff.run();
      return;
    }
    if (!executor.tryExecuteCallback(directCompletion, ignored -> fallbackHandoff.run())) {
      fallbackHandoff.run();

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated no new comments.

Suppressed comments (2)

Previously missed (1) — in code that hasn't changed since the last review.

internal/venice-common/src/main/java/com/linkedin/venice/writer/PartitionedVeniceWriteExecutor.java:230

• When both pools are disabled and no inline task is active, awaitWorkerTermination(0, NANOSECONDS) succeeds, but this remainingNanos > 0 check forces awaitTermination to return false without checking that callbacks are already terminated. The compatibility shim previously returned true in this fully inline state, and a zero timeout should still report resources that are already terminated.

    long remainingNanos = deadlineNanos - System.nanoTime();
    return remainingNanos > 0 && awaitCallbackTermination(remainingNanos, TimeUnit.NANOSECONDS);

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:471

• A failed STREAM shutdown sets isStarted to false while leaving cleanupComplete false so that stop() can retry physical cleanup. This guard nevertheless allows start() immediately afterward, which replaces veniceWriter and streamWriteDispatcher; the still-running old dispatcher is then orphaned and can no longer be cleaned up. Reject restart while prior STREAM cleanup is incomplete (or finish that cleanup before constructing replacement resources).

      if (this.isStarted) {
        return;
      }
      this.cleanupComplete = false;

[Copilot]

Pull request overview

Copilot reviewed 24 out of 24 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:553

• isStarted now remains false until the very end, but startup has no rollback path. If any step after creating clients or the writer fails (for example, monitor startup), a subsequent start() is allowed and overwrites those fields without closing the first attempt's resources; a failed attempt can therefore leak writer/client threads and connections. Add failure cleanup before rethrowing, or retain an explicit failed-start state that requires cleanup before retry.

      this.isStarted = true;

[Copilot]

Pull request overview

Copilot reviewed 27 out of 27 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:1029

• This repeats the STREAM shutdown race from send: after isStarted was observed true, stop() may null the dispatcher and writer before this branch, causing getInternalWriter().flush() to dereference null instead of participating in the dispatcher's fence. Select STREAM mode from pushType and use one captured dispatcher reference; an old dispatcher will reject/serialize with stop, while a missing dispatcher can produce the intended stopped-producer error.

    if (streamWriteDispatcher == null) {
      getInternalWriter().flush();
      return;
    }
    streamWriteDispatcher.flush();

[Copilot]

Pull request overview

Copilot reviewed 27 out of 27 changed files in this pull request and generated 1 comment.

Suppressed comments (3)

Previously missed (3) — in code that hasn't changed since the last review.

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducerWriteDispatcher.java:147

• stop() passes its remaining global deadline to an API that can wait that duration twice: once gracefully and again after shutdownNow() (shutdownWorkersAndAwait(timeout, unit) delegates the same timeout to both phases). An unresponsive active task can therefore keep this call blocked for nearly 120 seconds despite the dispatcher's 60-second shutdown deadline. Use a deadline-aware shutdown path whose forced phase receives only the time still remaining from the original deadline.

        workersTerminated = executor.shutdownWorkersAndAwait(remainingNanos(deadlineNanos), TimeUnit.NANOSECONDS);

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:178

• This mutable dispatcher is read without the lifecycle lock and is repeatedly dereferenced by send, put/delete, and flush, while cleanup assigns it to null. A call that already observed isStarted == true can be overtaken by stop(), then either dereference null or incorrectly switch to sendInline after veniceWriter was cleared. Capture one stable dispatcher reference for each operation and use it consistently for failure checks, dispatch, submission waiting, and flushing (or otherwise synchronize removal with those operations).

  private VeniceSystemProducerWriteDispatcher streamWriteDispatcher;

integrations/venice-samza/src/main/java/com/linkedin/venice/samza/VeniceSystemProducer.java:619

• Acquiring lifecycleLock without a deadline means a concurrent stop() can spend the full shutdown interval waiting for another start/stop and then begin an additional cleanup interval. This bypasses the bounded per-call shutdown behavior implemented by the dispatcher. Compute the stop deadline before locking, use timed interrupt-preserving acquisition, and pass only the remaining budget into cleanup.

    lifecycleLock.lock();

[Copilot]

Pull request overview

Copilot reviewed 27 out of 27 changed files in this pull request and generated no new comments.