From 6b94b8ba32459b67dd68094a78eec9fa5c912438 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Fri, 21 Aug 2026 15:17:26 +0200 Subject: [PATCH 01/17] Integration client --- build.gradle.kts | 1 + gradle/libs.versions.toml | 25 ++ integration-client/build.gradle.kts | 114 ++++++ .../restate/integration/AbstractProducer.java | 379 ++++++++++++++++++ .../restate/integration/AuthInterceptor.java | 43 ++ .../integration/ExactlyOnceProducer.java | 78 ++++ .../integration/ExactlyOnceProducerImpl.java | 45 +++ .../restate/integration/IngressEndpoint.java | 61 +++ .../integration/IntegrationClient.java | 88 ++++ .../IntegrationClientException.java | 42 ++ .../integration/IntegrationClientImpl.java | 97 +++++ .../dev/restate/integration/Invocation.java | 76 ++++ .../restate/integration/InvocationImpl.java | 156 +++++++ .../integration/InvocationMetadata.java | 64 +++ .../integration/InvocationMetadataImpl.java | 159 ++++++++ .../dev/restate/integration/Producer.java | 82 ++++ .../dev/restate/integration/ProducerBase.java | 71 ++++ .../dev/restate/integration/ProducerImpl.java | 34 ++ .../ProducerNotReadyException.java | 21 + .../dev/restate/integration/SendResult.java | 22 + .../ingress/ingestion/v1/ingestion_svc.proto | 283 +++++++++++++ .../integration/IntegrationClientTest.java | 320 +++++++++++++++ settings.gradle.kts | 1 + 23 files changed, 2262 insertions(+) create mode 100644 integration-client/build.gradle.kts create mode 100644 integration-client/src/main/java/dev/restate/integration/AbstractProducer.java create mode 100644 integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java create mode 100644 integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java create mode 100644 integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java create mode 100644 integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java create mode 100644 integration-client/src/main/java/dev/restate/integration/IntegrationClient.java create mode 100644 integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java create mode 100644 integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java create mode 100644 integration-client/src/main/java/dev/restate/integration/Invocation.java create mode 100644 integration-client/src/main/java/dev/restate/integration/InvocationImpl.java create mode 100644 integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java create mode 100644 integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java create mode 100644 integration-client/src/main/java/dev/restate/integration/Producer.java create mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerBase.java create mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerImpl.java create mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java create mode 100644 integration-client/src/main/java/dev/restate/integration/SendResult.java create mode 100644 integration-client/src/main/proto/dev/restate/ingress/ingestion/v1/ingestion_svc.proto create mode 100644 integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java diff --git a/build.gradle.kts b/build.gradle.kts index 9bca0b7f..1dda4152 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -68,6 +68,7 @@ val dokkaDocumentedProjects = "examples", "sdk-aggregated-javadocs", "admin-client", + "integration-client", "test-services", ) } diff --git a/gradle/libs.versions.toml b/gradle/libs.versions.toml index c7d28136..05190bab 100644 --- a/gradle/libs.versions.toml +++ b/gradle/libs.versions.toml @@ -16,6 +16,30 @@ bytebuddy = "net.bytebuddy:byte-buddy:1.18.3" objenesis = "org.objenesis:objenesis:3.4" + [libraries.grpc-protobuf] + module = 'io.grpc:grpc-protobuf' + + [libraries.grpc-protobuf.version] + ref = 'grpc' + + [libraries.grpc-stub] + module = 'io.grpc:grpc-stub' + + [libraries.grpc-stub.version] + ref = 'grpc' + + [libraries.grpc-netty-shaded] + module = 'io.grpc:grpc-netty-shaded' + + [libraries.grpc-netty-shaded.version] + ref = 'grpc' + + [libraries.grpc-inprocess] + module = 'io.grpc:grpc-inprocess' + + [libraries.grpc-inprocess.version] + ref = 'grpc' + [libraries.jackson-annotations] module = 'com.fasterxml.jackson.core:jackson-annotations' version = '2.22' @@ -262,6 +286,7 @@ ref = 'ksp' [versions] + grpc = '1.70.0' jackson = '2.22.0' junit = '5.14.1' kotlinx-coroutines = '1.10.2' diff --git a/integration-client/build.gradle.kts b/integration-client/build.gradle.kts new file mode 100644 index 00000000..deb94f76 --- /dev/null +++ b/integration-client/build.gradle.kts @@ -0,0 +1,114 @@ +import com.google.protobuf.gradle.id +import java.io.File +import java.util.concurrent.TimeUnit + +plugins { + `java-library` + `java-conventions` + `library-publishing-conventions` + alias(libs.plugins.protobuf) +} + +description = "Client for the Restate ingress integration (ingestion) API" + +dependencies { + // Generated protobuf messages + gRPC stubs are part of the public API surface. + api(libs.protobuf.java) + api(libs.grpc.protobuf) + api(libs.grpc.stub) + + // Netty transport shipped as the default runtime channel implementation. + implementation(libs.grpc.netty.shaded) + + // grpc-java generated stubs reference javax.annotation.Generated. + compileOnly(libs.tomcat.annotations) + + // @ApiStatus.Experimental markers on the public API. + compileOnly(libs.jetbrains.annotations) + + testImplementation(libs.junit.jupiter) + testImplementation(libs.assertj) + testImplementation(libs.protobuf.java) + // In-process transport to drive the client against a fake IngestionSvc in unit tests. + testImplementation(libs.grpc.inprocess) + testRuntimeOnly(libs.junit.platform.launcher) +} + +// Code generation: protobuf messages + gRPC Java stubs for the ingestion service. +protobuf { + protoc { artifact = "com.google.protobuf:protoc:${libs.versions.protobuf.get()}" } + plugins { id("grpc") { artifact = "io.grpc:protoc-gen-grpc-java:${libs.versions.grpc.get()}" } } + generateProtoTasks { all().forEach { it.plugins { id("grpc") } } } +} + +// Generate a `Version` class at build time. VERSION comes from the Gradle `version`, GIT_HASH from +// git; INTEGRATION is the default `name/version` identity stamped into the ingestion Start frame +// when the caller doesn't provide its own via IntegrationClient.Builder.integration(...). +val generatedVersionDir = layout.buildDirectory.dir("version") + +generatedVersionDir.get().asFile.mkdirs() + +// The protobuf plugin already registers the generated proto/grpc sources; only the version dir +// needs wiring here. +sourceSets { main { java { srcDir(generatedVersionDir) } } } + +// From https://discuss.kotlinlang.org/t/use-git-hash-as-version-number-in-build-gradle-kts/19818/4 +fun String.runCommand( + workingDir: File = File("."), + timeoutAmount: Long = 5, + timeoutUnit: TimeUnit = TimeUnit.SECONDS, +): String = + ProcessBuilder(split("\\s(?=(?:[^'\"`]*(['\"`])[^'\"`]*\\1)*[^'\"`]*$)".toRegex())) + .directory(workingDir) + .redirectOutput(ProcessBuilder.Redirect.PIPE) + .redirectError(ProcessBuilder.Redirect.PIPE) + .start() + .apply { waitFor(timeoutAmount, timeoutUnit) } + .run { + val error = errorStream.bufferedReader().readText().trim() + if (error.isNotEmpty()) { + throw IllegalStateException(error) + } + inputStream.bufferedReader().readText().trim() + } + +val generateVersionClass = + tasks.register("generateVersionClass") { + dependsOn(project.tasks.processResources) + outputs.dir(generatedVersionDir) + + doFirst { + // Tolerate a checkout with no commits yet (git rev-parse fails before the first commit). + val gitHash = + try { + "git rev-parse --short=8 HEAD".runCommand(workingDir = rootDir).ifBlank { "unknown" } + } catch (e: Exception) { + "unknown" + } + val containingDir = generatedVersionDir.get().dir("dev/restate/integration").asFile + assert(containingDir.exists() || containingDir.mkdirs()) + + file("$containingDir/Version.java") + .writeText( + """ + package dev.restate.integration; + + /** Generated at build time by the `generateVersionClass` Gradle task. Do not edit. */ + public final class Version { + private Version() {} + + public static final String VERSION = "$version"; + public static final String GIT_HASH = "$gitHash"; + // Default integration identifier for the ingestion Start frame: `name/version`. + public static final String INTEGRATION = "restate-integration-client/" + VERSION + "_" + GIT_HASH; + } + """ + .trimIndent() + ) + } + } + +tasks { + withType().configureEach { dependsOn(generateVersionClass) } + withType().configureEach { dependsOn(generateVersionClass) } +} diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java new file mode 100644 index 00000000..261f92dc --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java @@ -0,0 +1,379 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import dev.restate.ingestion.v1.DeduplicationMode; +import dev.restate.ingestion.v1.ErrorKind; +import dev.restate.ingestion.v1.IngestionDefaults; +import dev.restate.ingestion.v1.IngestionRequest; +import dev.restate.ingestion.v1.IngestionResponse; +import dev.restate.ingestion.v1.IngestionStart; +import dev.restate.ingestion.v1.IngestionSvcGrpc; +import io.grpc.stub.ClientCallStreamObserver; +import io.grpc.stub.ClientResponseObserver; +import java.util.ArrayList; +import java.util.ConcurrentModificationException; +import java.util.List; +import java.util.Map; +import java.util.TreeMap; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicReference; + +/** + * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and + * the module docs for the concurrency contract. + * + *

There is no client-side record queue: {@link #doSend} writes the record straight to the + * gRPC stream when the producer is ready — Restate's byte send-window has credit ({@code budget}) + * and the transport is writable ({@code callObserver.isReady()}) — or throws {@link + * ProducerNotReadyException} otherwise. In-flight records live on the wire; the only per-record + * client state until commit is a future parked in {@link #ackWaiters}, completed when the ack + * watermark passes its offset. + * + *

Two-tier concurrency: + * + *

    + *
  • A fail-fast, KafkaConsumer-style guard ({@link #acquire()}/{@link #release()}) rejects + * concurrent use from multiple threads; sequential hand-off between threads is fine. + *
  • A single monitor ({@link #lock}) guards the small set of fields genuinely shared between + * the caller thread and gRPC's callback threads. Futures are always completed outside the + * monitor. + *
+ */ +abstract class AbstractProducer implements ProducerBase { + + private final Object lock = new Object(); + + // Set once, synchronously, in beforeStart() before the constructor sends the Start frame. + private volatile ClientCallStreamObserver callObserver; + + // ---- fail-fast single-thread guard ---- + private final AtomicReference owner = new AtomicReference<>(); + private int reentrancy; + + // ---- state guarded by `lock` ---- + private long budget = 0; // remaining Restate send window, in bytes; may go one message negative + private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet + private final List> readyWaiters = new ArrayList<>(); + private final TreeMap>> ackWaiters = new TreeMap<>(); + private boolean closed = false; + private IntegrationClientException failure; + + // Written only by the (guarded) caller thread; never touched by gRPC callbacks. + long lastSent = -1; + + AbstractProducer( + IngestionSvcGrpc.IngestionSvcStub stub, + String producerId, + DeduplicationMode deduplicationMode, + IngestionDefaults defaults, + String integration) { + // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready + // handler. + stub.ingest(new ResponseObserver()); + // Mandatory Start handshake: the first frame on the stream (not flow-controlled). + IngestionRequest start = + IngestionRequest.newBuilder() + .setStart( + IngestionStart.newBuilder() + .setProducerId(producerId) + .setIntegration(integration) + .setDeduplicationMode(deduplicationMode) + .setDefaults(defaults)) + .build(); + synchronized (lock) { + callObserver.onNext(start); + } + } + + // ---- ProducerBase ---- + + @Override + public long lastSentOffset() { + acquire(); + try { + return lastSent; + } finally { + release(); + } + } + + @Override + public CompletableFuture waitReady() { + acquire(); + try { + synchronized (lock) { + if (closed) { + return CompletableFuture.failedFuture(failure); + } + if (isReadyLocked()) { + return CompletableFuture.completedFuture(null); + } + CompletableFuture f = new CompletableFuture<>(); + readyWaiters.add(f); + return f; + } + } finally { + release(); + } + } + + @Override + public CompletableFuture waitAcknowledged(long offset) { + acquire(); + try { + return registerAckWaiter(offset); + } finally { + release(); + } + } + + /** + * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark + * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are + * reentrant, so this is safe to call while already holding {@code lock}). + */ + private CompletableFuture registerAckWaiter(long offset) { + synchronized (lock) { + if (closed) { + return CompletableFuture.failedFuture(failure); + } + if (offset <= lastCommitted) { + return CompletableFuture.completedFuture(lastCommitted); + } + CompletableFuture f = new CompletableFuture<>(); + ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(f); + return f; + } + } + + @Override + public void close() { + acquire(); + try { + terminate( + new IntegrationClientException( + IntegrationClientException.Kind.UNKNOWN, "producer closed"), + true); + } finally { + release(); + } + } + + // ---- send path, shared by the subclasses (caller holds the guard) ---- + + /** + * Send at {@code offset}: write it to the stream now if the producer is ready, else throw {@link + * ProducerNotReadyException}. Returns a future that completes with a {@link SendResult} once the + * record is durably committed by Restate. There is no buffering — a not-ready producer refuses + * rather than parking the record. + */ + final CompletableFuture doSend(long offset, InvocationImpl invocation) + throws ProducerNotReadyException { + IngestionRequest req = + IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build(); + long debit = req.getInvocation().getSerializedSize(); + synchronized (lock) { + if (closed) { + throw new IllegalStateException("producer is closed", failure); + } + if (!isReadyLocked()) { + throw new ProducerNotReadyException("producer is not ready"); + } + writeLocked(req, debit); + lastSent = offset; + CompletableFuture committed = new CompletableFuture<>(); + ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(committed); + return committed.thenApply(watermark -> new SendResultImpl(offset)); + } + } + + // ---- internals (all `*Locked` methods require `lock`) ---- + + private boolean isReadyLocked() { + return !closed && budget > 0 && callObserver.isReady(); + } + + private void writeLocked(IngestionRequest req, long debit) { + callObserver.onNext(req); + budget -= debit; + } + + /** Wake readiness waiters once the stream can accept writes again (window credit + writable). */ + private void wakeReadyWaiters() { + List> wakeReady = null; + synchronized (lock) { + if (closed) { + return; + } + if (isReadyLocked() && !readyWaiters.isEmpty()) { + wakeReady = new ArrayList<>(readyWaiters); + readyWaiters.clear(); + } + } + if (wakeReady != null) { + for (CompletableFuture f : wakeReady) { + f.complete(null); + } + } + } + + private void onResponse(IngestionResponse resp) { + List> acksToComplete = null; + long watermark = -1; + boolean wakeReady = false; + IntegrationClientException err = null; + synchronized (lock) { + if (closed) { + return; + } + if (resp.hasLastCommitted() && resp.getLastCommitted() > lastCommitted) { + lastCommitted = resp.getLastCommitted(); + watermark = lastCommitted; + if (!ackWaiters.isEmpty()) { + acksToComplete = new ArrayList<>(); + Map>> head = ackWaiters.headMap(watermark, true); + for (List> waiters : head.values()) { + acksToComplete.addAll(waiters); + } + head.clear(); + } + } + if (resp.hasWindowUpdate()) { + // increment_bytes is a uint32; read it as unsigned. + budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes()); + wakeReady = true; + } else if (resp.hasError()) { + err = mapError(resp.getError()); + } + } + if (acksToComplete != null) { + for (CompletableFuture f : acksToComplete) { + f.complete(watermark); + } + } + if (err != null) { + terminate(err, false); + } else if (wakeReady) { + wakeReadyWaiters(); + } + } + + /** Mark the producer terminally closed and fail every pending future with {@code cause}. */ + private void terminate(IntegrationClientException cause, boolean halfClose) { + List> ready; + List> acks = new ArrayList<>(); + synchronized (lock) { + if (closed) { + return; + } + closed = true; + failure = cause; + ready = new ArrayList<>(readyWaiters); + readyWaiters.clear(); + for (List> waiters : ackWaiters.values()) { + acks.addAll(waiters); + } + ackWaiters.clear(); + } + if (halfClose) { + ClientCallStreamObserver obs = callObserver; + if (obs != null) { + try { + obs.onCompleted(); + } catch (RuntimeException ignored) { + // Already torn down transport-side; nothing to half-close. + } + } + } + for (CompletableFuture f : ready) { + f.completeExceptionally(cause); + } + for (CompletableFuture f : acks) { + f.completeExceptionally(cause); + } + } + + private static IntegrationClientException mapError(dev.restate.ingestion.v1.Error error) { + String detail = + error.hasInvocationOffset() + ? "[offset=" + error.getInvocationOffset() + "] " + error.getMessage() + : error.getMessage(); + return new IntegrationClientException(mapKind(error.getKind()), detail); + } + + private static IntegrationClientException.Kind mapKind(ErrorKind kind) { + switch (kind) { + case ERROR_KIND_SHUTTING_DOWN: + return IntegrationClientException.Kind.SHUTTING_DOWN; + case ERROR_KIND_GO_AWAY: + return IntegrationClientException.Kind.GO_AWAY; + case ERROR_KIND_NOT_FOUND: + return IntegrationClientException.Kind.NOT_FOUND; + case ERROR_KIND_BAD_REQUEST: + return IntegrationClientException.Kind.BAD_REQUEST; + default: + return IntegrationClientException.Kind.UNKNOWN; + } + } + + // ---- fail-fast guard ---- + + final void acquire() { + Thread current = Thread.currentThread(); + if (owner.get() == current) { + reentrancy++; + return; + } + if (!owner.compareAndSet(null, current)) { + throw new ConcurrentModificationException("Producer is not safe for multi-threaded access"); + } + reentrancy = 1; + } + + final void release() { + if (--reentrancy == 0) { + owner.set(null); + } + } + + private final class ResponseObserver + implements ClientResponseObserver { + @Override + public void beforeStart(ClientCallStreamObserver requestStream) { + callObserver = requestStream; + requestStream.setOnReadyHandler(AbstractProducer.this::wakeReadyWaiters); + } + + @Override + public void onNext(IngestionResponse value) { + onResponse(value); + } + + @Override + public void onError(Throwable t) { + terminate( + new IntegrationClientException( + IntegrationClientException.Kind.UNKNOWN, + "ingestion stream failed: " + t.getMessage(), + t), + false); + } + + @Override + public void onCompleted() { + terminate( + new IntegrationClientException( + IntegrationClientException.Kind.UNKNOWN, "ingestion stream closed by server"), + false); + } + } + + private record SendResultImpl(long offset) implements SendResult {} +} diff --git a/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java b/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java new file mode 100644 index 00000000..6ab9cafe --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java @@ -0,0 +1,43 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import io.grpc.CallOptions; +import io.grpc.Channel; +import io.grpc.ClientCall; +import io.grpc.ClientInterceptor; +import io.grpc.ForwardingClientCall; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; + +/** Adds an {@code Authorization: Bearer } header to every call. */ +final class AuthInterceptor implements ClientInterceptor { + + private static final Metadata.Key AUTHORIZATION = + Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER); + + private final String bearer; + + AuthInterceptor(String token) { + this.bearer = "Bearer " + token; + } + + @Override + public ClientCall interceptCall( + MethodDescriptor method, CallOptions callOptions, Channel next) { + return new ForwardingClientCall.SimpleForwardingClientCall<>( + next.newCall(method, callOptions)) { + @Override + public void start(Listener responseListener, Metadata headers) { + headers.put(AUTHORIZATION, bearer); + super.start(responseListener, headers); + } + }; + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java new file mode 100644 index 00000000..2526606c --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java @@ -0,0 +1,78 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.util.concurrent.CompletableFuture; + +/** + * Like {@link Producer}, but with exactly-once semantics. + * + *

Exactly once

+ * + * Pick a producer id that is stable across restarts and distinct per independent offset + * sequence. E.g., for a Kafka consumer {@code groupId/topic/partition}), for Postgres logical + * replication the slot name. Because deduplication happens on {@code (producerId, offset)}, it is + * then safe to replay from your last checkpoint after a crash: already-committed offsets are + * dropped, and {@link #flush} / {@link #waitAcknowledged(long)} reports how far Restate has durably + * caught up so you can advance the checkpoint. + * + *

Sending

+ * + * {@link #send} writes the record straight to the stream and returns a future that completes once + * Restate has durably committed it. If the producer is not ready {@code send} throws {@link + * ProducerNotReadyException} rather than queueing. Catch it, await {@link #waitReady()}, and retry. + * + *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To + * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in + * bulk. + * + * 

{@code
+ * while (true) {
+ *   try {
+ *     producer.send(lsn, Invocation.create().setBody(payload));
+ *     break;
+ *   } catch (ProducerNotReadyException notReady) {
+ *     producer.waitReady().get();
+ *   }
+ * }
+ * long committed = producer.flush().get();
+ * checkpoint.store(committed);
+ * }
+ * + *

Thread safety

+ * + * A producer is not thread-safe and fails fast with {@link + * java.util.ConcurrentModificationException} if used from more than one thread at once. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface ExactlyOnceProducer extends ProducerBase { + + /** + * Sends an invocation at {@code offset}. + * + *

If the internal buffer is full, or the producer doesn't have enough window credit, sending + * is refused with a {@link ProducerNotReadyException} exception, await {@link #waitReady()}, and + * retry. See the example in {@link ExactlyOnceProducer} for more details. + * + *

The returned future completes when the invocation is durably committed by Restate. + * + * @param offset the offset to assign to this record; must be strictly greater than the previous + * one + * @param invocation the invocation to send + * @return a future completing, once the record is durably committed by Restate, with the {@link + * SendResult} carrying {@code offset} + * @throws ProducerNotReadyException if the producer cannot accept a record right now + * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link + * #lastSentOffset()} + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + CompletableFuture send(long offset, Invocation invocation) + throws ProducerNotReadyException; +} diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java new file mode 100644 index 00000000..9cdad1a7 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java @@ -0,0 +1,45 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import dev.restate.ingestion.v1.DeduplicationMode; +import dev.restate.ingestion.v1.IngestionDefaults; +import dev.restate.ingestion.v1.IngestionSvcGrpc; +import java.util.concurrent.CompletableFuture; + +/** Exactly-once {@link ExactlyOnceProducer}: caller-supplied, strictly-increasing offsets. */ +final class ExactlyOnceProducerImpl extends AbstractProducer implements ExactlyOnceProducer { + + ExactlyOnceProducerImpl( + IngestionSvcGrpc.IngestionSvcStub stub, + String producerId, + IngestionDefaults defaults, + String integration) { + super(stub, producerId, DeduplicationMode.OFFSET_BASED, defaults, integration); + } + + @Override + public CompletableFuture send(long offset, Invocation invocation) + throws ProducerNotReadyException { + acquire(); + try { + checkOffset(offset); + return doSend(offset, (InvocationImpl) invocation); + } finally { + release(); + } + } + + private void checkOffset(long offset) { + if (offset <= lastSent) { + throw new IllegalArgumentException( + "offset must be strictly increasing; last sent " + lastSent + ", got " + offset); + } + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java new file mode 100644 index 00000000..401fc4f6 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java @@ -0,0 +1,61 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.net.URI; + +/** Where to reach the Restate ingestion gRPC endpoint, parsed from an http(s) URL. */ +final class IngressEndpoint { + + final String host; + final int port; + final boolean tls; + + private IngressEndpoint(String host, int port, boolean tls) { + this.host = host; + this.port = port; + this.tls = tls; + } + + /** + * Parse a single ingress URL. {@code https} selects TLS; the port defaults to 443 (TLS) or 80 + * otherwise. + */ + static IngressEndpoint parse(String raw) { + if (raw == null || raw.isBlank()) { + throw new IllegalArgumentException("ingress url must not be empty"); + } + URI uri; + try { + uri = new URI(raw.trim()); + } catch (Exception e) { + throw new IllegalArgumentException("ingress url is not a valid URL: '" + raw + "'", e); + } + boolean tls; + String scheme = uri.getScheme() == null ? null : uri.getScheme().toLowerCase(); + if ("https".equals(scheme)) { + tls = true; + } else if ("http".equals(scheme)) { + tls = false; + } else { + throw new IllegalArgumentException( + "ingress url must use http or https scheme, got '" + + uri.getScheme() + + "' in '" + + raw + + "'"); + } + String host = uri.getHost(); + if (host == null) { + throw new IllegalArgumentException("ingress url has no host: '" + raw + "'"); + } + int port = uri.getPort() != -1 ? uri.getPort() : (tls ? 443 : 80); + return new IngressEndpoint(host, port, tls); + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java new file mode 100644 index 00000000..ea439ca5 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java @@ -0,0 +1,88 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +/** Entry point for producing invocations to Restate ingress over the ingestion API. */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface IntegrationClient extends AutoCloseable { + + /** + * Creates at-least-once {@link Producer} with no stream defaults. + * + * @return a new producer + */ + Producer newProducer(); + + /** + * Creates at-least-once {@link Producer} with the given stream defaults. + * + * @param defaultMetadata invocation fields applied to every record unless overridden per record + * @return a new producer + */ + Producer newProducer(InvocationMetadata defaultMetadata); + + /** + * Creates an {@link ExactlyOnceProducer} identified by {@code producerId}. + * + * @param producerId stable identity of the producer; must be non-empty + * @return a new exactly-once producer + * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank + */ + ExactlyOnceProducer newExactlyOnceProducer(String producerId); + + /** + * Creates an {@link ExactlyOnceProducer} identified by {@code producerId} with the given stream + * defaults. + * + * @param producerId stable identity of the producer; must be non-empty + * @param defaultMetadata invocation fields applied to every record unless overridden per record + * @return a new exactly-once producer + * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank + */ + ExactlyOnceProducer newExactlyOnceProducer(String producerId, InvocationMetadata defaultMetadata); + + /** Shuts down the underlying client. */ + @Override + void close(); + + /** Start building a client that connects to the ingress at {@code target} (an http(s) URL). */ + static Builder builder(String ingressUrl) { + return new Builder(ingressUrl); + } + + /** Builder for {@link IntegrationClient}. */ + final class Builder { + private final String target; + private String authToken; + private String integration = Version.INTEGRATION; + + private Builder(String target) { + this.target = target; + } + + /** Bearer token sent as the {@code Authorization} header on the ingestion stream. */ + public Builder authToken(String authToken) { + this.authToken = authToken; + return this; + } + + /** + * Identify the integration in the ingestion {@code Start} frame as {@code name/version}. When + * not set, defaults to this client's own identity ({@link Version#INTEGRATION}). + */ + public Builder integration(String name, String version) { + this.integration = name + "/" + version; + return this; + } + + public IntegrationClient build() { + return IntegrationClientImpl.create(target, authToken, integration); + } + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java new file mode 100644 index 00000000..6d546e63 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java @@ -0,0 +1,42 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +/** + * The failure that pending producer futures complete with when the ingestion stream errors or + * closes. {@link #getKind()} maps the server's {@code ErrorKind}, or {@link Kind#UNKNOWN} for + * transport-level failures. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public class IntegrationClientException extends RuntimeException { + + /** Classification of an ingestion stream failure. */ + public enum Kind { + UNKNOWN, + SHUTTING_DOWN, + GO_AWAY, + NOT_FOUND, + BAD_REQUEST, + } + + private final Kind kind; + + public IntegrationClientException(Kind kind, String message) { + this(kind, message, null); + } + + public IntegrationClientException(Kind kind, String message, Throwable cause) { + super(message != null ? message : kind.name(), cause); + this.kind = kind; + } + + public Kind getKind() { + return kind; + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java new file mode 100644 index 00000000..11d46c04 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java @@ -0,0 +1,97 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import dev.restate.ingestion.v1.IngestionDefaults; +import dev.restate.ingestion.v1.IngestionSvcGrpc; +import io.grpc.ManagedChannel; +import io.grpc.ManagedChannelBuilder; +import java.util.concurrent.TimeUnit; + +/** {@link IntegrationClient} backed by a single gRPC {@link ManagedChannel} shared by producers. */ +final class IntegrationClientImpl implements IntegrationClient { + + private final ManagedChannel channel; + private final IngestionSvcGrpc.IngestionSvcStub stub; + private final String integration; + + private IntegrationClientImpl( + ManagedChannel channel, IngestionSvcGrpc.IngestionSvcStub stub, String integration) { + this.channel = channel; + this.stub = stub; + this.integration = integration; + } + + static IntegrationClient create(String target, String authToken, String integration) { + IngressEndpoint endpoint = IngressEndpoint.parse(target); + ManagedChannelBuilder builder = + ManagedChannelBuilder.forAddress(endpoint.host, endpoint.port); + if (endpoint.tls) { + builder.useTransportSecurity(); + } else { + builder.usePlaintext(); + } + ManagedChannel channel = builder.build(); + + IngestionSvcGrpc.IngestionSvcStub stub = IngestionSvcGrpc.newStub(channel); + if (authToken != null && !authToken.isBlank()) { + stub = stub.withInterceptors(new AuthInterceptor(authToken)); + } + return new IntegrationClientImpl(channel, stub, integration); + } + + /** Visible for testing: build a client over an already-created channel (e.g. gRPC in-process). */ + static IntegrationClient forChannel(ManagedChannel channel, String integration) { + return new IntegrationClientImpl(channel, IngestionSvcGrpc.newStub(channel), integration); + } + + @Override + public Producer newProducer() { + return newProducer(null); + } + + @Override + public Producer newProducer(InvocationMetadata defaultMetadata) { + return new ProducerImpl(stub, defaultsOf(defaultMetadata), integration); + } + + @Override + public ExactlyOnceProducer newExactlyOnceProducer(String producerId) { + return newExactlyOnceProducer(producerId, null); + } + + @Override + public ExactlyOnceProducer newExactlyOnceProducer( + String producerId, InvocationMetadata defaultMetadata) { + if (producerId == null || producerId.isBlank()) { + throw new IllegalArgumentException( + "producerId must be non-empty for an exactly-once producer"); + } + return new ExactlyOnceProducerImpl(stub, producerId, defaultsOf(defaultMetadata), integration); + } + + @Override + public void close() { + channel.shutdown(); + try { + if (!channel.awaitTermination(5, TimeUnit.SECONDS)) { + channel.shutdownNow(); + } + } catch (InterruptedException e) { + channel.shutdownNow(); + Thread.currentThread().interrupt(); + } + } + + private static IngestionDefaults defaultsOf(InvocationMetadata metadata) { + return metadata == null + ? IngestionDefaults.getDefaultInstance() + : ((InvocationMetadataImpl) metadata).toDefaults(); + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/Invocation.java b/integration-client/src/main/java/dev/restate/integration/Invocation.java new file mode 100644 index 00000000..1460b1ed --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/Invocation.java @@ -0,0 +1,76 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.time.Duration; +import java.time.Instant; +import java.util.Map; + +/** + * A single invocation to send through a {@link Producer} / {@link ExactlyOnceProducer}. + * + *

Instances are created via {@link #create()}. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface Invocation extends InvocationMetadata { + + /** Create a standalone invocation, not bound to any producer, to fill in and send. */ + static Invocation create() { + return new InvocationImpl(); + } + + /** The invocation payload. */ + Invocation setBody(byte[] body); + + byte[] getBody(); + + /** Schedule the invocation after a delay. Mutually exclusive with {@link #setInvokeTime}. */ + Invocation setDelay(Duration delay); + + Duration getDelay(); + + /** Schedule the invocation at an absolute time. Mutually exclusive with {@link #setDelay}. */ + Invocation setInvokeTime(Instant invokeTime); + + Instant getInvokeTime(); + + /** W3C {@code traceparent}. */ + Invocation setTraceparent(String traceparent); + + String getTraceparent(); + + /** W3C {@code tracestate}. */ + Invocation setTracestate(String tracestate); + + String getTracestate(); + + @Override + Invocation setServiceName(String serviceName); + + @Override + Invocation setHandlerName(String handlerName); + + @Override + Invocation setKey(String key); + + @Override + Invocation setScope(String scope); + + @Override + Invocation setLimitKey(String limitKey); + + @Override + Invocation setIdempotencyKey(String idempotencyKey); + + @Override + Invocation putHeader(String key, String value); + + @Override + Invocation setHeaders(Map headers); +} diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java new file mode 100644 index 00000000..095b03aa --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java @@ -0,0 +1,156 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import com.google.protobuf.ByteString; +import dev.restate.ingestion.v1.IngestionInvocation; +import java.time.Duration; +import java.time.Instant; +import java.util.Map; + +/** + * Mutable {@link Invocation} backed directly by the {@link IngestionInvocation.Builder} inherited + * from {@link InvocationMetadataImpl}; {@link #toProtoInvocation(long)} just stamps the offset and + * builds, with no field copying. + */ +class InvocationImpl extends InvocationMetadataImpl implements Invocation { + + @Override + public Invocation setBody(byte[] body) { + if (body == null) { + builder.clearPayload(); + } else { + builder.setPayload(ByteString.copyFrom(body)); + } + return this; + } + + @Override + public byte[] getBody() { + return builder.getPayload().toByteArray(); + } + + @Override + public Invocation setDelay(Duration delay) { + if (delay == null) { + builder.clearDelayMs(); + } else { + builder.setDelayMs(delay.toMillis()); + builder.clearInvokeTimeTsMs(); // mutually exclusive + } + return this; + } + + @Override + public Duration getDelay() { + return builder.hasDelayMs() ? Duration.ofMillis(builder.getDelayMs()) : null; + } + + @Override + public Invocation setInvokeTime(Instant invokeTime) { + if (invokeTime == null) { + builder.clearInvokeTimeTsMs(); + } else { + builder.setInvokeTimeTsMs(invokeTime.toEpochMilli()); + builder.clearDelayMs(); // mutually exclusive + } + return this; + } + + @Override + public Instant getInvokeTime() { + return builder.hasInvokeTimeTsMs() ? Instant.ofEpochMilli(builder.getInvokeTimeTsMs()) : null; + } + + @Override + public Invocation setTraceparent(String traceparent) { + if (traceparent == null) { + builder.clearTraceparent(); + } else { + builder.setTraceparent(traceparent); + } + return this; + } + + @Override + public String getTraceparent() { + return builder.hasTraceparent() ? builder.getTraceparent() : null; + } + + @Override + public Invocation setTracestate(String tracestate) { + if (tracestate == null) { + builder.clearTracestate(); + } else { + builder.setTracestate(tracestate); + } + return this; + } + + @Override + public String getTracestate() { + return builder.hasTracestate() ? builder.getTracestate() : null; + } + + // Covariant overrides so per-invocation chaining keeps the Invocation type. The mutation logic + // lives once in InvocationMetadataImpl (against the shared builder); these only refine the type. + + @Override + public Invocation setServiceName(String serviceName) { + super.setServiceName(serviceName); + return this; + } + + @Override + public Invocation setHandlerName(String handlerName) { + super.setHandlerName(handlerName); + return this; + } + + @Override + public Invocation setKey(String key) { + super.setKey(key); + return this; + } + + @Override + public Invocation setScope(String scope) { + super.setScope(scope); + return this; + } + + @Override + public Invocation setLimitKey(String limitKey) { + super.setLimitKey(limitKey); + return this; + } + + @Override + public Invocation setIdempotencyKey(String idempotencyKey) { + super.setIdempotencyKey(idempotencyKey); + return this; + } + + @Override + public Invocation putHeader(String key, String value) { + super.putHeader(key, value); + return this; + } + + @Override + public Invocation setHeaders(Map headers) { + super.setHeaders(headers); + return this; + } + + /** Stamp the producer-assigned offset and build the wire message (no field copying). */ + IngestionInvocation toProtoInvocation(long offset) { + return builder.setOffset(offset).build(); + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java new file mode 100644 index 00000000..608807ba --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java @@ -0,0 +1,64 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.util.Map; + +/** + * Invocation metadata. + * + *

When used as producer defaults, these will be used for all invocations sent through that + * producer. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface InvocationMetadata { + + /** Create a standalone metadata object, e.g. to use as producer defaults. */ + static InvocationMetadata create() { + return new InvocationMetadataImpl(); + } + + /** Target service name. */ + InvocationMetadata setServiceName(String serviceName); + + String getServiceName(); + + /** Target handler name. */ + InvocationMetadata setHandlerName(String handlerName); + + String getHandlerName(); + + /** Target key (required when the target is a Virtual Object or Workflow). */ + InvocationMetadata setKey(String key); + + String getKey(); + + /** Scope. */ + InvocationMetadata setScope(String scope); + + String getScope(); + + /** Rate/concurrency limit key. */ + InvocationMetadata setLimitKey(String limitKey); + + String getLimitKey(); + + /** Idempotency key used by Restate to deduplicate the invocation. */ + InvocationMetadata setIdempotencyKey(String idempotencyKey); + + String getIdempotencyKey(); + + /** Add or replace a single header. */ + InvocationMetadata putHeader(String key, String value); + + /** Replace the whole header map. */ + InvocationMetadata setHeaders(Map headers); + + Map getHeaders(); +} diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java new file mode 100644 index 00000000..432552f9 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java @@ -0,0 +1,159 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import dev.restate.ingestion.v1.IngestionDefaults; +import dev.restate.ingestion.v1.IngestionInvocation; +import java.util.Map; + +/** + * Mutable {@link InvocationMetadata} backed directly by an {@link IngestionInvocation.Builder}, so + * setters write straight to the wire object with no intermediate field copy. {@link InvocationImpl} + * extends this and reuses the same builder; when the object is used as producer defaults, {@link + * #toDefaults()} projects the shared fields onto an {@link IngestionDefaults}. + */ +class InvocationMetadataImpl implements InvocationMetadata { + + final IngestionInvocation.Builder builder = IngestionInvocation.newBuilder(); + + @Override + public InvocationMetadata setServiceName(String serviceName) { + if (serviceName == null) { + builder.clearService(); + } else { + builder.setService(serviceName); + } + return this; + } + + @Override + public String getServiceName() { + return builder.hasService() ? builder.getService() : null; + } + + @Override + public InvocationMetadata setHandlerName(String handlerName) { + if (handlerName == null) { + builder.clearHandler(); + } else { + builder.setHandler(handlerName); + } + return this; + } + + @Override + public String getHandlerName() { + return builder.hasHandler() ? builder.getHandler() : null; + } + + @Override + public InvocationMetadata setKey(String key) { + if (key == null) { + builder.clearKey(); + } else { + builder.setKey(key); + } + return this; + } + + @Override + public String getKey() { + return builder.hasKey() ? builder.getKey() : null; + } + + @Override + public InvocationMetadata setScope(String scope) { + if (scope == null) { + builder.clearScope(); + } else { + builder.setScope(scope); + } + return this; + } + + @Override + public String getScope() { + return builder.hasScope() ? builder.getScope() : null; + } + + @Override + public InvocationMetadata setLimitKey(String limitKey) { + if (limitKey == null) { + builder.clearLimitKey(); + } else { + builder.setLimitKey(limitKey); + } + return this; + } + + @Override + public String getLimitKey() { + return builder.hasLimitKey() ? builder.getLimitKey() : null; + } + + @Override + public InvocationMetadata setIdempotencyKey(String idempotencyKey) { + if (idempotencyKey == null) { + builder.clearIdempotencyKey(); + } else { + builder.setIdempotencyKey(idempotencyKey); + } + return this; + } + + @Override + public String getIdempotencyKey() { + return builder.hasIdempotencyKey() ? builder.getIdempotencyKey() : null; + } + + @Override + public InvocationMetadata putHeader(String key, String value) { + builder.putAdditionalHeaders(key, value); + return this; + } + + @Override + public InvocationMetadata setHeaders(Map headers) { + builder.clearAdditionalHeaders(); + if (headers != null) { + builder.putAllAdditionalHeaders(headers); + } + return this; + } + + @Override + public Map getHeaders() { + return builder.getAdditionalHeadersMap(); + } + + /** Project the shared fields onto an {@code IngestionDefaults} for the producer Start frame. */ + IngestionDefaults toDefaults() { + IngestionDefaults.Builder d = IngestionDefaults.newBuilder(); + if (builder.hasService()) { + d.setService(builder.getService()); + } + if (builder.hasHandler()) { + d.setHandler(builder.getHandler()); + } + if (builder.hasKey()) { + d.setKey(builder.getKey()); + } + if (builder.hasScope()) { + d.setScope(builder.getScope()); + } + if (builder.hasLimitKey()) { + d.setLimitKey(builder.getLimitKey()); + } + if (builder.hasIdempotencyKey()) { + d.setIdempotencyKey(builder.getIdempotencyKey()); + } + d.putAllHeaders(builder.getAdditionalHeadersMap()); + return d.build(); + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java new file mode 100644 index 00000000..f68443f8 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/Producer.java @@ -0,0 +1,82 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.util.concurrent.CompletableFuture; + +/** + * An at-least-once producer: the client assigns a monotonically increasing offset to each + * invocation. Deduplication is disabled (empty producer id); add an idempotency key on the + * invocations if you need handler-level dedup. + * + *

Sending

+ * + * {@link #send} writes the record straight to the stream and returns a future that completes once + * Restate has durably committed it. If the producer is not ready {@code send} throws {@link + * ProducerNotReadyException} rather than queueing. Catch it, await {@link #waitReady()}, and retry. + * + *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To + * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in + * bulk. + * + * 

{@code
+ * try (IntegrationClient client = IntegrationClient.builder("http://localhost:8080").build();
+ *     Producer producer = client.newProducer()) {
+ *   for (byte[] payload : payloads) {
+ *     Invocation invocation = Invocation.create().setBody(payload);
+ *     while (true) {
+ *       try {
+ *         producer.send(invocation);
+ *         break;
+ *       } catch (ProducerNotReadyException notReady) {
+ *         producer.waitReady().get(); // block until there is capacity, then retry
+ *       }
+ *     }
+ *   }
+ *   producer.flush().get(); // block until everything sent so far is durably committed
+ * }
+ * }
+ * + *

Stream defaults

+ * + * Pass an {@link InvocationMetadata} to {@link IntegrationClient#newProducer(InvocationMetadata)} + * to set fields shared by every record (e.g. the target service/handler) once; per-invocation + * fields override them. + * + * 
{@code
+ * Producer producer =
+ *     client.newProducer(
+ *         InvocationMetadata.create().setServiceName("Greeter").setHandlerName("greet"));
+ * }
+ * + *

Thread safety

+ * + * A producer is not thread-safe and fails fast with {@link + * java.util.ConcurrentModificationException}) if used from more than one thread at once. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface Producer extends ProducerBase { + + /** + * Sends an invocation. + * + *

If the internal buffer is full, or the producer doesn't have enough window credit, sending + * is refused with a {@link ProducerNotReadyException} exception, await {@link #waitReady()}, and + * retry. See the example in {@link Producer} for more details. + * + *

The returned future completes when the invocation is durably committed by Restate. + * + * @param invocation the invocation to send + * @return a future completing, once the invocation is durably committed by Restate. + * @throws ProducerNotReadyException if the producer cannot accept a record right now + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + CompletableFuture send(Invocation invocation) throws ProducerNotReadyException; +} diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java new file mode 100644 index 00000000..b087633b --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java @@ -0,0 +1,71 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import java.util.concurrent.CompletableFuture; + +/** + * @see Producer + * @see ExactlyOnceProducer + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface ProducerBase extends AutoCloseable { + + /** + * Returns the highest offset handed to {@code send} so far. + * + * @return the highest offset sent, or {@code -1} if nothing has been sent yet + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + long lastSentOffset(); + + /** + * Awaits capacity to send another invocation. Await this after {@code send} throws {@link + * ProducerNotReadyException}, then retry the send. + * + * @return a future completing once the producer can accept more invocations. + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + CompletableFuture waitReady(); + + /** + * Awaits durable acknowledgement of all invocations up to and including {@code offset}. + * + * @param offset the offset to wait for + * @return a future completing, once every invocation up to and including {@code offset} is + * durably acknowledged by Restate, with the highest acknowledged offset + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + CompletableFuture waitAcknowledged(long offset); + + /** + * Awaits durable acknowledgement of every invocation sent so far. + * + * @return a future completing, once all invocations sent so far are durably acknowledged by + * Restate, with the highest durably committed offset + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + default CompletableFuture flush() { + return waitAcknowledged(lastSentOffset()); + } + + /** + * Closes the producer and shuts down its stream. Any not-yet-acknowledged invocation completes + * its future exceptionally with an {@link IntegrationClientException}. + * + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + @Override + void close(); +} diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java new file mode 100644 index 00000000..aa787506 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java @@ -0,0 +1,34 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import dev.restate.ingestion.v1.DeduplicationMode; +import dev.restate.ingestion.v1.IngestionDefaults; +import dev.restate.ingestion.v1.IngestionSvcGrpc; +import java.util.concurrent.CompletableFuture; + +/** At-least-once {@link Producer}: dedup disabled, client-assigned monotonic offsets. */ +final class ProducerImpl extends AbstractProducer implements Producer { + + ProducerImpl( + IngestionSvcGrpc.IngestionSvcStub stub, IngestionDefaults defaults, String integration) { + super(stub, "", DeduplicationMode.DISABLED, defaults, integration); + } + + @Override + public CompletableFuture send(Invocation invocation) + throws ProducerNotReadyException { + acquire(); + try { + return doSend(lastSent + 1, (InvocationImpl) invocation); + } finally { + release(); + } + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java b/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java new file mode 100644 index 00000000..b1c36524 --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java @@ -0,0 +1,21 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +/** + * Thrown by {@code send} when the producer cannot accept a record right now (the send window is + * depleted or the transport is not writable). Unchecked: catch it to pace, or await {@link + * ProducerBase#waitReady()} before sending. + */ +@org.jetbrains.annotations.ApiStatus.Experimental +public class ProducerNotReadyException extends RuntimeException { + public ProducerNotReadyException(String message) { + super(message); + } +} diff --git a/integration-client/src/main/java/dev/restate/integration/SendResult.java b/integration-client/src/main/java/dev/restate/integration/SendResult.java new file mode 100644 index 00000000..3ccbca7b --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/SendResult.java @@ -0,0 +1,22 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +/** Send result metadata. */ +@org.jetbrains.annotations.ApiStatus.Experimental +public interface SendResult { + + /** + * Returns the offset assigned to this invocation within the producer's sequence. For an {@link + * ExactlyOnceProducer} this is the offset you supplied; otherwise it is autogenerated. + * + * @return the assigned offset + */ + long offset(); +} diff --git a/integration-client/src/main/proto/dev/restate/ingress/ingestion/v1/ingestion_svc.proto b/integration-client/src/main/proto/dev/restate/ingress/ingestion/v1/ingestion_svc.proto new file mode 100644 index 00000000..cd2960b6 --- /dev/null +++ b/integration-client/src/main/proto/dev/restate/ingress/ingestion/v1/ingestion_svc.proto @@ -0,0 +1,283 @@ +// Copyright (c) 2026 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate service protocol, which is +// released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/proto/blob/main/LICENSE + +syntax = "proto3"; + +package dev.restate.ingress.ingestion.v1; + +// The only local additions to this vendored file are the java_* options below, controlling code +// generation so the generated classes land in dev.restate.ingestion.v1. Re-add them when re-syncing. +option java_multiple_files = true; +option java_package = "dev.restate.ingestion.v1"; +option java_outer_classname = "IngestionProto"; + +// Required. +// A Start frame must be sent before any IngestionDefaults or +// IngestionInvocation frame. +// Sending a Start frame again during the lifetime of +// the stream is illegal: the server will return an +// error and terminate the stream. +// +// Invalid Start message will cause an error +// to be returned immediately, and stream to be teared down +// immediately +message IngestionStart { + // Stable identity of the *producer*, not of this particular stream. + // + // The ID is an opaque string: this API is generic, so the server + // attaches no format or meaning to it. It is entirely up to the producer + // to pick an ID that uniquely identifies a single source of sequential, + // ordered messages, since offsets are only meaningful relative to that + // source. For a Kafka-based producer, for instance, an ID composed of + // cluster + topic + consumer-group + partition would identify such a + // source. + // + // Depending on the DedupMode, the same producer must reuse the same + // ID across streams (including after reconnects), otherwise deduplication + // cannot work: the server tracks committed offsets per producer ID, + // so a new ID starts a fresh deduplication state and previously committed + // records may be duplicated. + // + // The ID must be unique across producers: two distinct producers + // sharing an ID will be treated as one and will shadow each other's + // offsets. + string producer_id = 1; + // Integration identification in format `name/version` + string integration = 2; + // Deduplication strategy for this stream. Defaults to `OFFSET_BASED` + // when unset, which requires `producer_id` to be set. + DeduplicationMode deduplication_mode = 3; + // Initial ingestion defaults for the stream, letting the client + // establish its defaults up front instead of following + // the Start frame with a separate IngestionDefaults frame. + // Clients can still send further IngestionDefaults frames later + // to update these defaults during the stream. + IngestionDefaults defaults = 4; +} + +// How the server deduplicates the records of a stream. +// +// The mode is fixed for the lifetime of the stream: it is chosen once in the +// `Start` frame and cannot be changed by a later `IngestionDefaults` frame. +// +// Note that this only controls deduplication. Regardless of the mode, every +// `Invocation` must still carry an `offset`, offsets must still increase +// monotonically within a stream, and `Response.last_committed` still reports +// progress in terms of those offsets. +enum DeduplicationMode { + // Disable deduplication: every record accepted by the server becomes an + // invocation. + // + // Nothing is tracked across streams, so resending records after a + // reconnect duplicates them. Pick this mode only when the producer cannot + // supply a stable ID, or when duplicates are acceptable or already + // suppressed elsewhere (for example by an `idempotency_key`). + // + // `producer_id` is optional in this mode, but is still recorded on traces + // and metrics when set. + DISABLED = 0; + // Deduplicate on `(producer_id, offset)`. + // + // The server tracks the highest committed offset per producer ID and + // silently drops any record whose offset is not above it, which makes + // resending a range of records after a reconnect safe: records the server + // already committed do not become duplicate invocations. + // + // This mode requires `producer_id` to be set; using it with an empty + // producer ID fails the stream with `GO_AWAY`. It also requires the + // producer to reuse the same ID across streams, since a new ID starts + // from a fresh deduplication state. + // + // This is the default: leaving `deduplication_mode` unset selects it. + OFFSET_BASED = 1; +} +// A IngestionDefaults message defines the default values applied to all +// subsequent records in the stream. +// Each IngestionDefaults message must specify every field it wants to set: +// its fields are not merged with those of previous IngestionDefaults messages. +// Leaving an optional field unset clears any previously established +// default, so later records must supply the value explicitly if needed. +// +// It's okay to send IngestionDefaults more than once to tweak and change the defaults +// during the life time of the stream. +// +// Every field here has a matching field on `Invocation`. For those fields the +// per-record value wins: if the `Invocation` sets the field, its value is used +// and the default from `IngestionDefaults` is ignored; if it leaves the field unset, the +// current default from `IngestionDefaults` is used instead. If neither is set, the field +// is treated as absent, which may be an error for fields that are required for +// the given target (e.g. `key` for virtual objects and workflows). +// +// `headers` is the exception: it is not overridden. `Invocation`'s +// `additional_headers` are appended to the headers established here, so the +// invocation sees both sets. +// +// Note: Keep the field tags uniform with the Invocation message +// This allows parsing IngestionInvocation message as IngestionDefaults. +message IngestionDefaults { + optional string scope = 2; + optional string limit_key = 3; + optional string service = 4; + optional string handler = 5; + optional string key = 6; + optional string idempotency_key = 7; + map headers = 8; +} + +message IngestionInvocation { + // Position of this record within the producer's sequence, and the value + // deduplication is based on. + // + // It must increase monotonically for a given producer ID, across the whole + // lifetime of that producer and not just within a single stream. On + // reconnect, resume from the `Response.last_committed` reported for the + // producer rather than restarting the numbering. + // + // Offsets do not have to be contiguous; gaps are allowed, so it is fine to + // skip records. Only the ordering matters. Offsets are 0-based, so 0 is a + // valid first offset. + // + // Within a single stream the server enforces this: sending an offset lower + // than one already sent on the same stream is a protocol violation and + // fails the stream with `GO_AWAY`. + // + // Across streams nothing is enforced, since the server does not remember + // what a previous stream sent. Reusing an offset already committed by an + // earlier stream of the same producer is not reported as an error: the + // record is treated as a duplicate and silently dropped, meaning it does + // not become an invocation and produces no `Error` frame. A client that + // restarts its numbering after a reconnect will therefore see its records + // disappear without any signal. + uint64 offset = 1; + + // Overrides + optional string scope = 2; + optional string limit_key = 3; + optional string service = 4; + optional string handler = 5; + // Required when the target service is a virtual object (VO). + // or workflow. This can also be set via a preceding IngestionDefaults message. + optional string key = 6; + optional string idempotency_key = 7; + map additional_headers = 8; + + // delay_ms and invoke_time_ts_ms are mutually + // exclusive. + optional uint64 delay_ms = 9; + // starttime timestamp in ms. + optional uint64 invoke_time_ts_ms = 10; + + // This uses the w3c tracecontext format + // https://www.w3.org/TR/trace-context/ + optional string traceparent = 11; + optional string tracestate = 12; + + bytes payload = 13; +} + +message IngestionRequest { + oneof payload { + IngestionStart start = 1; + IngestionDefaults defaults = 2; + IngestionInvocation invocation = 3; + } +} + +// Application-layer flow-control message. +// It tells the client how many more bytes it can send +// before it must wait for the next window update. +// +// WindowUpdate also doubles as an explicit ack: the +// server can send an update with a 0 increment to +// acknowledge commits up to `Response.last_committed` +// without changing the window size. +// +// If the next invocation is bigger than the available +// window, it's legal to send it anyway (letting the +// window go negative) so the client doesn't block. Once +// the window reaches 0 (or less) it's illegal to send +// more requests until the next window update; sending +// more invocations while the window is depleted can +// result in a GO_AWAY error. +message WindowUpdate { + uint32 increment_bytes = 1; +} + +enum ErrorKind { + ERROR_KIND_UNKNOWN = 0; + // The server is shutting down and can't accept more + // records. The caller must switch to a new ingress + // node and respect the reported last_committed. + ERROR_KIND_SHUTTING_DOWN = 1; + // A protocol violation, for example sending another + // Start message after the initial one or not + // respecting the send window. + ERROR_KIND_GO_AWAY = 2; + // The invocation target was not found (unknown + // service or handler). + ERROR_KIND_NOT_FOUND = 3; + // Covers a wide range of client errors, including: + // - unexpected key + // - missing key + // - invalid scope format + // and similar errors. The `Error.message` should + // carry more details. + ERROR_KIND_BAD_REQUEST = 4; +} + +// An error frame. It is always the last frame of the stream: whether the +// error is stream-scoped or record-scoped, the server tears the stream down +// right after sending it and processes no further messages. +// +// Record-scoped errors are the ones that carry `invocation_offset`: the +// record at that offset was rejected (for example a missing key, an unknown +// handler, or an internal ingestion failure) while the stream itself is +// still well-formed. Stream-scoped errors leave `invocation_offset` unset +// and describe a condition of the stream as a whole (shutdown, a protocol +// violation, or invalid `Start`/`IngestionDefaults`). +// +// Even a record-scoped error terminates the stream, on purpose. Skipping the +// faulty record and continuing would commit the records that follow it and +// move the deduplication high water mark past the failed offset, making it +// impossible to ever ingest that record afterwards. Terminating instead +// keeps the high water mark where it is and leaves the decision to the +// client. +// +// Recovery is therefore the client's responsibility: reconnect with the same +// producer ID, resume from the `last_committed` reported on the response +// that carried this error, and skip (or fix) the offending record before +// resending the rest. Note that records after the failed offset may have +// already been sent by the client but were not committed; they must be +// resent. +message Error { + // if set, the error is associated + // with the invocation with that offset + optional uint64 invocation_offset = 1; + ErrorKind kind = 2; + string message = 3; +} + +message IngestionResponse { + // Offsets are 0-based, so last_committed must be optional + // to distinguish "offset 0 committed" from "nothing committed yet". + // This matters when an error is returned on the first + // ingestion message. + optional uint64 last_committed = 1; + + oneof response { + WindowUpdate window_update = 2; + Error error = 3; + } +} + +service IngestionSvc { + // Opens a bidirectional node-to-node stream. + rpc Ingest(stream IngestionRequest) + returns (stream IngestionResponse); +} \ No newline at end of file diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java new file mode 100644 index 00000000..441edf0a --- /dev/null +++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java @@ -0,0 +1,320 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +package dev.restate.integration; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +import dev.restate.ingestion.v1.DeduplicationMode; +import dev.restate.ingestion.v1.ErrorKind; +import dev.restate.ingestion.v1.IngestionInvocation; +import dev.restate.ingestion.v1.IngestionRequest; +import dev.restate.ingestion.v1.IngestionResponse; +import dev.restate.ingestion.v1.IngestionSvcGrpc; +import dev.restate.ingestion.v1.WindowUpdate; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import io.grpc.stub.StreamObserver; +import java.io.IOException; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.BlockingQueue; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.TimeoutException; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Drives the producer client against an in-process fake {@code IngestionSvc}. */ +class IntegrationClientTest { + + private static final String INTEGRATION = "test-integration/1.0"; + + private Server server; + private ManagedChannel channel; + private FakeIngestionService fake; + private IntegrationClient client; + + @BeforeEach + void setUp() throws IOException { + String name = InProcessServerBuilder.generateName(); + fake = new FakeIngestionService(); + server = InProcessServerBuilder.forName(name).directExecutor().addService(fake).build().start(); + channel = InProcessChannelBuilder.forName(name).directExecutor().build(); + client = IntegrationClientImpl.forChannel(channel, INTEGRATION); + } + + @AfterEach + void tearDown() { + if (client != null) { + client.close(); + } + if (server != null) { + server.shutdownNow(); + } + } + + @Test + void producerSendsDisabledDedupHandshake() throws Exception { + InvocationMetadata defaults = InvocationMetadata.create().setServiceName("Svc"); + client.newProducer(defaults); + + IngestionRequest start = fake.take(); + assertThat(start.hasStart()).isTrue(); + assertThat(start.getStart().getProducerId()).isEmpty(); + assertThat(start.getStart().getIntegration()).isEqualTo(INTEGRATION); + assertThat(start.getStart().getDeduplicationMode()).isEqualTo(DeduplicationMode.DISABLED); + assertThat(start.getStart().getDefaults().getService()).isEqualTo("Svc"); + } + + @Test + void exactlyOnceProducerSendsOffsetBasedHandshake() throws Exception { + client.newExactlyOnceProducer("producer-1"); + + IngestionRequest start = fake.take(); + assertThat(start.hasStart()).isTrue(); + assertThat(start.getStart().getProducerId()).isEqualTo("producer-1"); + assertThat(start.getStart().getDeduplicationMode()).isEqualTo(DeduplicationMode.OFFSET_BASED); + } + + @Test + void exactlyOnceProducerRequiresProducerId() { + assertThatThrownBy(() -> client.newExactlyOnceProducer("")) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> client.newExactlyOnceProducer(null)) + .isInstanceOf(IllegalArgumentException.class); + } + + @Test + void producerAssignsMonotonicOffsetsAndFuturesCompleteOnCommit() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + + CompletableFuture a = producer.send(newBody("a")); + CompletableFuture b = producer.send(newBody("b")); + CompletableFuture c = producer.send(newBody("c")); + assertThat(producer.lastSentOffset()).isEqualTo(2L); + + assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L); + assertThat(fake.take().getInvocation().getOffset()).isEqualTo(1L); + assertThat(fake.take().getInvocation().getOffset()).isEqualTo(2L); + + // The send futures resolve on durable commit, each yielding its own offset. + assertThat(a).isNotDone(); + fake.ack(2L); + assertThat(get(a).offset()).isEqualTo(0L); + assertThat(get(b).offset()).isEqualTo(1L); + assertThat(get(c).offset()).isEqualTo(2L); + } + + @Test + void sendThrowsWhenNotReadyThenSucceedsAfterGrant() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + + Invocation inv = newBody("a"); + assertThatThrownBy(() -> producer.send(inv)).isInstanceOf(ProducerNotReadyException.class); + + fake.grantWindow(10_000); + CompletableFuture f = producer.send(newBody("b")); + assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L); + + fake.ack(0L); + assertThat(get(f).offset()).isEqualTo(0L); + } + + @Test + void waitReadyCompletesOnWindowGrant() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + + CompletableFuture ready = producer.waitReady(); + assertThat(ready).isNotDone(); + + fake.grantWindow(10_000); + get(ready); + } + + @Test + void waitAcknowledgedCompletesAtWatermark() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + producer.send(newBody("a")); // offset 0 + producer.send(newBody("b")); // offset 1 + + CompletableFuture acked = producer.waitAcknowledged(1L); + assertThat(acked).isNotDone(); + + fake.ack(0L); + assertThat(acked).isNotDone(); + + fake.ack(1L); + assertThat(get(acked)).isEqualTo(1L); + } + + @Test + void flushCompletesWhenEverythingSentIsCommitted() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + CompletableFuture a = producer.send(newBody("a")); // offset 0 + CompletableFuture b = producer.send(newBody("b")); // offset 1 + + CompletableFuture flushed = producer.flush(); // waits up to the last sent offset (1) + assertThat(a).isNotDone(); + assertThat(flushed).isNotDone(); + + fake.ack(0L); + assertThat(get(a).offset()).isEqualTo(0L); + assertThat(b).isNotDone(); + assertThat(flushed).isNotDone(); + + fake.ack(1L); + assertThat(get(b).offset()).isEqualTo(1L); + assertThat(get(flushed)).isEqualTo(1L); // last durably committed offset + } + + @Test + void streamErrorFailsPendingFuturesFast() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + + CompletableFuture pending = producer.send(newBody("a")); + CompletableFuture acked = producer.waitAcknowledged(0L); + + fake.error(ErrorKind.ERROR_KIND_BAD_REQUEST, "nope"); + + assertThatThrownBy(() -> get(pending)) + .isInstanceOf(ExecutionException.class) + .cause() + .isInstanceOf(IntegrationClientException.class) + .extracting(t -> ((IntegrationClientException) t).getKind()) + .isEqualTo(IntegrationClientException.Kind.BAD_REQUEST); + assertThatThrownBy(() -> get(acked)).isInstanceOf(ExecutionException.class); + + // Subsequent sends fail fast. + assertThatThrownBy(() -> producer.send(newBody("b"))).isInstanceOf(IllegalStateException.class); + } + + @Test + void exactlyOnceRejectsNonIncreasingOffsets() throws Exception { + ExactlyOnceProducer producer = client.newExactlyOnceProducer("p1"); + fake.take(); // Start + fake.grantWindow(10_000); + + producer.send(5L, newBody("a")); + assertThat(producer.lastSentOffset()).isEqualTo(5L); + + assertThatThrownBy(() -> producer.send(5L, newBody("b"))) + .isInstanceOf(IllegalArgumentException.class); + assertThatThrownBy(() -> producer.send(3L, newBody("c"))) + .isInstanceOf(IllegalArgumentException.class); + + producer.send(6L, newBody("d")); + assertThat(producer.lastSentOffset()).isEqualTo(6L); + } + + @Test + void invocationFieldsMapToProto() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + + Invocation inv = + Invocation.create() + .setServiceName("Svc") + .setHandlerName("handle") + .setKey("k") + .setIdempotencyKey("idem") + .putHeader("h1", "v1") + .setTraceparent("tp") + .setBody("hello".getBytes(StandardCharsets.UTF_8)); + producer.send(inv); + + IngestionInvocation sent = fake.take().getInvocation(); + assertThat(sent.getService()).isEqualTo("Svc"); + assertThat(sent.getHandler()).isEqualTo("handle"); + assertThat(sent.getKey()).isEqualTo("k"); + assertThat(sent.getIdempotencyKey()).isEqualTo("idem"); + assertThat(sent.getAdditionalHeadersMap()).containsEntry("h1", "v1"); + assertThat(sent.getTraceparent()).isEqualTo("tp"); + assertThat(sent.getPayload().toStringUtf8()).isEqualTo("hello"); + } + + // ---- helpers ---- + + private static Invocation newBody(String body) { + return Invocation.create().setBody(body.getBytes(StandardCharsets.UTF_8)); + } + + private static T get(CompletableFuture f) + throws InterruptedException, ExecutionException, TimeoutException { + return f.get(5, TimeUnit.SECONDS); + } + + /** Fake service capturing requests and scripting responses. */ + private static final class FakeIngestionService extends IngestionSvcGrpc.IngestionSvcImplBase { + + private final BlockingQueue received = new LinkedBlockingQueue<>(); + private volatile StreamObserver responses; + + @Override + public StreamObserver ingest( + StreamObserver responseObserver) { + this.responses = responseObserver; + return new StreamObserver<>() { + @Override + public void onNext(IngestionRequest value) { + received.add(value); + } + + @Override + public void onError(Throwable t) {} + + @Override + public void onCompleted() {} + }; + } + + IngestionRequest take() throws InterruptedException { + IngestionRequest req = received.poll(5, TimeUnit.SECONDS); + if (req == null) { + throw new AssertionError("timed out waiting for a request frame"); + } + return req; + } + + void grantWindow(long bytes) { + responses.onNext( + IngestionResponse.newBuilder() + .setWindowUpdate(WindowUpdate.newBuilder().setIncrementBytes((int) bytes)) + .build()); + } + + void ack(long lastCommitted) { + responses.onNext(IngestionResponse.newBuilder().setLastCommitted(lastCommitted).build()); + } + + void error(ErrorKind kind, String message) { + responses.onNext( + IngestionResponse.newBuilder() + .setError( + dev.restate.ingestion.v1.Error.newBuilder().setKind(kind).setMessage(message)) + .build()); + responses.onCompleted(); + } + } +} diff --git a/settings.gradle.kts b/settings.gradle.kts index f4a7f99d..f1818ebe 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -18,6 +18,7 @@ include( "common-kotlin", "client", "client-kotlin", + "integration-client", "sdk-common", "sdk-api", "sdk-api-kotlin", From d9909f2981db609ba8a99e49b65faa2c59dd5588 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 10:28:38 +0200 Subject: [PATCH 02/17] Smaller touchups --- .../restate/integration/AbstractProducer.java | 12 +++++++ .../integration/ExactlyOnceProducer.java | 3 +- .../dev/restate/integration/Invocation.java | 2 +- .../restate/integration/InvocationImpl.java | 2 +- .../integration/InvocationMetadata.java | 2 +- .../integration/InvocationMetadataImpl.java | 2 +- .../dev/restate/integration/ProducerBase.java | 13 +++++++ .../integration/IntegrationClientTest.java | 35 +++++++++++++++++++ 8 files changed, 66 insertions(+), 5 deletions(-) diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java index 261f92dc..ddc194ec 100644 --- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java @@ -104,6 +104,18 @@ public long lastSentOffset() { } } + @Override + public long lastAcknowledgedOffset() { + acquire(); + try { + synchronized (lock) { + return lastCommitted; + } + } finally { + release(); + } + } + @Override public CompletableFuture waitReady() { acquire(); diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java index 2526606c..36c6189d 100644 --- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java @@ -20,7 +20,8 @@ * replication the slot name. Because deduplication happens on {@code (producerId, offset)}, it is * then safe to replay from your last checkpoint after a crash: already-committed offsets are * dropped, and {@link #flush} / {@link #waitAcknowledged(long)} reports how far Restate has durably - * caught up so you can advance the checkpoint. + * caught up so you can advance the checkpoint. After a stream failure, {@link + * #lastAcknowledgedOffset()} remains available so you can determine where to resume. * *

Sending

* diff --git a/integration-client/src/main/java/dev/restate/integration/Invocation.java b/integration-client/src/main/java/dev/restate/integration/Invocation.java index 1460b1ed..5a74c6d8 100644 --- a/integration-client/src/main/java/dev/restate/integration/Invocation.java +++ b/integration-client/src/main/java/dev/restate/integration/Invocation.java @@ -18,7 +18,7 @@ *

Instances are created via {@link #create()}. */ @org.jetbrains.annotations.ApiStatus.Experimental -public interface Invocation extends InvocationMetadata { +public sealed interface Invocation extends InvocationMetadata permits InvocationImpl { /** Create a standalone invocation, not bound to any producer, to fill in and send. */ static Invocation create() { diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java index 095b03aa..38654354 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java @@ -19,7 +19,7 @@ * from {@link InvocationMetadataImpl}; {@link #toProtoInvocation(long)} just stamps the offset and * builds, with no field copying. */ -class InvocationImpl extends InvocationMetadataImpl implements Invocation { +final class InvocationImpl extends InvocationMetadataImpl implements Invocation { @Override public Invocation setBody(byte[] body) { diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java index 608807ba..d827f3b5 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java @@ -17,7 +17,7 @@ * producer. */ @org.jetbrains.annotations.ApiStatus.Experimental -public interface InvocationMetadata { +public sealed interface InvocationMetadata permits Invocation, InvocationMetadataImpl { /** Create a standalone metadata object, e.g. to use as producer defaults. */ static InvocationMetadata create() { diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java index 432552f9..add2654d 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java @@ -18,7 +18,7 @@ * extends this and reuses the same builder; when the object is used as producer defaults, {@link * #toDefaults()} projects the shared fields onto an {@link IngestionDefaults}. */ -class InvocationMetadataImpl implements InvocationMetadata { +sealed class InvocationMetadataImpl implements InvocationMetadata permits InvocationImpl { final IngestionInvocation.Builder builder = IngestionInvocation.newBuilder(); diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java index b087633b..51851af2 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java @@ -26,6 +26,19 @@ public interface ProducerBase extends AutoCloseable { */ long lastSentOffset(); + /** + * Returns the highest offset durably acknowledged by Restate. + * + *

This value remains available after the producer closes or fails, so an exactly-once + * producer can use it to determine where to resume. + * + * @return the highest durably acknowledged offset, or {@code -1} if nothing has been acknowledged + * yet + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + long lastAcknowledgedOffset(); + /** * Awaits capacity to send another invocation. Await this after {@code send} throws {@link * ProducerNotReadyException}, then retry the send. diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java index 441edf0a..7bb9bd81 100644 --- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java +++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java @@ -101,6 +101,8 @@ void producerAssignsMonotonicOffsetsAndFuturesCompleteOnCommit() throws Exceptio fake.take(); // Start fake.grantWindow(10_000); + assertThat(producer.lastAcknowledgedOffset()).isEqualTo(-1L); + CompletableFuture a = producer.send(newBody("a")); CompletableFuture b = producer.send(newBody("b")); CompletableFuture c = producer.send(newBody("c")); @@ -116,6 +118,29 @@ void producerAssignsMonotonicOffsetsAndFuturesCompleteOnCommit() throws Exceptio assertThat(get(a).offset()).isEqualTo(0L); assertThat(get(b).offset()).isEqualTo(1L); assertThat(get(c).offset()).isEqualTo(2L); + assertThat(producer.lastAcknowledgedOffset()).isEqualTo(2L); + } + + @Test + void lastAcknowledgedOffsetRemainsAvailableAfterFailure() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + producer.send(newBody("a")); + producer.send(newBody("b")); + + fake.error(ErrorKind.ERROR_KIND_BAD_REQUEST, "nope", 0L); + + assertThat(producer.lastAcknowledgedOffset()).isEqualTo(0L); + } + + @Test + void invocationTypesAreSealedToSdkImplementations() { + assertThat(Invocation.class.getPermittedSubclasses()).containsExactly(InvocationImpl.class); + assertThat(InvocationMetadata.class.getPermittedSubclasses()) + .containsExactlyInAnyOrder(Invocation.class, InvocationMetadataImpl.class); + assertThat(InvocationMetadataImpl.class.getPermittedSubclasses()) + .containsExactly(InvocationImpl.class); } @Test @@ -316,5 +341,15 @@ void error(ErrorKind kind, String message) { .build()); responses.onCompleted(); } + + void error(ErrorKind kind, String message, long lastCommitted) { + responses.onNext( + IngestionResponse.newBuilder() + .setLastCommitted(lastCommitted) + .setError( + dev.restate.ingestion.v1.Error.newBuilder().setKind(kind).setMessage(message)) + .build()); + responses.onCompleted(); + } } } From 63190f123a2cb9a60f97aca8e3f8880d03708a40 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 11:17:01 +0200 Subject: [PATCH 03/17] New API design again --- .../restate/integration/AbstractProducer.java | 259 ++++++++++++------ .../integration/ExactlyOnceProducer.java | 47 ++-- .../integration/ExactlyOnceProducerImpl.java | 16 +- .../integration/IntegrationClient.java | 22 ++ .../integration/IntegrationClientImpl.java | 28 +- .../integration/InvocationMetadataImpl.java | 24 ++ .../dev/restate/integration/Producer.java | 52 ++-- .../dev/restate/integration/ProducerBase.java | 20 +- .../dev/restate/integration/ProducerImpl.java | 15 +- .../ProducerNotReadyException.java | 9 +- .../restate/integration/ProducerOptions.java | 139 ++++++++++ .../dev/restate/integration/SendAttempt.java | 40 +++ .../integration/IntegrationClientTest.java | 169 +++++++++++- 13 files changed, 682 insertions(+), 158 deletions(-) create mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerOptions.java create mode 100644 integration-client/src/main/java/dev/restate/integration/SendAttempt.java diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java index ddc194ec..ad179a1c 100644 --- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java @@ -10,15 +10,17 @@ import dev.restate.ingestion.v1.DeduplicationMode; import dev.restate.ingestion.v1.ErrorKind; -import dev.restate.ingestion.v1.IngestionDefaults; import dev.restate.ingestion.v1.IngestionRequest; import dev.restate.ingestion.v1.IngestionResponse; import dev.restate.ingestion.v1.IngestionStart; import dev.restate.ingestion.v1.IngestionSvcGrpc; import io.grpc.stub.ClientCallStreamObserver; import io.grpc.stub.ClientResponseObserver; +import java.time.Duration; +import java.util.ArrayDeque; import java.util.ArrayList; import java.util.ConcurrentModificationException; +import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.TreeMap; @@ -29,12 +31,9 @@ * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and * the module docs for the concurrency contract. * - *

There is no client-side record queue: {@link #doSend} writes the record straight to the - * gRPC stream when the producer is ready — Restate's byte send-window has credit ({@code budget}) - * and the transport is writable ({@code callObserver.isReady()}) — or throws {@link - * ProducerNotReadyException} otherwise. In-flight records live on the wire; the only per-record - * client state until commit is a future parked in {@link #ackWaiters}, completed when the ack - * watermark passes its offset. + *

Accepted records wait in a byte-bounded queue until Restate's send-window has credit ({@code + * budget}) and the transport is writable ({@code callObserver.isReady()}). Once handed to gRPC, + * only their acknowledgement futures remain until the commit watermark passes their offsets. * *

Two-tier concurrency: * @@ -60,11 +59,17 @@ abstract class AbstractProducer implements ProducerBase { // ---- state guarded by `lock` ---- private long budget = 0; // remaining Restate send window, in bytes; may go one message negative private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet - private final List> readyWaiters = new ArrayList<>(); + private final ArrayDeque bufferedSends = new ArrayDeque<>(); + private long bufferedBytes = 0; + private final List capacityWaiters = new ArrayList<>(); private final TreeMap>> ackWaiters = new TreeMap<>(); private boolean closed = false; private IntegrationClientException failure; + private final long bufferMemory; + private final Duration maxBlockTime; + private final long maxBlockNanos; + // Written only by the (guarded) caller thread; never touched by gRPC callbacks. long lastSent = -1; @@ -72,8 +77,11 @@ abstract class AbstractProducer implements ProducerBase { IngestionSvcGrpc.IngestionSvcStub stub, String producerId, DeduplicationMode deduplicationMode, - IngestionDefaults defaults, + ProducerOptions options, String integration) { + this.bufferMemory = options.bufferMemory(); + this.maxBlockTime = options.maxBlockTime(); + this.maxBlockNanos = toNanosSaturated(maxBlockTime); // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready // handler. stub.ingest(new ResponseObserver()); @@ -85,7 +93,7 @@ abstract class AbstractProducer implements ProducerBase { .setProducerId(producerId) .setIntegration(integration) .setDeduplicationMode(deduplicationMode) - .setDefaults(defaults)) + .setDefaults(options.toDefaults())) .build(); synchronized (lock) { callObserver.onNext(start); @@ -116,26 +124,6 @@ public long lastAcknowledgedOffset() { } } - @Override - public CompletableFuture waitReady() { - acquire(); - try { - synchronized (lock) { - if (closed) { - return CompletableFuture.failedFuture(failure); - } - if (isReadyLocked()) { - return CompletableFuture.completedFuture(null); - } - CompletableFuture f = new CompletableFuture<>(); - readyWaiters.add(f); - return f; - } - } finally { - release(); - } - } - @Override public CompletableFuture waitAcknowledged(long offset) { acquire(); @@ -180,66 +168,166 @@ public void close() { // ---- send path, shared by the subclasses (caller holds the guard) ---- - /** - * Send at {@code offset}: write it to the stream now if the producer is ready, else throw {@link - * ProducerNotReadyException}. Returns a future that completes with a {@link SendResult} once the - * record is durably committed by Restate. There is no buffering — a not-ready producer refuses - * rather than parking the record. - */ + /** Admit a record, blocking up to the configured maximum when the local buffer is full. */ final CompletableFuture doSend(long offset, InvocationImpl invocation) throws ProducerNotReadyException { - IngestionRequest req = - IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build(); - long debit = req.getInvocation().getSerializedSize(); + PreparedSend prepared = prepare(offset, invocation); + List> ready; + CompletableFuture acknowledgement; + long waitStarted = System.nanoTime(); synchronized (lock) { - if (closed) { - throw new IllegalStateException("producer is closed", failure); + ensureOpenLocked(); + while (!hasCapacityLocked(prepared.bufferSize())) { + if (maxBlockNanos == 0) { + throw admissionTimeout(); + } + long remaining = maxBlockNanos - (System.nanoTime() - waitStarted); + if (remaining <= 0) { + throw admissionTimeout(); + } + try { + long millis = remaining / 1_000_000; + int nanos = (int) (remaining % 1_000_000); + lock.wait(millis, nanos); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new ProducerNotReadyException( + "interrupted while waiting for producer buffer capacity", e); + } + ensureOpenLocked(); } - if (!isReadyLocked()) { - throw new ProducerNotReadyException("producer is not ready"); + acknowledgement = acceptLocked(prepared); + ready = drainLocked(); + } + completeReady(ready); + return acknowledgement; + } + + /** Attempt to admit a record without blocking or consuming an offset when capacity is absent. */ + final SendAttempt doTrySend(long offset, InvocationImpl invocation) { + PreparedSend prepared = prepare(offset, invocation); + List> ready; + SendAttempt result; + synchronized (lock) { + ensureOpenLocked(); + if (hasCapacityLocked(prepared.bufferSize())) { + result = new SendAttempt.Accepted(acceptLocked(prepared)); + ready = drainLocked(); + } else { + CompletableFuture future = new CompletableFuture<>(); + CapacityWaiter waiter = new CapacityWaiter(prepared.bufferSize(), future); + capacityWaiters.add(waiter); + future.whenComplete( + (ignored, failure) -> { + if (future.isCancelled()) { + synchronized (lock) { + capacityWaiters.remove(waiter); + } + } + }); + result = new SendAttempt.Backpressured(future); + ready = List.of(); } - writeLocked(req, debit); - lastSent = offset; - CompletableFuture committed = new CompletableFuture<>(); - ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(committed); - return committed.thenApply(watermark -> new SendResultImpl(offset)); } + completeReady(ready); + return result; } // ---- internals (all `*Locked` methods require `lock`) ---- - private boolean isReadyLocked() { - return !closed && budget > 0 && callObserver.isReady(); + private PreparedSend prepare(long offset, InvocationImpl invocation) { + IngestionRequest request = + IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build(); + long bufferSize = request.getSerializedSize(); + if (bufferSize > bufferMemory) { + throw new IllegalArgumentException( + "serialized invocation requires " + + bufferSize + + " bytes, exceeding bufferMemory " + + bufferMemory); + } + return new PreparedSend( + offset, request, request.getInvocation().getSerializedSize(), bufferSize); } - private void writeLocked(IngestionRequest req, long debit) { - callObserver.onNext(req); - budget -= debit; + private boolean hasCapacityLocked(long requiredBytes) { + return requiredBytes <= bufferMemory - bufferedBytes; } - /** Wake readiness waiters once the stream can accept writes again (window credit + writable). */ - private void wakeReadyWaiters() { - List> wakeReady = null; - synchronized (lock) { - if (closed) { - return; - } - if (isReadyLocked() && !readyWaiters.isEmpty()) { - wakeReady = new ArrayList<>(readyWaiters); - readyWaiters.clear(); - } + private CompletableFuture acceptLocked(PreparedSend prepared) { + bufferedSends.addLast( + new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize())); + bufferedBytes += prepared.bufferSize(); + lastSent = prepared.offset(); + CompletableFuture committed = new CompletableFuture<>(); + ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed); + return committed.thenApply(ignored -> new SendResultImpl(prepared.offset())); + } + + /** Write as many queued records as transport and protocol flow control currently permit. */ + private List> drainLocked() { + boolean freedCapacity = false; + while (!closed && budget > 0 && callObserver.isReady() && !bufferedSends.isEmpty()) { + BufferedSend send = bufferedSends.removeFirst(); + bufferedBytes -= send.bufferSize(); + budget -= send.windowDebit(); + freedCapacity = true; + callObserver.onNext(send.request()); } - if (wakeReady != null) { - for (CompletableFuture f : wakeReady) { - f.complete(null); + + if (!freedCapacity) { + return List.of(); + } + + lock.notifyAll(); + long available = bufferMemory - bufferedBytes; + List> ready = new ArrayList<>(); + for (Iterator it = capacityWaiters.iterator(); it.hasNext(); ) { + CapacityWaiter waiter = it.next(); + if (waiter.requiredBytes() <= available) { + ready.add(waiter.future()); + it.remove(); } } + return ready; + } + + private void drainAndWake() { + List> ready; + synchronized (lock) { + ready = drainLocked(); + } + completeReady(ready); + } + + private static void completeReady(List> ready) { + for (CompletableFuture future : ready) { + future.complete(null); + } + } + + private void ensureOpenLocked() { + if (closed) { + throw new IllegalStateException("producer is closed", failure); + } + } + + private ProducerNotReadyException admissionTimeout() { + return new ProducerNotReadyException("producer buffer remained full for " + maxBlockTime); + } + + private static long toNanosSaturated(Duration duration) { + try { + return duration.toNanos(); + } catch (ArithmeticException ignored) { + return Long.MAX_VALUE; + } } private void onResponse(IngestionResponse resp) { List> acksToComplete = null; long watermark = -1; - boolean wakeReady = false; + boolean drain = false; IntegrationClientException err = null; synchronized (lock) { if (closed) { @@ -260,26 +348,26 @@ private void onResponse(IngestionResponse resp) { if (resp.hasWindowUpdate()) { // increment_bytes is a uint32; read it as unsigned. budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes()); - wakeReady = true; + drain = true; } else if (resp.hasError()) { err = mapError(resp.getError()); } } + if (err != null) { + terminate(err, false); + } else if (drain) { + drainAndWake(); + } if (acksToComplete != null) { for (CompletableFuture f : acksToComplete) { f.complete(watermark); } } - if (err != null) { - terminate(err, false); - } else if (wakeReady) { - wakeReadyWaiters(); - } } /** Mark the producer terminally closed and fail every pending future with {@code cause}. */ private void terminate(IntegrationClientException cause, boolean halfClose) { - List> ready; + List> capacity; List> acks = new ArrayList<>(); synchronized (lock) { if (closed) { @@ -287,8 +375,14 @@ private void terminate(IntegrationClientException cause, boolean halfClose) { } closed = true; failure = cause; - ready = new ArrayList<>(readyWaiters); - readyWaiters.clear(); + capacity = new ArrayList<>(capacityWaiters.size()); + for (CapacityWaiter waiter : capacityWaiters) { + capacity.add(waiter.future()); + } + capacityWaiters.clear(); + bufferedSends.clear(); + bufferedBytes = 0; + lock.notifyAll(); for (List> waiters : ackWaiters.values()) { acks.addAll(waiters); } @@ -304,7 +398,7 @@ private void terminate(IntegrationClientException cause, boolean halfClose) { } } } - for (CompletableFuture f : ready) { + for (CompletableFuture f : capacity) { f.completeExceptionally(cause); } for (CompletableFuture f : acks) { @@ -360,7 +454,7 @@ private final class ResponseObserver @Override public void beforeStart(ClientCallStreamObserver requestStream) { callObserver = requestStream; - requestStream.setOnReadyHandler(AbstractProducer.this::wakeReadyWaiters); + requestStream.setOnReadyHandler(AbstractProducer.this::drainAndWake); } @Override @@ -387,5 +481,12 @@ public void onCompleted() { } } + private record PreparedSend( + long offset, IngestionRequest request, long windowDebit, long bufferSize) {} + + private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {} + + private record CapacityWaiter(long requiredBytes, CompletableFuture future) {} + private record SendResultImpl(long offset) implements SendResult {} } diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java index 36c6189d..9a57ee31 100644 --- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java @@ -16,7 +16,7 @@ *

Exactly once

* * Pick a producer id that is stable across restarts and distinct per independent offset - * sequence. E.g., for a Kafka consumer {@code groupId/topic/partition}), for Postgres logical + * sequence. E.g., for a Kafka consumer {@code groupId/topic/partition}, for Postgres logical * replication the slot name. Because deduplication happens on {@code (producerId, offset)}, it is * then safe to replay from your last checkpoint after a crash: already-committed offsets are * dropped, and {@link #flush} / {@link #waitAcknowledged(long)} reports how far Restate has durably @@ -25,23 +25,16 @@ * *

Sending

* - * {@link #send} writes the record straight to the stream and returns a future that completes once - * Restate has durably committed it. If the producer is not ready {@code send} throws {@link - * ProducerNotReadyException} rather than queueing. Catch it, await {@link #waitReady()}, and retry. + * {@link #send} admits the record into a byte-bounded local buffer, waiting up to {@link + * ProducerOptions#maxBlockTime()} for capacity, and returns a future that completes once Restate + * has durably committed it. Use {@link #trySend} when the calling thread must never block. * *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in * bulk. * * 

{@code
- * while (true) {
- *   try {
- *     producer.send(lsn, Invocation.create().setBody(payload));
- *     break;
- *   } catch (ProducerNotReadyException notReady) {
- *     producer.waitReady().get();
- *   }
- * }
+ * producer.send(lsn, Invocation.create().setBody(payload));
  * long committed = producer.flush().get();
  * checkpoint.store(committed);
  * }
@@ -57,9 +50,9 @@ public interface ExactlyOnceProducer extends ProducerBase { /** * Sends an invocation at {@code offset}. * - *

If the internal buffer is full, or the producer doesn't have enough window credit, sending - * is refused with a {@link ProducerNotReadyException} exception, await {@link #waitReady()}, and - * retry. See the example in {@link ExactlyOnceProducer} for more details. + *

If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()} + * for capacity. The invocation is refused with {@link ProducerNotReadyException} if the timeout + * elapses. A zero duration makes this method fail immediately under backpressure. * *

The returned future completes when the invocation is durably committed by Restate. * @@ -68,12 +61,32 @@ public interface ExactlyOnceProducer extends ProducerBase { * @param invocation the invocation to send * @return a future completing, once the record is durably committed by Restate, with the {@link * SendResult} carrying {@code offset} - * @throws ProducerNotReadyException if the producer cannot accept a record right now + * @throws ProducerNotReadyException if buffer capacity does not become available before the + * configured maximum blocking time elapses * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link - * #lastSentOffset()} + * #lastSentOffset()}, or the serialized invocation is larger than {@link + * ProducerOptions#bufferMemory()} * @throws java.util.ConcurrentModificationException if the producer is used concurrently from * another thread */ CompletableFuture send(long offset, Invocation invocation) throws ProducerNotReadyException; + + /** + * Attempts to send an invocation at {@code offset} without blocking. + * + *

An {@link SendAttempt.Accepted} carries the durable-acknowledgement future. A {@link + * SendAttempt.Backpressured} carries a future that completes when retrying may succeed; the + * notification does not reserve capacity. + * + * @param offset the offset to assign; must be strictly greater than the previous accepted offset + * @param invocation the invocation to send + * @return the admission result + * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link + * #lastSentOffset()}, or the serialized invocation is larger than {@link + * ProducerOptions#bufferMemory()} + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + SendAttempt trySend(long offset, Invocation invocation); } diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java index 9cdad1a7..1c7371a7 100644 --- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java @@ -9,7 +9,6 @@ package dev.restate.integration; import dev.restate.ingestion.v1.DeduplicationMode; -import dev.restate.ingestion.v1.IngestionDefaults; import dev.restate.ingestion.v1.IngestionSvcGrpc; import java.util.concurrent.CompletableFuture; @@ -19,9 +18,20 @@ final class ExactlyOnceProducerImpl extends AbstractProducer implements ExactlyO ExactlyOnceProducerImpl( IngestionSvcGrpc.IngestionSvcStub stub, String producerId, - IngestionDefaults defaults, + ProducerOptions options, String integration) { - super(stub, producerId, DeduplicationMode.OFFSET_BASED, defaults, integration); + super(stub, producerId, DeduplicationMode.OFFSET_BASED, options, integration); + } + + @Override + public SendAttempt trySend(long offset, Invocation invocation) { + acquire(); + try { + checkOffset(offset); + return doTrySend(offset, (InvocationImpl) invocation); + } finally { + release(); + } } @Override diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java index ea439ca5..ed6a70af 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java @@ -24,9 +24,19 @@ public interface IntegrationClient extends AutoCloseable { * * @param defaultMetadata invocation fields applied to every record unless overridden per record * @return a new producer + * @throws NullPointerException if {@code defaultMetadata} is {@code null} */ Producer newProducer(InvocationMetadata defaultMetadata); + /** + * Creates an at-least-once {@link Producer} with the given options. + * + * @param options producer buffering, blocking, and invocation-default options + * @return a new producer + * @throws NullPointerException if {@code options} is {@code null} + */ + Producer newProducer(ProducerOptions options); + /** * Creates an {@link ExactlyOnceProducer} identified by {@code producerId}. * @@ -44,9 +54,21 @@ public interface IntegrationClient extends AutoCloseable { * @param defaultMetadata invocation fields applied to every record unless overridden per record * @return a new exactly-once producer * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank + * @throws NullPointerException if {@code defaultMetadata} is {@code null} */ ExactlyOnceProducer newExactlyOnceProducer(String producerId, InvocationMetadata defaultMetadata); + /** + * Creates an {@link ExactlyOnceProducer} identified by {@code producerId} with the given options. + * + * @param producerId stable identity of the producer; must be non-empty + * @param options producer buffering, blocking, and invocation-default options + * @return a new exactly-once producer + * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank + * @throws NullPointerException if {@code options} is {@code null} + */ + ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOptions options); + /** Shuts down the underlying client. */ @Override void close(); diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java index 11d46c04..0431231f 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java @@ -8,10 +8,10 @@ // https://github.com/restatedev/sdk-java/blob/main/LICENSE package dev.restate.integration; -import dev.restate.ingestion.v1.IngestionDefaults; import dev.restate.ingestion.v1.IngestionSvcGrpc; import io.grpc.ManagedChannel; import io.grpc.ManagedChannelBuilder; +import java.util.Objects; import java.util.concurrent.TimeUnit; /** {@link IntegrationClient} backed by a single gRPC {@link ManagedChannel} shared by producers. */ @@ -53,27 +53,39 @@ static IntegrationClient forChannel(ManagedChannel channel, String integration) @Override public Producer newProducer() { - return newProducer(null); + return newProducer(ProducerOptions.defaults()); } @Override public Producer newProducer(InvocationMetadata defaultMetadata) { - return new ProducerImpl(stub, defaultsOf(defaultMetadata), integration); + return newProducer(ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); + } + + @Override + public Producer newProducer(ProducerOptions options) { + return new ProducerImpl(stub, Objects.requireNonNull(options, "options"), integration); } @Override public ExactlyOnceProducer newExactlyOnceProducer(String producerId) { - return newExactlyOnceProducer(producerId, null); + return newExactlyOnceProducer(producerId, ProducerOptions.defaults()); } @Override public ExactlyOnceProducer newExactlyOnceProducer( String producerId, InvocationMetadata defaultMetadata) { + return newExactlyOnceProducer( + producerId, ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); + } + + @Override + public ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOptions options) { if (producerId == null || producerId.isBlank()) { throw new IllegalArgumentException( "producerId must be non-empty for an exactly-once producer"); } - return new ExactlyOnceProducerImpl(stub, producerId, defaultsOf(defaultMetadata), integration); + return new ExactlyOnceProducerImpl( + stub, producerId, Objects.requireNonNull(options, "options"), integration); } @Override @@ -88,10 +100,4 @@ public void close() { Thread.currentThread().interrupt(); } } - - private static IngestionDefaults defaultsOf(InvocationMetadata metadata) { - return metadata == null - ? IngestionDefaults.getDefaultInstance() - : ((InvocationMetadataImpl) metadata).toDefaults(); - } } diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java index add2654d..b01a1ccc 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java @@ -22,6 +22,30 @@ sealed class InvocationMetadataImpl implements InvocationMetadata permits Invoca final IngestionInvocation.Builder builder = IngestionInvocation.newBuilder(); + static InvocationMetadataImpl fromDefaults(IngestionDefaults defaults) { + InvocationMetadataImpl metadata = new InvocationMetadataImpl(); + if (defaults.hasService()) { + metadata.builder.setService(defaults.getService()); + } + if (defaults.hasHandler()) { + metadata.builder.setHandler(defaults.getHandler()); + } + if (defaults.hasKey()) { + metadata.builder.setKey(defaults.getKey()); + } + if (defaults.hasScope()) { + metadata.builder.setScope(defaults.getScope()); + } + if (defaults.hasLimitKey()) { + metadata.builder.setLimitKey(defaults.getLimitKey()); + } + if (defaults.hasIdempotencyKey()) { + metadata.builder.setIdempotencyKey(defaults.getIdempotencyKey()); + } + metadata.builder.putAllAdditionalHeaders(defaults.getHeadersMap()); + return metadata; + } + @Override public InvocationMetadata setServiceName(String serviceName) { if (serviceName == null) { diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java index f68443f8..a65cf7a6 100644 --- a/integration-client/src/main/java/dev/restate/integration/Producer.java +++ b/integration-client/src/main/java/dev/restate/integration/Producer.java @@ -17,9 +17,9 @@ * *

Sending

* - * {@link #send} writes the record straight to the stream and returns a future that completes once - * Restate has durably committed it. If the producer is not ready {@code send} throws {@link - * ProducerNotReadyException} rather than queueing. Catch it, await {@link #waitReady()}, and retry. + * {@link #send} admits the record into a byte-bounded local buffer, waiting up to {@link + * ProducerOptions#maxBlockTime()} for capacity, and returns a future that completes once Restate + * has durably committed it. Use {@link #trySend} when the calling thread must never block. * *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in @@ -29,15 +29,7 @@ * try (IntegrationClient client = IntegrationClient.builder("http://localhost:8080").build(); * Producer producer = client.newProducer()) { * for (byte[] payload : payloads) { - * Invocation invocation = Invocation.create().setBody(payload); - * while (true) { - * try { - * producer.send(invocation); - * break; - * } catch (ProducerNotReadyException notReady) { - * producer.waitReady().get(); // block until there is capacity, then retry - * } - * } + * producer.send(Invocation.create().setBody(payload)); * } * producer.flush().get(); // block until everything sent so far is durably committed * } @@ -45,9 +37,10 @@ * *

Stream defaults

* - * Pass an {@link InvocationMetadata} to {@link IntegrationClient#newProducer(InvocationMetadata)} - * to set fields shared by every record (e.g. the target service/handler) once; per-invocation - * fields override them. + * Pass an {@link InvocationMetadata} to {@link IntegrationClient#newProducer(InvocationMetadata)}, + * or set {@link ProducerOptions.Builder#defaultMetadata(InvocationMetadata)}, to configure fields + * shared by every record (e.g. the target service/handler) once; per-invocation fields override + * them. * * 
{@code
  * Producer producer =
@@ -58,7 +51,7 @@
  * 
Thread safety

  *
  * A producer is not thread-safe and fails fast with {@link
- * java.util.ConcurrentModificationException}) if used from more than one thread at once.
+ * java.util.ConcurrentModificationException} if used from more than one thread at once.
  */
 @org.jetbrains.annotations.ApiStatus.Experimental
 public interface Producer extends ProducerBase {
@@ -66,17 +59,36 @@ public interface Producer extends ProducerBase {
   /**
    * Sends an invocation.
    *
-   * 
If the internal buffer is full, or the producer doesn't have enough window credit, sending
-   * is refused with a {@link ProducerNotReadyException} exception, await {@link #waitReady()}, and
-   * retry. See the example in {@link Producer} for more details.
+   * 
If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()}
+   * for capacity. The invocation is refused with {@link ProducerNotReadyException} if the timeout
+   * elapses. A zero duration makes this method fail immediately under backpressure.
    *
    * 
The returned future completes when the invocation is durably committed by Restate.
    *
    * @param invocation the invocation to send
    * @return a future completing, once the invocation is durably committed by Restate.
-   * @throws ProducerNotReadyException if the producer cannot accept a record right now
+   * @throws ProducerNotReadyException if buffer capacity does not become available before the
+   *     configured maximum blocking time elapses
+   * @throws IllegalArgumentException if the serialized invocation is larger than {@link
+   *     ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
   CompletableFuture send(Invocation invocation) throws ProducerNotReadyException;
+
+  /**
+   * Attempts to send an invocation without blocking.
+   *
+   * 
An {@link SendAttempt.Accepted} carries the durable-acknowledgement future. A {@link
+   * SendAttempt.Backpressured} carries a future that completes when retrying may succeed; the
+   * notification does not reserve capacity.
+   *
+   * @param invocation the invocation to send
+   * @return the admission result
+   * @throws IllegalArgumentException if the serialized invocation is larger than {@link
+   *     ProducerOptions#bufferMemory()}
+   * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
+   *     another thread
+   */
+  SendAttempt trySend(Invocation invocation);
 }
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
index 51851af2..153801bd 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
@@ -11,6 +11,8 @@
 import java.util.concurrent.CompletableFuture;
 
 /**
+ * Common producer offsets, acknowledgement, flushing, and lifecycle operations.
+ *
  * @see Producer
  * @see ExactlyOnceProducer
  */
@@ -18,9 +20,9 @@
 public interface ProducerBase extends AutoCloseable {
 
   /**
-   * Returns the highest offset handed to {@code send} so far.
+   * Returns the highest offset successfully accepted by {@code send} or {@code trySend} so far.
    *
-   * @return the highest offset sent, or {@code -1} if nothing has been sent yet
+   * @return the highest offset accepted, or {@code -1} if nothing has been accepted yet
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
@@ -29,8 +31,8 @@ public interface ProducerBase extends AutoCloseable {
   /**
    * Returns the highest offset durably acknowledged by Restate.
    *
-   * 
This value remains available after the producer closes or fails, so an exactly-once
-   * producer can use it to determine where to resume.
+   * 
This value remains available after the producer closes or fails, so an exactly-once producer
+   * can use it to determine where to resume.
    *
    * @return the highest durably acknowledged offset, or {@code -1} if nothing has been acknowledged
    *     yet
@@ -39,16 +41,6 @@ public interface ProducerBase extends AutoCloseable {
    */
   long lastAcknowledgedOffset();
 
-  /**
-   * Awaits capacity to send another invocation. Await this after {@code send} throws {@link
-   * ProducerNotReadyException}, then retry the send.
-   *
-   * @return a future completing once the producer can accept more invocations.
-   * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
-   *     another thread
-   */
-  CompletableFuture waitReady();
-
   /**
    * Awaits durable acknowledgement of all invocations up to and including {@code offset}.
    *
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
index aa787506..62b2815b 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
@@ -9,7 +9,6 @@
 package dev.restate.integration;
 
 import dev.restate.ingestion.v1.DeduplicationMode;
-import dev.restate.ingestion.v1.IngestionDefaults;
 import dev.restate.ingestion.v1.IngestionSvcGrpc;
 import java.util.concurrent.CompletableFuture;
 
@@ -17,8 +16,8 @@
 final class ProducerImpl extends AbstractProducer implements Producer {
 
   ProducerImpl(
-      IngestionSvcGrpc.IngestionSvcStub stub, IngestionDefaults defaults, String integration) {
-    super(stub, "", DeduplicationMode.DISABLED, defaults, integration);
+      IngestionSvcGrpc.IngestionSvcStub stub, ProducerOptions options, String integration) {
+    super(stub, "", DeduplicationMode.DISABLED, options, integration);
   }
 
   @Override
@@ -31,4 +30,14 @@ public CompletableFuture send(Invocation invocation)
       release();
     }
   }
+
+  @Override
+  public SendAttempt trySend(Invocation invocation) {
+    acquire();
+    try {
+      return doTrySend(lastSent + 1, (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
 }
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java b/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java
index b1c36524..3335c766 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java
@@ -9,13 +9,16 @@
 package dev.restate.integration;
 
 /**
- * Thrown by {@code send} when the producer cannot accept a record right now (the send window is
- * depleted or the transport is not writable). Unchecked: catch it to pace, or await {@link
- * ProducerBase#waitReady()} before sending.
+ * Thrown by {@code send} when local buffer capacity does not become available within the configured
+ * {@link ProducerOptions#maxBlockTime()}.
  */
 @org.jetbrains.annotations.ApiStatus.Experimental
 public class ProducerNotReadyException extends RuntimeException {
   public ProducerNotReadyException(String message) {
     super(message);
   }
+
+  public ProducerNotReadyException(String message, Throwable cause) {
+    super(message, cause);
+  }
 }
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
new file mode 100644
index 00000000..4597ede8
--- /dev/null
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
@@ -0,0 +1,139 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.integration;
+
+import dev.restate.ingestion.v1.IngestionDefaults;
+import java.time.Duration;
+import java.util.Objects;
+
+/** Configuration shared by at-least-once and exactly-once producers. */
+@org.jetbrains.annotations.ApiStatus.Experimental
+public final class ProducerOptions {
+
+  /** Kafka-compatible default producer buffer size: 32 MiB. */
+  public static final long DEFAULT_BUFFER_MEMORY = 32L * 1024 * 1024;
+
+  /** Kafka-compatible default maximum admission wait: one minute. */
+  public static final Duration DEFAULT_MAX_BLOCK_TIME = Duration.ofMinutes(1);
+
+  private static final ProducerOptions DEFAULTS = builder().build();
+
+  private final long bufferMemory;
+  private final Duration maxBlockTime;
+  private final IngestionDefaults defaultMetadata;
+
+  private ProducerOptions(Builder builder) {
+    this.bufferMemory = builder.bufferMemory;
+    this.maxBlockTime = builder.maxBlockTime;
+    this.defaultMetadata =
+        builder.defaultMetadata == null
+            ? IngestionDefaults.getDefaultInstance()
+            : ((InvocationMetadataImpl) builder.defaultMetadata).toDefaults();
+  }
+
+  /**
+   * Returns options with the standard defaults.
+   *
+   * @return the shared default options
+   */
+  public static ProducerOptions defaults() {
+    return DEFAULTS;
+  }
+
+  /**
+   * Starts building producer options.
+   *
+   * @return a new builder
+   */
+  public static Builder builder() {
+    return new Builder();
+  }
+
+  /**
+   * Maximum serialized bytes retained while invocations wait to be handed to the transport.
+   *
+   * @return the local buffer limit in bytes
+   */
+  public long bufferMemory() {
+    return bufferMemory;
+  }
+
+  /**
+   * Maximum time {@code send} waits for buffer capacity before refusing an invocation.
+   *
+   * @return the maximum admission wait
+   */
+  public Duration maxBlockTime() {
+    return maxBlockTime;
+  }
+
+  /**
+   * Returns a mutable copy of the default invocation metadata captured by these options. Mutating
+   * the returned object does not change these options.
+   *
+   * @return a mutable copy of the default invocation metadata
+   */
+  public InvocationMetadata defaultMetadata() {
+    return InvocationMetadataImpl.fromDefaults(defaultMetadata);
+  }
+
+  IngestionDefaults toDefaults() {
+    return defaultMetadata;
+  }
+
+  /** Builder for {@link ProducerOptions}. */
+  public static final class Builder {
+    private long bufferMemory = DEFAULT_BUFFER_MEMORY;
+    private Duration maxBlockTime = DEFAULT_MAX_BLOCK_TIME;
+    private InvocationMetadata defaultMetadata;
+
+    private Builder() {}
+
+    /**
+     * Sets the maximum serialized bytes retained while invocations wait to be handed to the
+     * transport.
+     *
+     * @param bytes a positive byte count
+     */
+    public Builder bufferMemory(long bytes) {
+      if (bytes <= 0) {
+        throw new IllegalArgumentException("bufferMemory must be greater than zero");
+      }
+      this.bufferMemory = bytes;
+      return this;
+    }
+
+    /**
+     * Sets how long {@code send} waits for buffer capacity. {@link Duration#ZERO} makes {@code
+     * send} fail immediately when the buffer is full.
+     */
+    public Builder maxBlockTime(Duration duration) {
+      Objects.requireNonNull(duration, "maxBlockTime");
+      if (duration.isNegative()) {
+        throw new IllegalArgumentException("maxBlockTime must not be negative");
+      }
+      this.maxBlockTime = duration;
+      return this;
+    }
+
+    /**
+     * Sets invocation fields applied to every record unless overridden by that invocation. The
+     * metadata is snapshotted when {@link #build()} is called.
+     */
+    public Builder defaultMetadata(InvocationMetadata metadata) {
+      Objects.requireNonNull(metadata, "defaultMetadata");
+      this.defaultMetadata = metadata;
+      return this;
+    }
+
+    public ProducerOptions build() {
+      return new ProducerOptions(this);
+    }
+  }
+}
diff --git a/integration-client/src/main/java/dev/restate/integration/SendAttempt.java b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java
new file mode 100644
index 00000000..6d59e5b3
--- /dev/null
+++ b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java
@@ -0,0 +1,40 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.integration;
+
+import java.util.Objects;
+import java.util.concurrent.CompletableFuture;
+
+/** Result of a non-blocking producer admission attempt. */
+@org.jetbrains.annotations.ApiStatus.Experimental
+public sealed interface SendAttempt {
+
+  /**
+   * The invocation was accepted; the future completes on durable acknowledgement.
+   *
+   * @param acknowledgement future completed when Restate durably acknowledges the invocation
+   */
+  record Accepted(CompletableFuture acknowledgement) implements SendAttempt {
+    public Accepted {
+      Objects.requireNonNull(acknowledgement, "acknowledgement");
+    }
+  }
+
+  /**
+   * The invocation was not accepted because the local buffer was full. {@code ready} completes when
+   * retrying may succeed; it is a notification, not a capacity reservation.
+   *
+   * @param ready future completed when retrying may succeed
+   */
+  record Backpressured(CompletableFuture ready) implements SendAttempt {
+    public Backpressured {
+      Objects.requireNonNull(ready, "ready");
+    }
+  }
+}
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 7bb9bd81..c22e4371 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -25,8 +25,10 @@
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
+import java.time.Duration;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
@@ -77,6 +79,46 @@ void producerSendsDisabledDedupHandshake() throws Exception {
     assertThat(start.getStart().getDefaults().getService()).isEqualTo("Svc");
   }
 
+  @Test
+  void producerOptionsHaveKafkaCompatibleDefaultsAndSnapshotMetadata() throws Exception {
+    InvocationMetadata metadata = InvocationMetadata.create().setServiceName("Original");
+    ProducerOptions options = ProducerOptions.builder().defaultMetadata(metadata).build();
+    metadata.setServiceName("Changed");
+
+    assertThat(options.bufferMemory()).isEqualTo(32L * 1024 * 1024);
+    assertThat(options.maxBlockTime()).isEqualTo(Duration.ofMinutes(1));
+    assertThat(options.defaultMetadata().getServiceName()).isEqualTo("Original");
+
+    client.newProducer(options);
+    assertThat(fake.take().getStart().getDefaults().getService()).isEqualTo("Original");
+  }
+
+  @Test
+  void exactlyOnceProducerAcceptsProducerOptions() throws Exception {
+    ProducerOptions options =
+        ProducerOptions.builder()
+            .defaultMetadata(InvocationMetadata.create().setHandlerName("handle"))
+            .build();
+
+    client.newExactlyOnceProducer("producer-1", options);
+
+    IngestionRequest start = fake.take();
+    assertThat(start.getStart().getProducerId()).isEqualTo("producer-1");
+    assertThat(start.getStart().getDefaults().getHandler()).isEqualTo("handle");
+  }
+
+  @Test
+  void producerOptionsValidateBufferAndBlockTime() {
+    assertThatThrownBy(() -> ProducerOptions.builder().bufferMemory(0))
+        .isInstanceOf(IllegalArgumentException.class);
+    assertThatThrownBy(() -> ProducerOptions.builder().maxBlockTime(Duration.ofMillis(-1)))
+        .isInstanceOf(IllegalArgumentException.class);
+    assertThatThrownBy(() -> ProducerOptions.builder().maxBlockTime(null))
+        .isInstanceOf(NullPointerException.class);
+    assertThatThrownBy(() -> ProducerOptions.builder().defaultMetadata(null))
+        .isInstanceOf(NullPointerException.class);
+  }
+
   @Test
   void exactlyOnceProducerSendsOffsetBasedHandshake() throws Exception {
     client.newExactlyOnceProducer("producer-1");
@@ -144,31 +186,98 @@ void invocationTypesAreSealedToSdkImplementations() {
   }
 
   @Test
-  void sendThrowsWhenNotReadyThenSucceedsAfterGrant() throws Exception {
+  void sendBuffersBeforeInitialWindowGrant() throws Exception {
     Producer producer = client.newProducer();
     fake.take(); // Start
 
-    Invocation inv = newBody("a");
-    assertThatThrownBy(() -> producer.send(inv)).isInstanceOf(ProducerNotReadyException.class);
+    CompletableFuture acknowledgement = producer.send(newBody("a"));
+    assertThat(producer.lastSentOffset()).isEqualTo(0L);
+    assertThat(acknowledgement).isNotDone();
+    fake.assertNoRequest();
 
     fake.grantWindow(10_000);
-    CompletableFuture f = producer.send(newBody("b"));
     assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
 
     fake.ack(0L);
-    assertThat(get(f).offset()).isEqualTo(0L);
+    assertThat(get(acknowledgement).offset()).isEqualTo(0L);
   }
 
   @Test
-  void waitReadyCompletesOnWindowGrant() throws Exception {
-    Producer producer = client.newProducer();
+  void trySendReportsBackpressureAndSignalsWhenCapacityReturns() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(128).maxBlockTime(Duration.ZERO).build());
     fake.take(); // Start
 
-    CompletableFuture ready = producer.waitReady();
+    SendAttempt first = producer.trySend(newBody("a".repeat(80)));
+    assertThat(first).isInstanceOf(SendAttempt.Accepted.class);
+
+    SendAttempt second = producer.trySend(newBody("b".repeat(80)));
+    assertThat(second).isInstanceOf(SendAttempt.Backpressured.class);
+    CompletableFuture ready = ((SendAttempt.Backpressured) second).ready();
+    assertThat(producer.lastSentOffset()).isEqualTo(0L);
     assertThat(ready).isNotDone();
 
     fake.grantWindow(10_000);
     get(ready);
+
+    SendAttempt retried = producer.trySend(newBody("b".repeat(80)));
+    assertThat(retried).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(1L);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(1L);
+  }
+
+  @Test
+  void sendBlocksUntilBufferCapacityReturns() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder()
+                .bufferMemory(128)
+                .maxBlockTime(Duration.ofSeconds(5))
+                .build());
+    fake.take(); // Start
+    producer.send(newBody("a".repeat(80)));
+
+    CountDownLatch attempting = new CountDownLatch(1);
+    CompletableFuture> blocked =
+        CompletableFuture.supplyAsync(
+            () -> {
+              attempting.countDown();
+              return producer.send(newBody("b".repeat(80)));
+            });
+    assertThat(attempting.await(5, TimeUnit.SECONDS)).isTrue();
+    Thread.sleep(50);
+    assertThat(blocked).isNotDone();
+
+    fake.grantWindow(10_000);
+    get(blocked);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(1L);
+  }
+
+  @Test
+  void zeroMaxBlockTimeFailsWithoutConsumingOffset() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(128).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+    producer.send(newBody("a".repeat(80)));
+
+    assertThatThrownBy(() -> producer.send(newBody("b".repeat(80))))
+        .isInstanceOf(ProducerNotReadyException.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(0L);
+  }
+
+  @Test
+  void oversizedInvocationIsRejectedWithoutConsumingOffset() throws Exception {
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(64).build());
+    fake.take(); // Start
+
+    assertThatThrownBy(() -> producer.send(newBody("a".repeat(100))))
+        .isInstanceOf(IllegalArgumentException.class)
+        .hasMessageContaining("exceeding bufferMemory");
+    assertThat(producer.lastSentOffset()).isEqualTo(-1L);
   }
 
   @Test
@@ -234,6 +343,25 @@ void streamErrorFailsPendingFuturesFast() throws Exception {
     assertThatThrownBy(() -> producer.send(newBody("b"))).isInstanceOf(IllegalStateException.class);
   }
 
+  @Test
+  void streamErrorFailsBufferedAcknowledgementsAndBackpressureWaiters() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(128).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    SendAttempt.Accepted accepted =
+        (SendAttempt.Accepted) producer.trySend(newBody("a".repeat(80)));
+    SendAttempt.Backpressured backpressured =
+        (SendAttempt.Backpressured) producer.trySend(newBody("b".repeat(80)));
+
+    fake.error(ErrorKind.ERROR_KIND_BAD_REQUEST, "nope");
+
+    assertThatThrownBy(() -> get(accepted.acknowledgement()))
+        .isInstanceOf(ExecutionException.class);
+    assertThatThrownBy(() -> get(backpressured.ready())).isInstanceOf(ExecutionException.class);
+  }
+
   @Test
   void exactlyOnceRejectsNonIncreasingOffsets() throws Exception {
     ExactlyOnceProducer producer = client.newExactlyOnceProducer("p1");
@@ -252,6 +380,27 @@ void exactlyOnceRejectsNonIncreasingOffsets() throws Exception {
     assertThat(producer.lastSentOffset()).isEqualTo(6L);
   }
 
+  @Test
+  void exactlyOnceTrySendDoesNotConsumeBackpressuredOffset() throws Exception {
+    ExactlyOnceProducer producer =
+        client.newExactlyOnceProducer(
+            "p1", ProducerOptions.builder().bufferMemory(128).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    assertThat(producer.trySend(5L, newBody("a".repeat(80))))
+        .isInstanceOf(SendAttempt.Accepted.class);
+    SendAttempt rejected = producer.trySend(6L, newBody("b".repeat(80)));
+    assertThat(rejected).isInstanceOf(SendAttempt.Backpressured.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(5L);
+
+    fake.grantWindow(10_000);
+    get(((SendAttempt.Backpressured) rejected).ready());
+
+    assertThat(producer.trySend(6L, newBody("b".repeat(80))))
+        .isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(6L);
+  }
+
   @Test
   void invocationFieldsMapToProto() throws Exception {
     Producer producer = client.newProducer();
@@ -322,6 +471,10 @@ IngestionRequest take() throws InterruptedException {
       return req;
     }
 
+    void assertNoRequest() {
+      assertThat(received.poll()).isNull();
+    }
+
     void grantWindow(long bytes) {
       responses.onNext(
           IngestionResponse.newBuilder()

From 680749920126282179b65bbc38e64a51d4536449 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 11:38:21 +0200
Subject: [PATCH 04/17] Improvements

---
 .../restate/integration/AbstractProducer.java | 50 +++++++++++++++++--
 .../integration/ExactlyOnceProducer.java      | 19 ++++---
 .../integration/ExactlyOnceProducerImpl.java  |  2 +-
 .../integration/IntegrationClient.java        |  5 +-
 .../dev/restate/integration/Producer.java     | 15 +++---
 .../dev/restate/integration/ProducerBase.java | 29 ++++++++---
 ... => ProducerBufferExhaustedException.java} |  8 +--
 .../dev/restate/integration/ProducerImpl.java |  2 +-
 .../integration/IntegrationClientTest.java    | 28 +++++++++--
 9 files changed, 122 insertions(+), 36 deletions(-)
 rename integration-client/src/main/java/dev/restate/integration/{ProducerNotReadyException.java => ProducerBufferExhaustedException.java} (65%)

diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
index ad179a1c..254ac008 100644
--- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
+++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
@@ -25,6 +25,7 @@
 import java.util.Map;
 import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.atomic.AtomicReference;
 
 /**
@@ -134,6 +135,26 @@ public CompletableFuture waitAcknowledged(long offset) {
     }
   }
 
+  @Override
+  public long flush() {
+    acquire();
+    try {
+      return awaitFlush(registerAckWaiter(lastSent));
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture flushAsync() {
+    acquire();
+    try {
+      return registerAckWaiter(lastSent);
+    } finally {
+      release();
+    }
+  }
+
   /**
    * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark
    * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are
@@ -170,7 +191,7 @@ public void close() {
 
   /** Admit a record, blocking up to the configured maximum when the local buffer is full. */
   final CompletableFuture doSend(long offset, InvocationImpl invocation)
-      throws ProducerNotReadyException {
+      throws ProducerBufferExhaustedException {
     PreparedSend prepared = prepare(offset, invocation);
     List> ready;
     CompletableFuture acknowledgement;
@@ -191,7 +212,7 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio
           lock.wait(millis, nanos);
         } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
-          throw new ProducerNotReadyException(
+          throw new ProducerBufferExhaustedException(
               "interrupted while waiting for producer buffer capacity", e);
         }
         ensureOpenLocked();
@@ -312,8 +333,29 @@ private void ensureOpenLocked() {
     }
   }
 
-  private ProducerNotReadyException admissionTimeout() {
-    return new ProducerNotReadyException("producer buffer remained full for " + maxBlockTime);
+  private ProducerBufferExhaustedException admissionTimeout() {
+    return new ProducerBufferExhaustedException(
+        "producer buffer remained full for " + maxBlockTime);
+  }
+
+  private static long awaitFlush(CompletableFuture flush) {
+    try {
+      return flush.get();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e);
+    } catch (ExecutionException e) {
+      Throwable cause = e.getCause();
+      if (cause instanceof RuntimeException runtimeException) {
+        throw runtimeException;
+      }
+      if (cause instanceof Error error) {
+        throw error;
+      }
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause);
+    }
   }
 
   private static long toNanosSaturated(Duration duration) {
diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
index 9a57ee31..c6deefda 100644
--- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
+++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
@@ -19,8 +19,8 @@
  * sequence. E.g., for a Kafka consumer {@code groupId/topic/partition}, for Postgres logical
  * replication the slot name. Because deduplication happens on {@code (producerId, offset)}, it is
  * then safe to replay from your last checkpoint after a crash: already-committed offsets are
- * dropped, and {@link #flush} / {@link #waitAcknowledged(long)} reports how far Restate has durably
- * caught up so you can advance the checkpoint. After a stream failure, {@link
+ * dropped, and {@link #flush()} / {@link #waitAcknowledged(long)} reports how far Restate has
+ * durably caught up so you can advance the checkpoint. After a stream failure, {@link
  * #lastAcknowledgedOffset()} remains available so you can determine where to resume.
  *
  * 
Sending

@@ -35,10 +35,13 @@
  *
  * 
{@code
  * producer.send(lsn, Invocation.create().setBody(payload));
- * long committed = producer.flush().get();
+ * long committed = producer.flush();
  * checkpoint.store(committed);
  * }

  *
+ * 
{@link #close()} does not flush. Call {@link #flush()} before closing, or await {@link
+ * #flushAsync()}, when accepted invocations must be durably committed.
+ *
  * 
Thread safety

  *
  * A producer is not thread-safe and fails fast with {@link
@@ -51,8 +54,8 @@ public interface ExactlyOnceProducer extends ProducerBase {
    * Sends an invocation at {@code offset}.
    *
    * 
If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()}
-   * for capacity. The invocation is refused with {@link ProducerNotReadyException} if the timeout
-   * elapses. A zero duration makes this method fail immediately under backpressure.
+   * for capacity. The invocation is refused with {@link ProducerBufferExhaustedException} if the
+   * timeout elapses. A zero duration makes this method fail immediately under backpressure.
    *
    * 
The returned future completes when the invocation is durably committed by Restate.
    *
@@ -61,8 +64,8 @@ public interface ExactlyOnceProducer extends ProducerBase {
    * @param invocation the invocation to send
    * @return a future completing, once the record is durably committed by Restate, with the {@link
    *     SendResult} carrying {@code offset}
-   * @throws ProducerNotReadyException if buffer capacity does not become available before the
-   *     configured maximum blocking time elapses
+   * @throws ProducerBufferExhaustedException if buffer capacity does not become available before
+   *     the configured maximum blocking time elapses, or the thread is interrupted while waiting
    * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link
    *     #lastSentOffset()}, or the serialized invocation is larger than {@link
    *     ProducerOptions#bufferMemory()}
@@ -70,7 +73,7 @@ public interface ExactlyOnceProducer extends ProducerBase {
    *     another thread
    */
   CompletableFuture send(long offset, Invocation invocation)
-      throws ProducerNotReadyException;
+      throws ProducerBufferExhaustedException;
 
   /**
    * Attempts to send an invocation at {@code offset} without blocking.
diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java
index 1c7371a7..03b90f1c 100644
--- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java
@@ -36,7 +36,7 @@ public SendAttempt trySend(long offset, Invocation invocation) {
 
   @Override
   public CompletableFuture send(long offset, Invocation invocation)
-      throws ProducerNotReadyException {
+      throws ProducerBufferExhaustedException {
     acquire();
     try {
       checkOffset(offset);
diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
index ed6a70af..634d5ea7 100644
--- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
+++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
@@ -69,7 +69,10 @@ public interface IntegrationClient extends AutoCloseable {
    */
   ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOptions options);
 
-  /** Shuts down the underlying client. */
+  /**
+   * Shuts down the underlying client. Flush and close any producers first when their accepted
+   * invocations must be durably committed.
+   */
   @Override
   void close();
 
diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java
index a65cf7a6..c5995f6b 100644
--- a/integration-client/src/main/java/dev/restate/integration/Producer.java
+++ b/integration-client/src/main/java/dev/restate/integration/Producer.java
@@ -31,10 +31,13 @@
  *   for (byte[] payload : payloads) {
  *     producer.send(Invocation.create().setBody(payload));
  *   }
- *   producer.flush().get(); // block until everything sent so far is durably committed
+ *   producer.flush(); // block until everything sent so far is durably committed
  * }
  * }
* + *

{@link #close()} does not flush. Call {@link #flush()} before closing, or await {@link + * #flushAsync()}, when accepted invocations must be durably committed. + * *

Stream defaults

* * Pass an {@link InvocationMetadata} to {@link IntegrationClient#newProducer(InvocationMetadata)}, @@ -60,21 +63,21 @@ public interface Producer extends ProducerBase { * Sends an invocation. * *

If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()} - * for capacity. The invocation is refused with {@link ProducerNotReadyException} if the timeout - * elapses. A zero duration makes this method fail immediately under backpressure. + * for capacity. The invocation is refused with {@link ProducerBufferExhaustedException} if the + * timeout elapses. A zero duration makes this method fail immediately under backpressure. * *

The returned future completes when the invocation is durably committed by Restate. * * @param invocation the invocation to send * @return a future completing, once the invocation is durably committed by Restate. - * @throws ProducerNotReadyException if buffer capacity does not become available before the - * configured maximum blocking time elapses + * @throws ProducerBufferExhaustedException if buffer capacity does not become available before + * the configured maximum blocking time elapses, or the thread is interrupted while waiting * @throws IllegalArgumentException if the serialized invocation is larger than {@link * ProducerOptions#bufferMemory()} * @throws java.util.ConcurrentModificationException if the producer is used concurrently from * another thread */ - CompletableFuture send(Invocation invocation) throws ProducerNotReadyException; + CompletableFuture send(Invocation invocation) throws ProducerBufferExhaustedException; /** * Attempts to send an invocation without blocking. diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java index 153801bd..03db1994 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java @@ -53,20 +53,33 @@ public interface ProducerBase extends AutoCloseable { CompletableFuture waitAcknowledged(long offset); /** - * Awaits durable acknowledgement of every invocation sent so far. + * Blocks until every invocation accepted before this call is durably acknowledged. * - * @return a future completing, once all invocations sent so far are durably acknowledged by - * Restate, with the highest durably committed offset + * @return the highest durably committed offset + * @throws IntegrationClientException if the producer fails before all invocations are + * acknowledged * @throws java.util.ConcurrentModificationException if the producer is used concurrently from * another thread */ - default CompletableFuture flush() { - return waitAcknowledged(lastSentOffset()); - } + long flush(); /** - * Closes the producer and shuts down its stream. Any not-yet-acknowledged invocation completes - * its future exceptionally with an {@link IntegrationClientException}. + * Asynchronously awaits durable acknowledgement of every invocation accepted before this call. + * + * @return a future completing with the highest durably committed offset once all invocations sent + * so far are acknowledged + * @throws java.util.ConcurrentModificationException if the producer is used concurrently from + * another thread + */ + CompletableFuture flushAsync(); + + /** + * Immediately closes the producer and shuts down its stream without flushing. Any + * not-yet-acknowledged invocation completes its future exceptionally with an {@link + * IntegrationClientException}. + * + *

Call {@link #flush()} before closing, or await {@link #flushAsync()}, when accepted + * invocations must be durably committed. * * @throws java.util.ConcurrentModificationException if the producer is used concurrently from * another thread diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java b/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java similarity index 65% rename from integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java rename to integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java index 3335c766..a52e7936 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerNotReadyException.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java @@ -10,15 +10,15 @@ /** * Thrown by {@code send} when local buffer capacity does not become available within the configured - * {@link ProducerOptions#maxBlockTime()}. + * {@link ProducerOptions#maxBlockTime()}, or the thread is interrupted while waiting for capacity. */ @org.jetbrains.annotations.ApiStatus.Experimental -public class ProducerNotReadyException extends RuntimeException { - public ProducerNotReadyException(String message) { +public class ProducerBufferExhaustedException extends RuntimeException { + public ProducerBufferExhaustedException(String message) { super(message); } - public ProducerNotReadyException(String message, Throwable cause) { + public ProducerBufferExhaustedException(String message, Throwable cause) { super(message, cause); } } diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java index 62b2815b..642be6bd 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java @@ -22,7 +22,7 @@ final class ProducerImpl extends AbstractProducer implements Producer { @Override public CompletableFuture send(Invocation invocation) - throws ProducerNotReadyException { + throws ProducerBufferExhaustedException { acquire(); try { return doSend(lastSent + 1, (InvocationImpl) invocation); diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java index c22e4371..ed5ce170 100644 --- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java +++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java @@ -265,7 +265,7 @@ void zeroMaxBlockTimeFailsWithoutConsumingOffset() throws Exception { producer.send(newBody("a".repeat(80))); assertThatThrownBy(() -> producer.send(newBody("b".repeat(80)))) - .isInstanceOf(ProducerNotReadyException.class); + .isInstanceOf(ProducerBufferExhaustedException.class); assertThat(producer.lastSentOffset()).isEqualTo(0L); } @@ -299,14 +299,14 @@ void waitAcknowledgedCompletesAtWatermark() throws Exception { } @Test - void flushCompletesWhenEverythingSentIsCommitted() throws Exception { + void flushAsyncCompletesWhenEverythingSentIsCommitted() throws Exception { Producer producer = client.newProducer(); fake.take(); // Start fake.grantWindow(10_000); CompletableFuture a = producer.send(newBody("a")); // offset 0 CompletableFuture b = producer.send(newBody("b")); // offset 1 - CompletableFuture flushed = producer.flush(); // waits up to the last sent offset (1) + CompletableFuture flushed = producer.flushAsync(); // waits up to the last sent offset (1) assertThat(a).isNotDone(); assertThat(flushed).isNotDone(); @@ -320,6 +320,28 @@ void flushCompletesWhenEverythingSentIsCommitted() throws Exception { assertThat(get(flushed)).isEqualTo(1L); // last durably committed offset } + @Test + void flushBlocksUntilEverythingSentIsCommitted() throws Exception { + Producer producer = client.newProducer(); + fake.take(); // Start + fake.grantWindow(10_000); + producer.send(newBody("a")); + + CountDownLatch flushing = new CountDownLatch(1); + CompletableFuture flushed = + CompletableFuture.supplyAsync( + () -> { + flushing.countDown(); + return producer.flush(); + }); + assertThat(flushing.await(5, TimeUnit.SECONDS)).isTrue(); + Thread.sleep(50); + assertThat(flushed).isNotDone(); + + fake.ack(0L); + assertThat(get(flushed)).isEqualTo(0L); + } + @Test void streamErrorFailsPendingFuturesFast() throws Exception { Producer producer = client.newProducer(); From 391c250581da3e0e076f6dc11b305c9a14eaddd4 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 11:48:18 +0200 Subject: [PATCH 05/17] More improvements --- .../integration/IntegrationClient.java | 18 +++++++++++---- .../integration/IntegrationClientImpl.java | 22 ------------------- .../restate/integration/ProducerOptions.java | 21 +----------------- 3 files changed, 15 insertions(+), 46 deletions(-) diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java index 634d5ea7..8bca0161 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java @@ -17,7 +17,9 @@ public interface IntegrationClient extends AutoCloseable { * * @return a new producer */ - Producer newProducer(); + default Producer newProducer() { + return newProducer(ProducerOptions.DEFAULTS); + } /** * Creates at-least-once {@link Producer} with the given stream defaults. @@ -26,7 +28,9 @@ public interface IntegrationClient extends AutoCloseable { * @return a new producer * @throws NullPointerException if {@code defaultMetadata} is {@code null} */ - Producer newProducer(InvocationMetadata defaultMetadata); + default Producer newProducer(InvocationMetadata defaultMetadata) { + return newProducer(ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); + } /** * Creates an at-least-once {@link Producer} with the given options. @@ -44,7 +48,9 @@ public interface IntegrationClient extends AutoCloseable { * @return a new exactly-once producer * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank */ - ExactlyOnceProducer newExactlyOnceProducer(String producerId); + default ExactlyOnceProducer newExactlyOnceProducer(String producerId) { + return newExactlyOnceProducer(producerId, ProducerOptions.DEFAULTS); + } /** * Creates an {@link ExactlyOnceProducer} identified by {@code producerId} with the given stream @@ -56,7 +62,11 @@ public interface IntegrationClient extends AutoCloseable { * @throws IllegalArgumentException if {@code producerId} is {@code null} or blank * @throws NullPointerException if {@code defaultMetadata} is {@code null} */ - ExactlyOnceProducer newExactlyOnceProducer(String producerId, InvocationMetadata defaultMetadata); + default ExactlyOnceProducer newExactlyOnceProducer( + String producerId, InvocationMetadata defaultMetadata) { + return newExactlyOnceProducer( + producerId, ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); + } /** * Creates an {@link ExactlyOnceProducer} identified by {@code producerId} with the given options. diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java index 0431231f..56f40efc 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java @@ -51,33 +51,11 @@ static IntegrationClient forChannel(ManagedChannel channel, String integration) return new IntegrationClientImpl(channel, IngestionSvcGrpc.newStub(channel), integration); } - @Override - public Producer newProducer() { - return newProducer(ProducerOptions.defaults()); - } - - @Override - public Producer newProducer(InvocationMetadata defaultMetadata) { - return newProducer(ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); - } - @Override public Producer newProducer(ProducerOptions options) { return new ProducerImpl(stub, Objects.requireNonNull(options, "options"), integration); } - @Override - public ExactlyOnceProducer newExactlyOnceProducer(String producerId) { - return newExactlyOnceProducer(producerId, ProducerOptions.defaults()); - } - - @Override - public ExactlyOnceProducer newExactlyOnceProducer( - String producerId, InvocationMetadata defaultMetadata) { - return newExactlyOnceProducer( - producerId, ProducerOptions.builder().defaultMetadata(defaultMetadata).build()); - } - @Override public ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOptions options) { if (producerId == null || producerId.isBlank()) { diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java index 4597ede8..c6bbef11 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java @@ -22,7 +22,7 @@ public final class ProducerOptions { /** Kafka-compatible default maximum admission wait: one minute. */ public static final Duration DEFAULT_MAX_BLOCK_TIME = Duration.ofMinutes(1); - private static final ProducerOptions DEFAULTS = builder().build(); + static final ProducerOptions DEFAULTS = builder().build(); private final long bufferMemory; private final Duration maxBlockTime; @@ -37,15 +37,6 @@ private ProducerOptions(Builder builder) { : ((InvocationMetadataImpl) builder.defaultMetadata).toDefaults(); } - /** - * Returns options with the standard defaults. - * - * @return the shared default options - */ - public static ProducerOptions defaults() { - return DEFAULTS; - } - /** * Starts building producer options. * @@ -73,16 +64,6 @@ public Duration maxBlockTime() { return maxBlockTime; } - /** - * Returns a mutable copy of the default invocation metadata captured by these options. Mutating - * the returned object does not change these options. - * - * @return a mutable copy of the default invocation metadata - */ - public InvocationMetadata defaultMetadata() { - return InvocationMetadataImpl.fromDefaults(defaultMetadata); - } - IngestionDefaults toDefaults() { return defaultMetadata; } From 75e6a13f3eb5268ceb4990db41245b3fe9655596 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 12:04:03 +0200 Subject: [PATCH 06/17] nullability checks --- integration-client/build.gradle.kts | 1 + .../restate/integration/AbstractProducer.java | 53 +++++++++++-------- .../restate/integration/IngressEndpoint.java | 5 +- .../integration/IntegrationClient.java | 4 +- .../IntegrationClientException.java | 3 +- .../integration/IntegrationClientImpl.java | 3 +- .../dev/restate/integration/Invocation.java | 36 +++++++------ .../restate/integration/InvocationImpl.java | 33 ++++++------ .../integration/InvocationMetadata.java | 30 ++++++----- .../integration/InvocationMetadataImpl.java | 27 +++++----- .../restate/integration/ProducerOptions.java | 3 +- .../dev/restate/integration/SendAttempt.java | 3 +- .../dev/restate/integration/package-info.java | 13 +++++ .../integration/IntegrationClientTest.java | 1 - 14 files changed, 127 insertions(+), 88 deletions(-) create mode 100644 integration-client/src/main/java/dev/restate/integration/package-info.java diff --git a/integration-client/build.gradle.kts b/integration-client/build.gradle.kts index deb94f76..bebbc1da 100644 --- a/integration-client/build.gradle.kts +++ b/integration-client/build.gradle.kts @@ -25,6 +25,7 @@ dependencies { // @ApiStatus.Experimental markers on the public API. compileOnly(libs.jetbrains.annotations) + compileOnly(libs.jspecify) testImplementation(libs.junit.jupiter) testImplementation(libs.assertj) diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java index 254ac008..6f1ad524 100644 --- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java @@ -23,10 +23,12 @@ import java.util.Iterator; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.TreeMap; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutionException; import java.util.concurrent.atomic.AtomicReference; +import org.jspecify.annotations.Nullable; /** * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and @@ -51,10 +53,10 @@ abstract class AbstractProducer implements ProducerBase { private final Object lock = new Object(); // Set once, synchronously, in beforeStart() before the constructor sends the Start frame. - private volatile ClientCallStreamObserver callObserver; + private volatile @Nullable ClientCallStreamObserver callObserver; // ---- fail-fast single-thread guard ---- - private final AtomicReference owner = new AtomicReference<>(); + private final AtomicReference<@Nullable Thread> owner = new AtomicReference<>(); private int reentrancy; // ---- state guarded by `lock` ---- @@ -65,7 +67,7 @@ abstract class AbstractProducer implements ProducerBase { private final List capacityWaiters = new ArrayList<>(); private final TreeMap>> ackWaiters = new TreeMap<>(); private boolean closed = false; - private IntegrationClientException failure; + private @Nullable IntegrationClientException failure; private final long bufferMemory; private final Duration maxBlockTime; @@ -97,7 +99,7 @@ abstract class AbstractProducer implements ProducerBase { .setDefaults(options.toDefaults())) .build(); synchronized (lock) { - callObserver.onNext(start); + Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(start); } } @@ -163,7 +165,8 @@ public CompletableFuture flushAsync() { private CompletableFuture registerAckWaiter(long offset) { synchronized (lock) { if (closed) { - return CompletableFuture.failedFuture(failure); + return CompletableFuture.failedFuture( + Objects.requireNonNull(failure, "closed producer has no failure")); } if (offset <= lastCommitted) { return CompletableFuture.completedFuture(lastCommitted); @@ -193,7 +196,7 @@ public void close() { final CompletableFuture doSend(long offset, InvocationImpl invocation) throws ProducerBufferExhaustedException { PreparedSend prepared = prepare(offset, invocation); - List> ready; + List> ready; CompletableFuture acknowledgement; long waitStarted = System.nanoTime(); synchronized (lock) { @@ -227,7 +230,7 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio /** Attempt to admit a record without blocking or consuming an offset when capacity is absent. */ final SendAttempt doTrySend(long offset, InvocationImpl invocation) { PreparedSend prepared = prepare(offset, invocation); - List> ready; + List> ready; SendAttempt result; synchronized (lock) { ensureOpenLocked(); @@ -235,7 +238,7 @@ final SendAttempt doTrySend(long offset, InvocationImpl invocation) { result = new SendAttempt.Accepted(acceptLocked(prepared)); ready = drainLocked(); } else { - CompletableFuture future = new CompletableFuture<>(); + CompletableFuture<@Nullable Void> future = new CompletableFuture<>(); CapacityWaiter waiter = new CapacityWaiter(prepared.bufferSize(), future); capacityWaiters.add(waiter); future.whenComplete( @@ -286,14 +289,16 @@ private CompletableFuture acceptLocked(PreparedSend prepared) { } /** Write as many queued records as transport and protocol flow control currently permit. */ - private List> drainLocked() { + private List> drainLocked() { boolean freedCapacity = false; - while (!closed && budget > 0 && callObserver.isReady() && !bufferedSends.isEmpty()) { + ClientCallStreamObserver observer = + Objects.requireNonNull(callObserver, "gRPC request stream was not initialized"); + while (!closed && budget > 0 && observer.isReady() && !bufferedSends.isEmpty()) { BufferedSend send = bufferedSends.removeFirst(); bufferedBytes -= send.bufferSize(); budget -= send.windowDebit(); freedCapacity = true; - callObserver.onNext(send.request()); + observer.onNext(send.request()); } if (!freedCapacity) { @@ -302,7 +307,7 @@ private List> drainLocked() { lock.notifyAll(); long available = bufferMemory - bufferedBytes; - List> ready = new ArrayList<>(); + List> ready = new ArrayList<>(); for (Iterator it = capacityWaiters.iterator(); it.hasNext(); ) { CapacityWaiter waiter = it.next(); if (waiter.requiredBytes() <= available) { @@ -314,15 +319,15 @@ private List> drainLocked() { } private void drainAndWake() { - List> ready; + List> ready; synchronized (lock) { ready = drainLocked(); } completeReady(ready); } - private static void completeReady(List> ready) { - for (CompletableFuture future : ready) { + private static void completeReady(List> ready) { + for (CompletableFuture<@Nullable Void> future : ready) { future.complete(null); } } @@ -346,13 +351,17 @@ private static long awaitFlush(CompletableFuture flush) { throw new IntegrationClientException( IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e); } catch (ExecutionException e) { - Throwable cause = e.getCause(); + @Nullable Throwable cause = e.getCause(); if (cause instanceof RuntimeException runtimeException) { throw runtimeException; } if (cause instanceof Error error) { throw error; } + if (cause == null) { + throw new IntegrationClientException( + IntegrationClientException.Kind.UNKNOWN, "producer flush failed"); + } throw new IntegrationClientException( IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause); } @@ -367,10 +376,10 @@ private static long toNanosSaturated(Duration duration) { } private void onResponse(IngestionResponse resp) { - List> acksToComplete = null; + @Nullable List> acksToComplete = null; long watermark = -1; boolean drain = false; - IntegrationClientException err = null; + @Nullable IntegrationClientException err = null; synchronized (lock) { if (closed) { return; @@ -409,7 +418,7 @@ private void onResponse(IngestionResponse resp) { /** Mark the producer terminally closed and fail every pending future with {@code cause}. */ private void terminate(IntegrationClientException cause, boolean halfClose) { - List> capacity; + List> capacity; List> acks = new ArrayList<>(); synchronized (lock) { if (closed) { @@ -431,7 +440,7 @@ private void terminate(IntegrationClientException cause, boolean halfClose) { ackWaiters.clear(); } if (halfClose) { - ClientCallStreamObserver obs = callObserver; + @Nullable ClientCallStreamObserver obs = callObserver; if (obs != null) { try { obs.onCompleted(); @@ -440,7 +449,7 @@ private void terminate(IntegrationClientException cause, boolean halfClose) { } } } - for (CompletableFuture f : capacity) { + for (CompletableFuture<@Nullable Void> f : capacity) { f.completeExceptionally(cause); } for (CompletableFuture f : acks) { @@ -528,7 +537,7 @@ private record PreparedSend( private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {} - private record CapacityWaiter(long requiredBytes, CompletableFuture future) {} + private record CapacityWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {} private record SendResultImpl(long offset) implements SendResult {} } diff --git a/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java index 401fc4f6..edc66734 100644 --- a/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java +++ b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java @@ -9,6 +9,7 @@ package dev.restate.integration; import java.net.URI; +import org.jspecify.annotations.Nullable; /** Where to reach the Restate ingestion gRPC endpoint, parsed from an http(s) URL. */ final class IngressEndpoint { @@ -38,7 +39,7 @@ static IngressEndpoint parse(String raw) { throw new IllegalArgumentException("ingress url is not a valid URL: '" + raw + "'", e); } boolean tls; - String scheme = uri.getScheme() == null ? null : uri.getScheme().toLowerCase(); + @Nullable String scheme = uri.getScheme() == null ? null : uri.getScheme().toLowerCase(); if ("https".equals(scheme)) { tls = true; } else if ("http".equals(scheme)) { @@ -51,7 +52,7 @@ static IngressEndpoint parse(String raw) { + raw + "'"); } - String host = uri.getHost(); + @Nullable String host = uri.getHost(); if (host == null) { throw new IllegalArgumentException("ingress url has no host: '" + raw + "'"); } diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java index 8bca0161..04779444 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java @@ -8,6 +8,8 @@ // https://github.com/restatedev/sdk-java/blob/main/LICENSE package dev.restate.integration; +import org.jspecify.annotations.Nullable; + /** Entry point for producing invocations to Restate ingress over the ingestion API. */ @org.jetbrains.annotations.ApiStatus.Experimental public interface IntegrationClient extends AutoCloseable { @@ -94,7 +96,7 @@ static Builder builder(String ingressUrl) { /** Builder for {@link IntegrationClient}. */ final class Builder { private final String target; - private String authToken; + private @Nullable String authToken; private String integration = Version.INTEGRATION; private Builder(String target) { diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java index 6d546e63..60d36d21 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientException.java @@ -28,7 +28,8 @@ public enum Kind { private final Kind kind; public IntegrationClientException(Kind kind, String message) { - this(kind, message, null); + super(message != null ? message : kind.name()); + this.kind = kind; } public IntegrationClientException(Kind kind, String message, Throwable cause) { diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java index 56f40efc..c1a57ce9 100644 --- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java @@ -13,6 +13,7 @@ import io.grpc.ManagedChannelBuilder; import java.util.Objects; import java.util.concurrent.TimeUnit; +import org.jspecify.annotations.Nullable; /** {@link IntegrationClient} backed by a single gRPC {@link ManagedChannel} shared by producers. */ final class IntegrationClientImpl implements IntegrationClient { @@ -28,7 +29,7 @@ private IntegrationClientImpl( this.integration = integration; } - static IntegrationClient create(String target, String authToken, String integration) { + static IntegrationClient create(String target, @Nullable String authToken, String integration) { IngressEndpoint endpoint = IngressEndpoint.parse(target); ManagedChannelBuilder builder = ManagedChannelBuilder.forAddress(endpoint.host, endpoint.port); diff --git a/integration-client/src/main/java/dev/restate/integration/Invocation.java b/integration-client/src/main/java/dev/restate/integration/Invocation.java index 5a74c6d8..370e059d 100644 --- a/integration-client/src/main/java/dev/restate/integration/Invocation.java +++ b/integration-client/src/main/java/dev/restate/integration/Invocation.java @@ -11,11 +11,15 @@ import java.time.Duration; import java.time.Instant; import java.util.Map; +import org.jspecify.annotations.Nullable; /** * A single invocation to send through a {@link Producer} / {@link ExactlyOnceProducer}. * *

Instances are created via {@link #create()}. + * + *

Passing {@code null} to a setter clears that field. Optional getters return {@code null} when + * the corresponding field is not set. */ @org.jetbrains.annotations.ApiStatus.Experimental public sealed interface Invocation extends InvocationMetadata permits InvocationImpl { @@ -26,51 +30,51 @@ static Invocation create() { } /** The invocation payload. */ - Invocation setBody(byte[] body); + Invocation setBody(byte @Nullable [] body); byte[] getBody(); /** Schedule the invocation after a delay. Mutually exclusive with {@link #setInvokeTime}. */ - Invocation setDelay(Duration delay); + Invocation setDelay(@Nullable Duration delay); - Duration getDelay(); + @Nullable Duration getDelay(); /** Schedule the invocation at an absolute time. Mutually exclusive with {@link #setDelay}. */ - Invocation setInvokeTime(Instant invokeTime); + Invocation setInvokeTime(@Nullable Instant invokeTime); - Instant getInvokeTime(); + @Nullable Instant getInvokeTime(); /** W3C {@code traceparent}. */ - Invocation setTraceparent(String traceparent); + Invocation setTraceparent(@Nullable String traceparent); - String getTraceparent(); + @Nullable String getTraceparent(); /** W3C {@code tracestate}. */ - Invocation setTracestate(String tracestate); + Invocation setTracestate(@Nullable String tracestate); - String getTracestate(); + @Nullable String getTracestate(); @Override - Invocation setServiceName(String serviceName); + Invocation setServiceName(@Nullable String serviceName); @Override - Invocation setHandlerName(String handlerName); + Invocation setHandlerName(@Nullable String handlerName); @Override - Invocation setKey(String key); + Invocation setKey(@Nullable String key); @Override - Invocation setScope(String scope); + Invocation setScope(@Nullable String scope); @Override - Invocation setLimitKey(String limitKey); + Invocation setLimitKey(@Nullable String limitKey); @Override - Invocation setIdempotencyKey(String idempotencyKey); + Invocation setIdempotencyKey(@Nullable String idempotencyKey); @Override Invocation putHeader(String key, String value); @Override - Invocation setHeaders(Map headers); + Invocation setHeaders(@Nullable Map headers); } diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java index 38654354..fc422bdb 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationImpl.java @@ -13,6 +13,7 @@ import java.time.Duration; import java.time.Instant; import java.util.Map; +import org.jspecify.annotations.Nullable; /** * Mutable {@link Invocation} backed directly by the {@link IngestionInvocation.Builder} inherited @@ -22,7 +23,7 @@ final class InvocationImpl extends InvocationMetadataImpl implements Invocation { @Override - public Invocation setBody(byte[] body) { + public Invocation setBody(byte @Nullable [] body) { if (body == null) { builder.clearPayload(); } else { @@ -37,7 +38,7 @@ public byte[] getBody() { } @Override - public Invocation setDelay(Duration delay) { + public Invocation setDelay(@Nullable Duration delay) { if (delay == null) { builder.clearDelayMs(); } else { @@ -48,12 +49,12 @@ public Invocation setDelay(Duration delay) { } @Override - public Duration getDelay() { + public @Nullable Duration getDelay() { return builder.hasDelayMs() ? Duration.ofMillis(builder.getDelayMs()) : null; } @Override - public Invocation setInvokeTime(Instant invokeTime) { + public Invocation setInvokeTime(@Nullable Instant invokeTime) { if (invokeTime == null) { builder.clearInvokeTimeTsMs(); } else { @@ -64,12 +65,12 @@ public Invocation setInvokeTime(Instant invokeTime) { } @Override - public Instant getInvokeTime() { + public @Nullable Instant getInvokeTime() { return builder.hasInvokeTimeTsMs() ? Instant.ofEpochMilli(builder.getInvokeTimeTsMs()) : null; } @Override - public Invocation setTraceparent(String traceparent) { + public Invocation setTraceparent(@Nullable String traceparent) { if (traceparent == null) { builder.clearTraceparent(); } else { @@ -79,12 +80,12 @@ public Invocation setTraceparent(String traceparent) { } @Override - public String getTraceparent() { + public @Nullable String getTraceparent() { return builder.hasTraceparent() ? builder.getTraceparent() : null; } @Override - public Invocation setTracestate(String tracestate) { + public Invocation setTracestate(@Nullable String tracestate) { if (tracestate == null) { builder.clearTracestate(); } else { @@ -94,7 +95,7 @@ public Invocation setTracestate(String tracestate) { } @Override - public String getTracestate() { + public @Nullable String getTracestate() { return builder.hasTracestate() ? builder.getTracestate() : null; } @@ -102,37 +103,37 @@ public String getTracestate() { // lives once in InvocationMetadataImpl (against the shared builder); these only refine the type. @Override - public Invocation setServiceName(String serviceName) { + public Invocation setServiceName(@Nullable String serviceName) { super.setServiceName(serviceName); return this; } @Override - public Invocation setHandlerName(String handlerName) { + public Invocation setHandlerName(@Nullable String handlerName) { super.setHandlerName(handlerName); return this; } @Override - public Invocation setKey(String key) { + public Invocation setKey(@Nullable String key) { super.setKey(key); return this; } @Override - public Invocation setScope(String scope) { + public Invocation setScope(@Nullable String scope) { super.setScope(scope); return this; } @Override - public Invocation setLimitKey(String limitKey) { + public Invocation setLimitKey(@Nullable String limitKey) { super.setLimitKey(limitKey); return this; } @Override - public Invocation setIdempotencyKey(String idempotencyKey) { + public Invocation setIdempotencyKey(@Nullable String idempotencyKey) { super.setIdempotencyKey(idempotencyKey); return this; } @@ -144,7 +145,7 @@ public Invocation putHeader(String key, String value) { } @Override - public Invocation setHeaders(Map headers) { + public Invocation setHeaders(@Nullable Map headers) { super.setHeaders(headers); return this; } diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java index d827f3b5..3e79cb02 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadata.java @@ -9,12 +9,16 @@ package dev.restate.integration; import java.util.Map; +import org.jspecify.annotations.Nullable; /** * Invocation metadata. * *

When used as producer defaults, these will be used for all invocations sent through that * producer. + * + *

Passing {@code null} to a setter clears that field. Getters return {@code null} when the + * corresponding field is not set. */ @org.jetbrains.annotations.ApiStatus.Experimental public sealed interface InvocationMetadata permits Invocation, InvocationMetadataImpl { @@ -25,40 +29,40 @@ static InvocationMetadata create() { } /** Target service name. */ - InvocationMetadata setServiceName(String serviceName); + InvocationMetadata setServiceName(@Nullable String serviceName); - String getServiceName(); + @Nullable String getServiceName(); /** Target handler name. */ - InvocationMetadata setHandlerName(String handlerName); + InvocationMetadata setHandlerName(@Nullable String handlerName); - String getHandlerName(); + @Nullable String getHandlerName(); /** Target key (required when the target is a Virtual Object or Workflow). */ - InvocationMetadata setKey(String key); + InvocationMetadata setKey(@Nullable String key); - String getKey(); + @Nullable String getKey(); /** Scope. */ - InvocationMetadata setScope(String scope); + InvocationMetadata setScope(@Nullable String scope); - String getScope(); + @Nullable String getScope(); /** Rate/concurrency limit key. */ - InvocationMetadata setLimitKey(String limitKey); + InvocationMetadata setLimitKey(@Nullable String limitKey); - String getLimitKey(); + @Nullable String getLimitKey(); /** Idempotency key used by Restate to deduplicate the invocation. */ - InvocationMetadata setIdempotencyKey(String idempotencyKey); + InvocationMetadata setIdempotencyKey(@Nullable String idempotencyKey); - String getIdempotencyKey(); + @Nullable String getIdempotencyKey(); /** Add or replace a single header. */ InvocationMetadata putHeader(String key, String value); /** Replace the whole header map. */ - InvocationMetadata setHeaders(Map headers); + InvocationMetadata setHeaders(@Nullable Map headers); Map getHeaders(); } diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java index b01a1ccc..9d3294ab 100644 --- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java +++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java @@ -11,6 +11,7 @@ import dev.restate.ingestion.v1.IngestionDefaults; import dev.restate.ingestion.v1.IngestionInvocation; import java.util.Map; +import org.jspecify.annotations.Nullable; /** * Mutable {@link InvocationMetadata} backed directly by an {@link IngestionInvocation.Builder}, so @@ -47,7 +48,7 @@ static InvocationMetadataImpl fromDefaults(IngestionDefaults defaults) { } @Override - public InvocationMetadata setServiceName(String serviceName) { + public InvocationMetadata setServiceName(@Nullable String serviceName) { if (serviceName == null) { builder.clearService(); } else { @@ -57,12 +58,12 @@ public InvocationMetadata setServiceName(String serviceName) { } @Override - public String getServiceName() { + public @Nullable String getServiceName() { return builder.hasService() ? builder.getService() : null; } @Override - public InvocationMetadata setHandlerName(String handlerName) { + public InvocationMetadata setHandlerName(@Nullable String handlerName) { if (handlerName == null) { builder.clearHandler(); } else { @@ -72,12 +73,12 @@ public InvocationMetadata setHandlerName(String handlerName) { } @Override - public String getHandlerName() { + public @Nullable String getHandlerName() { return builder.hasHandler() ? builder.getHandler() : null; } @Override - public InvocationMetadata setKey(String key) { + public InvocationMetadata setKey(@Nullable String key) { if (key == null) { builder.clearKey(); } else { @@ -87,12 +88,12 @@ public InvocationMetadata setKey(String key) { } @Override - public String getKey() { + public @Nullable String getKey() { return builder.hasKey() ? builder.getKey() : null; } @Override - public InvocationMetadata setScope(String scope) { + public InvocationMetadata setScope(@Nullable String scope) { if (scope == null) { builder.clearScope(); } else { @@ -102,12 +103,12 @@ public InvocationMetadata setScope(String scope) { } @Override - public String getScope() { + public @Nullable String getScope() { return builder.hasScope() ? builder.getScope() : null; } @Override - public InvocationMetadata setLimitKey(String limitKey) { + public InvocationMetadata setLimitKey(@Nullable String limitKey) { if (limitKey == null) { builder.clearLimitKey(); } else { @@ -117,12 +118,12 @@ public InvocationMetadata setLimitKey(String limitKey) { } @Override - public String getLimitKey() { + public @Nullable String getLimitKey() { return builder.hasLimitKey() ? builder.getLimitKey() : null; } @Override - public InvocationMetadata setIdempotencyKey(String idempotencyKey) { + public InvocationMetadata setIdempotencyKey(@Nullable String idempotencyKey) { if (idempotencyKey == null) { builder.clearIdempotencyKey(); } else { @@ -132,7 +133,7 @@ public InvocationMetadata setIdempotencyKey(String idempotencyKey) { } @Override - public String getIdempotencyKey() { + public @Nullable String getIdempotencyKey() { return builder.hasIdempotencyKey() ? builder.getIdempotencyKey() : null; } @@ -143,7 +144,7 @@ public InvocationMetadata putHeader(String key, String value) { } @Override - public InvocationMetadata setHeaders(Map headers) { + public InvocationMetadata setHeaders(@Nullable Map headers) { builder.clearAdditionalHeaders(); if (headers != null) { builder.putAllAdditionalHeaders(headers); diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java index c6bbef11..4b207bca 100644 --- a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java +++ b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java @@ -11,6 +11,7 @@ import dev.restate.ingestion.v1.IngestionDefaults; import java.time.Duration; import java.util.Objects; +import org.jspecify.annotations.Nullable; /** Configuration shared by at-least-once and exactly-once producers. */ @org.jetbrains.annotations.ApiStatus.Experimental @@ -72,7 +73,7 @@ IngestionDefaults toDefaults() { public static final class Builder { private long bufferMemory = DEFAULT_BUFFER_MEMORY; private Duration maxBlockTime = DEFAULT_MAX_BLOCK_TIME; - private InvocationMetadata defaultMetadata; + private @Nullable InvocationMetadata defaultMetadata; private Builder() {} diff --git a/integration-client/src/main/java/dev/restate/integration/SendAttempt.java b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java index 6d59e5b3..416e2ff3 100644 --- a/integration-client/src/main/java/dev/restate/integration/SendAttempt.java +++ b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java @@ -10,6 +10,7 @@ import java.util.Objects; import java.util.concurrent.CompletableFuture; +import org.jspecify.annotations.Nullable; /** Result of a non-blocking producer admission attempt. */ @org.jetbrains.annotations.ApiStatus.Experimental @@ -32,7 +33,7 @@ record Accepted(CompletableFuture acknowledgement) implements SendAt * * @param ready future completed when retrying may succeed */ - record Backpressured(CompletableFuture ready) implements SendAttempt { + record Backpressured(CompletableFuture<@Nullable Void> ready) implements SendAttempt { public Backpressured { Objects.requireNonNull(ready, "ready"); } diff --git a/integration-client/src/main/java/dev/restate/integration/package-info.java b/integration-client/src/main/java/dev/restate/integration/package-info.java new file mode 100644 index 00000000..a3271a8c --- /dev/null +++ b/integration-client/src/main/java/dev/restate/integration/package-info.java @@ -0,0 +1,13 @@ +// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH +// +// This file is part of the Restate Java SDK, +// which is released under the MIT license. +// +// You can find a copy of the license in file LICENSE in the root +// directory of this repository or package, or at +// https://github.com/restatedev/sdk-java/blob/main/LICENSE +/** APIs for producing invocations to Restate through the ingestion protocol. */ +@NullMarked +package dev.restate.integration; + +import org.jspecify.annotations.NullMarked; diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java index ed5ce170..63b05489 100644 --- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java +++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java @@ -87,7 +87,6 @@ void producerOptionsHaveKafkaCompatibleDefaultsAndSnapshotMetadata() throws Exce assertThat(options.bufferMemory()).isEqualTo(32L * 1024 * 1024); assertThat(options.maxBlockTime()).isEqualTo(Duration.ofMinutes(1)); - assertThat(options.defaultMetadata().getServiceName()).isEqualTo("Original"); client.newProducer(options); assertThat(fake.take().getStart().getDefaults().getService()).isEqualTo("Original"); From 70d93b8330e3a0471ebd9860158efc61fd8e97a2 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 12:11:27 +0200 Subject: [PATCH 07/17] Better javadocs --- .../integration/ExactlyOnceProducer.java | 72 +++++++++++++------ .../dev/restate/integration/Producer.java | 63 ++++++++-------- 2 files changed, 83 insertions(+), 52 deletions(-) diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java index c6deefda..cdee4b1d 100644 --- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java @@ -11,40 +11,66 @@ import java.util.concurrent.CompletableFuture; /** - * Like {@link Producer}, but with exactly-once semantics. + * Sends invocations to Restate with exactly-once deduplication. * - *

Exactly once

+ * 
{@code
+ * try (IntegrationClient client = IntegrationClient.builder("http://localhost:8080").build();
+ *     ExactlyOnceProducer producer =
+ *         client.newExactlyOnceProducer("group-a/orders/0")) {
+ *   long offset = checkpoint.load() + 1;
+ *   producer.send(
+ *       offset,
+ *       Invocation.create()
+ *           .setServiceName("Orders")
+ *           .setHandlerName("ingest")
+ *           .setBody(payload));
+ *   checkpoint.store(producer.flush());
+ * }
+ * }
* - * Pick a producer id that is stable across restarts and distinct per independent offset - * sequence. E.g., for a Kafka consumer {@code groupId/topic/partition}, for Postgres logical - * replication the slot name. Because deduplication happens on {@code (producerId, offset)}, it is - * then safe to replay from your last checkpoint after a crash: already-committed offsets are - * dropped, and {@link #flush()} / {@link #waitAcknowledged(long)} reports how far Restate has - * durably caught up so you can advance the checkpoint. After a stream failure, {@link - * #lastAcknowledgedOffset()} remains available so you can determine where to resume. + *

Buffering

* - *

Sending

+ * {@link #send} first admits the invocation to a local buffer, bounded by {@link + * ProducerOptions#bufferMemory()}, while it waits to be handed to the transport. If the buffer is + * full, {@code send} waits up to {@link ProducerOptions#maxBlockTime()} and then throws {@link + * ProducerBufferExhaustedException}. The returned future tracks durable acknowledgement, not buffer + * admission. Send several invocations without awaiting each future, then use {@link #flush()} or + * {@link #flushAsync()} to await them in bulk. {@link #close()} does not flush. * - * {@link #send} admits the record into a byte-bounded local buffer, waiting up to {@link - * ProducerOptions#maxBlockTime()} for capacity, and returns a future that completes once Restate - * has durably committed it. Use {@link #trySend} when the calling thread must never block. + *

Non-blocking admission

* - *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To - * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in - * bulk. + * For event-loop or callback-based code, {@link #trySend} does not wait for buffer capacity. {@link + * SendAttempt.Accepted} contains the durable-acknowledgement future. On {@link + * SendAttempt.Backpressured}, use {@link SendAttempt.Backpressured#ready()} to schedule a retry of + * the same offset and invocation on the event loop; readiness is a notification, not a capacity + * reservation. * * 

{@code
- * producer.send(lsn, Invocation.create().setBody(payload));
- * long committed = producer.flush();
- * checkpoint.store(committed);
+ * static CompletableFuture sendWithoutBlocking(
+ *     ExactlyOnceProducer producer, long offset, Invocation invocation, Executor eventLoop) {
+ *   SendAttempt attempt = producer.trySend(offset, invocation);
+ *   if (attempt instanceof SendAttempt.Accepted accepted) {
+ *     return accepted.acknowledgement();
+ *   }
+ *   return ((SendAttempt.Backpressured) attempt)
+ *       .ready()
+ *       .thenComposeAsync(
+ *           ignored -> sendWithoutBlocking(producer, offset, invocation, eventLoop), eventLoop);
+ * }
  * }
* - *

{@link #close()} does not flush. Call {@link #flush()} before closing, or await {@link - * #flushAsync()}, when accepted invocations must be durably committed. + *

Producer identity and deduplication

+ * + * Each invocation has a strictly increasing offset. Restate deduplicates on {@code (producerId, + * offset)}, so choose a producer id that is stable across restarts and distinct per + * independent offset sequence: for example, a Kafka {@code groupId/topic/partition} or a + * Postgres logical-replication slot. * - *

Thread safety

+ *

After a crash, replay from the last checkpoint; Restate drops already-committed offsets. + * {@link #flush()} and {@link #waitAcknowledged(long)} report the durable watermark to checkpoint, + * and {@link #lastAcknowledgedOffset()} remains available after a stream failure. * - * A producer is not thread-safe and fails fast with {@link + *

A producer is not thread-safe and fails fast with {@link * java.util.ConcurrentModificationException} if used from more than one thread at once. */ @org.jetbrains.annotations.ApiStatus.Experimental diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java index c5995f6b..86c0e8d5 100644 --- a/integration-client/src/main/java/dev/restate/integration/Producer.java +++ b/integration-client/src/main/java/dev/restate/integration/Producer.java @@ -11,49 +11,54 @@ import java.util.concurrent.CompletableFuture; /** - * An at-least-once producer: the client assigns a monotonically increasing offset to each - * invocation. Deduplication is disabled (empty producer id); add an idempotency key on the - * invocations if you need handler-level dedup. - * - *

Sending

- * - * {@link #send} admits the record into a byte-bounded local buffer, waiting up to {@link - * ProducerOptions#maxBlockTime()} for capacity, and returns a future that completes once Restate - * has durably committed it. Use {@link #trySend} when the calling thread must never block. - * - *

Awaiting each {@code send} future before the next send serializes to one in-flight record. To - * parallelize sending, just keep {@code send}ing and use {@link #flush} to await durability in - * bulk. + * Sends invocations to Restate with at-least-once delivery. * * 

{@code
  * try (IntegrationClient client = IntegrationClient.builder("http://localhost:8080").build();
  *     Producer producer = client.newProducer()) {
- *   for (byte[] payload : payloads) {
- *     producer.send(Invocation.create().setBody(payload));
- *   }
- *   producer.flush(); // block until everything sent so far is durably committed
+ *   producer.send(
+ *       Invocation.create()
+ *           .setServiceName("Greeter")
+ *           .setHandlerName("greet")
+ *           .setBody(payload));
+ *   producer.flush();
  * }
  * }
* - *

{@link #close()} does not flush. Call {@link #flush()} before closing, or await {@link - * #flushAsync()}, when accepted invocations must be durably committed. + *

Buffering

* - *

Stream defaults

+ * {@link #send} first admits the invocation to a local buffer, bounded by {@link + * ProducerOptions#bufferMemory()}, while it waits to be handed to the transport. If the buffer is + * full, {@code send} waits up to {@link ProducerOptions#maxBlockTime()} and then throws {@link + * ProducerBufferExhaustedException}. The returned future tracks durable acknowledgement, not buffer + * admission. Send several invocations without awaiting each future, then use {@link #flush()} or + * {@link #flushAsync()} to await them in bulk. {@link #close()} does not flush. * - * Pass an {@link InvocationMetadata} to {@link IntegrationClient#newProducer(InvocationMetadata)}, - * or set {@link ProducerOptions.Builder#defaultMetadata(InvocationMetadata)}, to configure fields - * shared by every record (e.g. the target service/handler) once; per-invocation fields override - * them. + *

Non-blocking admission

+ * + * For event-loop or callback-based code, {@link #trySend} does not wait for buffer capacity. {@link + * SendAttempt.Accepted} contains the durable-acknowledgement future. On {@link + * SendAttempt.Backpressured}, use {@link SendAttempt.Backpressured#ready()} to schedule a retry on + * the event loop; readiness is a notification, not a capacity reservation. * * 
{@code
- * Producer producer =
- *     client.newProducer(
- *         InvocationMetadata.create().setServiceName("Greeter").setHandlerName("greet"));
+ * static CompletableFuture sendWithoutBlocking(
+ *     Producer producer, Invocation invocation, Executor eventLoop) {
+ *   SendAttempt attempt = producer.trySend(invocation);
+ *   if (attempt instanceof SendAttempt.Accepted accepted) {
+ *     return accepted.acknowledgement();
+ *   }
+ *   return ((SendAttempt.Backpressured) attempt)
+ *       .ready()
+ *       .thenComposeAsync(
+ *           ignored -> sendWithoutBlocking(producer, invocation, eventLoop), eventLoop);
+ * }
  * }
* - *

Thread safety

+ *

The client assigns monotonically increasing offsets. Producer-level deduplication is disabled; + * set an idempotency key on an invocation when handler-level deduplication is required. * - * A producer is not thread-safe and fails fast with {@link + *

A producer is not thread-safe and fails fast with {@link * java.util.ConcurrentModificationException} if used from more than one thread at once. */ @org.jetbrains.annotations.ApiStatus.Experimental From 4473d564724b0a2a55231aed8f055b4924b92db9 Mon Sep 17 00:00:00 2001 From: slinkydeveloper Date: Mon, 24 Aug 2026 12:42:23 +0200 Subject: [PATCH 08/17] blabla --- .../restate/integration/AbstractProducer.java | 77 +++++++++++++------ .../integration/ExactlyOnceProducer.java | 38 ++++----- .../dev/restate/integration/Producer.java | 38 ++++----- .../ProducerBufferExhaustedException.java | 6 +- .../restate/integration/ProducerOptions.java | 20 +++-- .../dev/restate/integration/SendAttempt.java | 4 +- .../integration/IntegrationClientTest.java | 76 +++++++++++++++++- 7 files changed, 185 insertions(+), 74 deletions(-) diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java index 6f1ad524..a4467b93 100644 --- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java @@ -34,9 +34,11 @@ * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and * the module docs for the concurrency contract. * - *

Accepted records wait in a byte-bounded queue until Restate's send-window has credit ({@code - * budget}) and the transport is writable ({@code callObserver.isReady()}). Once handed to gRPC, - * only their acknowledgement futures remain until the commit watermark passes their offsets. + *

When buffering is enabled, accepted records wait in a byte-bounded queue until Restate's + * send-window has credit ({@code budget}) and the transport is writable ({@code + * callObserver.isReady()}). With buffering disabled, records are accepted only when they can be + * handed directly to gRPC. Once handed off, only their acknowledgement futures remain until the + * commit watermark passes their offsets. * *

Two-tier concurrency: * @@ -64,7 +66,7 @@ abstract class AbstractProducer implements ProducerBase { private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet private final ArrayDeque bufferedSends = new ArrayDeque<>(); private long bufferedBytes = 0; - private final List capacityWaiters = new ArrayList<>(); + private final List admissionWaiters = new ArrayList<>(); private final TreeMap>> ackWaiters = new TreeMap<>(); private boolean closed = false; private @Nullable IntegrationClientException failure; @@ -192,7 +194,7 @@ public void close() { // ---- send path, shared by the subclasses (caller holds the guard) ---- - /** Admit a record, blocking up to the configured maximum when the local buffer is full. */ + /** Admit a record, blocking up to the configured maximum while the producer is backpressured. */ final CompletableFuture doSend(long offset, InvocationImpl invocation) throws ProducerBufferExhaustedException { PreparedSend prepared = prepare(offset, invocation); @@ -201,7 +203,7 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio long waitStarted = System.nanoTime(); synchronized (lock) { ensureOpenLocked(); - while (!hasCapacityLocked(prepared.bufferSize())) { + while (!canAdmitLocked(prepared.bufferSize())) { if (maxBlockNanos == 0) { throw admissionTimeout(); } @@ -216,7 +218,7 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio } catch (InterruptedException e) { Thread.currentThread().interrupt(); throw new ProducerBufferExhaustedException( - "interrupted while waiting for producer buffer capacity", e); + "interrupted while waiting for producer admission", e); } ensureOpenLocked(); } @@ -227,25 +229,25 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio return acknowledgement; } - /** Attempt to admit a record without blocking or consuming an offset when capacity is absent. */ + /** Attempt to admit a record without blocking or consuming an offset under backpressure. */ final SendAttempt doTrySend(long offset, InvocationImpl invocation) { PreparedSend prepared = prepare(offset, invocation); List> ready; SendAttempt result; synchronized (lock) { ensureOpenLocked(); - if (hasCapacityLocked(prepared.bufferSize())) { + if (canAdmitLocked(prepared.bufferSize())) { result = new SendAttempt.Accepted(acceptLocked(prepared)); ready = drainLocked(); } else { CompletableFuture<@Nullable Void> future = new CompletableFuture<>(); - CapacityWaiter waiter = new CapacityWaiter(prepared.bufferSize(), future); - capacityWaiters.add(waiter); + AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future); + admissionWaiters.add(waiter); future.whenComplete( (ignored, failure) -> { if (future.isCancelled()) { synchronized (lock) { - capacityWaiters.remove(waiter); + admissionWaiters.remove(waiter); } } }); @@ -263,7 +265,7 @@ private PreparedSend prepare(long offset, InvocationImpl invocation) { IngestionRequest request = IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build(); long bufferSize = request.getSerializedSize(); - if (bufferSize > bufferMemory) { + if (bufferMemory > 0 && bufferSize > bufferMemory) { throw new IllegalArgumentException( "serialized invocation requires " + bufferSize @@ -274,22 +276,46 @@ private PreparedSend prepare(long offset, InvocationImpl invocation) { offset, request, request.getInvocation().getSerializedSize(), bufferSize); } - private boolean hasCapacityLocked(long requiredBytes) { + private boolean canAdmitLocked(long requiredBytes) { + if (bufferMemory == 0) { + ClientCallStreamObserver observer = + Objects.requireNonNull(callObserver, "gRPC request stream was not initialized"); + return budget > 0 && observer.isReady(); + } return requiredBytes <= bufferMemory - bufferedBytes; } private CompletableFuture acceptLocked(PreparedSend prepared) { - bufferedSends.addLast( - new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize())); - bufferedBytes += prepared.bufferSize(); lastSent = prepared.offset(); CompletableFuture committed = new CompletableFuture<>(); ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed); + if (bufferMemory == 0) { + budget -= prepared.windowDebit(); + Objects.requireNonNull(callObserver, "gRPC request stream was not initialized") + .onNext(prepared.request()); + } else { + bufferedSends.addLast( + new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize())); + bufferedBytes += prepared.bufferSize(); + } return committed.thenApply(ignored -> new SendResultImpl(prepared.offset())); } /** Write as many queued records as transport and protocol flow control currently permit. */ private List> drainLocked() { + if (bufferMemory == 0) { + if (!canAdmitLocked(0)) { + return List.of(); + } + lock.notifyAll(); + List> ready = new ArrayList<>(admissionWaiters.size()); + for (AdmissionWaiter waiter : admissionWaiters) { + ready.add(waiter.future()); + } + admissionWaiters.clear(); + return ready; + } + boolean freedCapacity = false; ClientCallStreamObserver observer = Objects.requireNonNull(callObserver, "gRPC request stream was not initialized"); @@ -308,8 +334,8 @@ private CompletableFuture acceptLocked(PreparedSend prepared) { lock.notifyAll(); long available = bufferMemory - bufferedBytes; List> ready = new ArrayList<>(); - for (Iterator it = capacityWaiters.iterator(); it.hasNext(); ) { - CapacityWaiter waiter = it.next(); + for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) { + AdmissionWaiter waiter = it.next(); if (waiter.requiredBytes() <= available) { ready.add(waiter.future()); it.remove(); @@ -339,8 +365,9 @@ private void ensureOpenLocked() { } private ProducerBufferExhaustedException admissionTimeout() { - return new ProducerBufferExhaustedException( - "producer buffer remained full for " + maxBlockTime); + String condition = + bufferMemory == 0 ? "producer remained backpressured" : "producer buffer remained full"; + return new ProducerBufferExhaustedException(condition + " for " + maxBlockTime); } private static long awaitFlush(CompletableFuture flush) { @@ -426,11 +453,11 @@ private void terminate(IntegrationClientException cause, boolean halfClose) { } closed = true; failure = cause; - capacity = new ArrayList<>(capacityWaiters.size()); - for (CapacityWaiter waiter : capacityWaiters) { + capacity = new ArrayList<>(admissionWaiters.size()); + for (AdmissionWaiter waiter : admissionWaiters) { capacity.add(waiter.future()); } - capacityWaiters.clear(); + admissionWaiters.clear(); bufferedSends.clear(); bufferedBytes = 0; lock.notifyAll(); @@ -537,7 +564,7 @@ private record PreparedSend( private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {} - private record CapacityWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {} + private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {} private record SendResultImpl(long offset) implements SendResult {} } diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java index cdee4b1d..47acfdcd 100644 --- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java +++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java @@ -30,19 +30,20 @@ * *

Buffering

* - * {@link #send} first admits the invocation to a local buffer, bounded by {@link - * ProducerOptions#bufferMemory()}, while it waits to be handed to the transport. If the buffer is - * full, {@code send} waits up to {@link ProducerOptions#maxBlockTime()} and then throws {@link - * ProducerBufferExhaustedException}. The returned future tracks durable acknowledgement, not buffer - * admission. Send several invocations without awaiting each future, then use {@link #flush()} or - * {@link #flushAsync()} to await them in bulk. {@link #close()} does not flush. + * With a positive {@link ProducerOptions#bufferMemory()}, {@link #send} first admits the invocation + * to a byte-bounded local buffer while it waits to be handed to the transport. A value of zero + * disables local buffering, so {@code send} instead waits until the protocol and transport are + * writable. Either wait is bounded by {@link ProducerOptions#maxBlockTime()}. The returned future + * tracks durable acknowledgement, not admission. Send several invocations without awaiting each + * future, then use {@link #flush()} or {@link #flushAsync()} to await them in bulk. {@link + * #close()} does not flush. * *

Non-blocking admission

* - * For event-loop or callback-based code, {@link #trySend} does not wait for buffer capacity. {@link + * For event-loop or callback-based code, {@link #trySend} does not wait for admission. {@link * SendAttempt.Accepted} contains the durable-acknowledgement future. On {@link * SendAttempt.Backpressured}, use {@link SendAttempt.Backpressured#ready()} to schedule a retry of - * the same offset and invocation on the event loop; readiness is a notification, not a capacity + * the same offset and invocation on the event loop; readiness is a notification, not an admission * reservation. * * 
{@code
@@ -79,9 +80,10 @@ public interface ExactlyOnceProducer extends ProducerBase {
   /**
    * Sends an invocation at {@code offset}.
    *
-   * 
If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()}
-   * for capacity. The invocation is refused with {@link ProducerBufferExhaustedException} if the
-   * timeout elapses. A zero duration makes this method fail immediately under backpressure.
+   * 
This method waits up to {@link ProducerOptions#maxBlockTime()} when the local buffer is
+   * full, or, when buffering is disabled, until protocol and transport readiness permit a direct
+   * write. The invocation is refused with {@link ProducerBufferExhaustedException} if the timeout
+   * elapses. A zero duration makes this method fail immediately under backpressure.
    *
    * 
The returned future completes when the invocation is durably committed by Restate.
    *
@@ -90,11 +92,11 @@ public interface ExactlyOnceProducer extends ProducerBase {
    * @param invocation the invocation to send
    * @return a future completing, once the record is durably committed by Restate, with the {@link
    *     SendResult} carrying {@code offset}
-   * @throws ProducerBufferExhaustedException if buffer capacity does not become available before
-   *     the configured maximum blocking time elapses, or the thread is interrupted while waiting
+   * @throws ProducerBufferExhaustedException if the producer cannot admit the invocation before the
+   *     configured maximum blocking time elapses, or the thread is interrupted while waiting
    * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link
-   *     #lastSentOffset()}, or the serialized invocation is larger than {@link
-   *     ProducerOptions#bufferMemory()}
+   *     #lastSentOffset()}, or buffering is enabled and the serialized invocation is larger than
+   *     {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
@@ -106,14 +108,14 @@ CompletableFuture send(long offset, Invocation invocation)
    *
    * 
An {@link SendAttempt.Accepted} carries the durable-acknowledgement future. A {@link
    * SendAttempt.Backpressured} carries a future that completes when retrying may succeed; the
-   * notification does not reserve capacity.
+   * notification does not reserve admission.
    *
    * @param offset the offset to assign; must be strictly greater than the previous accepted offset
    * @param invocation the invocation to send
    * @return the admission result
    * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link
-   *     #lastSentOffset()}, or the serialized invocation is larger than {@link
-   *     ProducerOptions#bufferMemory()}
+   *     #lastSentOffset()}, or buffering is enabled and the serialized invocation is larger than
+   *     {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java
index 86c0e8d5..55eeebf0 100644
--- a/integration-client/src/main/java/dev/restate/integration/Producer.java
+++ b/integration-client/src/main/java/dev/restate/integration/Producer.java
@@ -27,19 +27,20 @@
  *
  * 
Buffering

  *
- * {@link #send} first admits the invocation to a local buffer, bounded by {@link
- * ProducerOptions#bufferMemory()}, while it waits to be handed to the transport. If the buffer is
- * full, {@code send} waits up to {@link ProducerOptions#maxBlockTime()} and then throws {@link
- * ProducerBufferExhaustedException}. The returned future tracks durable acknowledgement, not buffer
- * admission. Send several invocations without awaiting each future, then use {@link #flush()} or
- * {@link #flushAsync()} to await them in bulk. {@link #close()} does not flush.
+ * With a positive {@link ProducerOptions#bufferMemory()}, {@link #send} first admits the invocation
+ * to a byte-bounded local buffer while it waits to be handed to the transport. A value of zero
+ * disables local buffering, so {@code send} instead waits until the protocol and transport are
+ * writable. Either wait is bounded by {@link ProducerOptions#maxBlockTime()}. The returned future
+ * tracks durable acknowledgement, not admission. Send several invocations without awaiting each
+ * future, then use {@link #flush()} or {@link #flushAsync()} to await them in bulk. {@link
+ * #close()} does not flush.
  *
  * 
Non-blocking admission

  *
- * For event-loop or callback-based code, {@link #trySend} does not wait for buffer capacity. {@link
+ * For event-loop or callback-based code, {@link #trySend} does not wait for admission. {@link
  * SendAttempt.Accepted} contains the durable-acknowledgement future. On {@link
  * SendAttempt.Backpressured}, use {@link SendAttempt.Backpressured#ready()} to schedule a retry on
- * the event loop; readiness is a notification, not a capacity reservation.
+ * the event loop; readiness is a notification, not an admission reservation.
  *
  * 
{@code
  * static CompletableFuture sendWithoutBlocking(
@@ -67,18 +68,19 @@ public interface Producer extends ProducerBase {
   /**
    * Sends an invocation.
    *
-   * 
If the local buffer is full, this method waits up to {@link ProducerOptions#maxBlockTime()}
-   * for capacity. The invocation is refused with {@link ProducerBufferExhaustedException} if the
-   * timeout elapses. A zero duration makes this method fail immediately under backpressure.
+   * 
This method waits up to {@link ProducerOptions#maxBlockTime()} when the local buffer is
+   * full, or, when buffering is disabled, until protocol and transport readiness permit a direct
+   * write. The invocation is refused with {@link ProducerBufferExhaustedException} if the timeout
+   * elapses. A zero duration makes this method fail immediately under backpressure.
    *
    * 
The returned future completes when the invocation is durably committed by Restate.
    *
    * @param invocation the invocation to send
    * @return a future completing, once the invocation is durably committed by Restate.
-   * @throws ProducerBufferExhaustedException if buffer capacity does not become available before
-   *     the configured maximum blocking time elapses, or the thread is interrupted while waiting
-   * @throws IllegalArgumentException if the serialized invocation is larger than {@link
-   *     ProducerOptions#bufferMemory()}
+   * @throws ProducerBufferExhaustedException if the producer cannot admit the invocation before the
+   *     configured maximum blocking time elapses, or the thread is interrupted while waiting
+   * @throws IllegalArgumentException if buffering is enabled and the serialized invocation is
+   *     larger than {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
@@ -89,12 +91,12 @@ public interface Producer extends ProducerBase {
    *
    * 
An {@link SendAttempt.Accepted} carries the durable-acknowledgement future. A {@link
    * SendAttempt.Backpressured} carries a future that completes when retrying may succeed; the
-   * notification does not reserve capacity.
+   * notification does not reserve admission.
    *
    * @param invocation the invocation to send
    * @return the admission result
-   * @throws IllegalArgumentException if the serialized invocation is larger than {@link
-   *     ProducerOptions#bufferMemory()}
+   * @throws IllegalArgumentException if buffering is enabled and the serialized invocation is
+   *     larger than {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java b/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java
index a52e7936..3c01f86f 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerBufferExhaustedException.java
@@ -9,8 +9,10 @@
 package dev.restate.integration;
 
 /**
- * Thrown by {@code send} when local buffer capacity does not become available within the configured
- * {@link ProducerOptions#maxBlockTime()}, or the thread is interrupted while waiting for capacity.
+ * Thrown by {@code send} when the producer cannot admit an invocation within the configured {@link
+ * ProducerOptions#maxBlockTime()}, or the thread is interrupted while waiting. Admission may be
+ * blocked by a full local buffer, or by protocol or transport backpressure when local buffering is
+ * disabled.
  */
 @org.jetbrains.annotations.ApiStatus.Experimental
 public class ProducerBufferExhaustedException extends RuntimeException {
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
index 4b207bca..17607bea 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
@@ -48,7 +48,9 @@ public static Builder builder() {
   }
 
   /**
-   * Maximum serialized bytes retained while invocations wait to be handed to the transport.
+   * Maximum serialized bytes retained while invocations wait to be handed to the transport. A value
+   * of zero disables local buffering: {@code send} waits until the invocation can be handed
+   * directly to the transport, and {@code trySend} reports backpressure until that is possible.
    *
    * @return the local buffer limit in bytes
    */
@@ -57,7 +59,9 @@ public long bufferMemory() {
   }
 
   /**
-   * Maximum time {@code send} waits for buffer capacity before refusing an invocation.
+   * Maximum time {@code send} waits for admission before refusing an invocation. Admission requires
+   * buffer capacity when buffering is enabled, or protocol and transport readiness when {@link
+   * #bufferMemory()} is zero.
    *
    * @return the maximum admission wait
    */
@@ -79,21 +83,21 @@ private Builder() {}
 
     /**
      * Sets the maximum serialized bytes retained while invocations wait to be handed to the
-     * transport.
+     * transport. Set this to zero to disable local buffering.
      *
-     * @param bytes a positive byte count
+     * @param bytes a non-negative byte count
      */
     public Builder bufferMemory(long bytes) {
-      if (bytes <= 0) {
-        throw new IllegalArgumentException("bufferMemory must be greater than zero");
+      if (bytes < 0) {
+        throw new IllegalArgumentException("bufferMemory must not be negative");
       }
       this.bufferMemory = bytes;
       return this;
     }
 
     /**
-     * Sets how long {@code send} waits for buffer capacity. {@link Duration#ZERO} makes {@code
-     * send} fail immediately when the buffer is full.
+     * Sets how long {@code send} waits for admission. {@link Duration#ZERO} makes {@code send} fail
+     * immediately under backpressure.
      */
     public Builder maxBlockTime(Duration duration) {
       Objects.requireNonNull(duration, "maxBlockTime");
diff --git a/integration-client/src/main/java/dev/restate/integration/SendAttempt.java b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java
index 416e2ff3..cbcef7f9 100644
--- a/integration-client/src/main/java/dev/restate/integration/SendAttempt.java
+++ b/integration-client/src/main/java/dev/restate/integration/SendAttempt.java
@@ -28,8 +28,8 @@ record Accepted(CompletableFuture acknowledgement) implements SendAt
   }
 
   /**
-   * The invocation was not accepted because the local buffer was full. {@code ready} completes when
-   * retrying may succeed; it is a notification, not a capacity reservation.
+   * The invocation was not accepted because the producer was backpressured. {@code ready} completes
+   * when retrying may succeed; it is a notification, not an admission reservation.
    *
    * @param ready future completed when retrying may succeed
    */
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 63b05489..8a3e4c67 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -108,7 +108,8 @@ void exactlyOnceProducerAcceptsProducerOptions() throws Exception {
 
   @Test
   void producerOptionsValidateBufferAndBlockTime() {
-    assertThatThrownBy(() -> ProducerOptions.builder().bufferMemory(0))
+    assertThat(ProducerOptions.builder().bufferMemory(0).build().bufferMemory()).isZero();
+    assertThatThrownBy(() -> ProducerOptions.builder().bufferMemory(-1))
         .isInstanceOf(IllegalArgumentException.class);
     assertThatThrownBy(() -> ProducerOptions.builder().maxBlockTime(Duration.ofMillis(-1)))
         .isInstanceOf(IllegalArgumentException.class);
@@ -201,6 +202,79 @@ void sendBuffersBeforeInitialWindowGrant() throws Exception {
     assertThat(get(acknowledgement).offset()).isEqualTo(0L);
   }
 
+  @Test
+  void zeroBufferTrySendWaitsForDirectWriteReadiness() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(0).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    SendAttempt first = producer.trySend(newBody("a".repeat(100)));
+    assertThat(first).isInstanceOf(SendAttempt.Backpressured.class);
+    CompletableFuture ready = ((SendAttempt.Backpressured) first).ready();
+    assertThat(ready).isNotDone();
+    assertThat(producer.lastSentOffset()).isEqualTo(-1L);
+    fake.assertNoRequest();
+
+    // Any positive protocol credit permits one direct write, even when the invocation overshoots
+    // the remaining byte window.
+    fake.grantWindow(1);
+    get(ready);
+
+    SendAttempt.Accepted accepted =
+        (SendAttempt.Accepted) producer.trySend(newBody("a".repeat(100)));
+    assertThat(producer.lastSentOffset()).isEqualTo(0L);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
+
+    // The first invocation exhausted the window, so another direct write is backpressured.
+    assertThat(producer.trySend(newBody("b"))).isInstanceOf(SendAttempt.Backpressured.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(0L);
+
+    fake.ack(0L);
+    assertThat(get(accepted.acknowledgement()).offset()).isEqualTo(0L);
+  }
+
+  @Test
+  void zeroBufferSendBlocksUntilDirectWriteReadiness() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(0).maxBlockTime(Duration.ofSeconds(5)).build());
+    fake.take(); // Start
+
+    CountDownLatch attempting = new CountDownLatch(1);
+    CompletableFuture> blocked =
+        CompletableFuture.supplyAsync(
+            () -> {
+              attempting.countDown();
+              return producer.send(newBody("a"));
+            });
+    assertThat(attempting.await(5, TimeUnit.SECONDS)).isTrue();
+    Thread.sleep(50);
+    assertThat(blocked).isNotDone();
+    fake.assertNoRequest();
+
+    fake.grantWindow(10_000);
+    CompletableFuture acknowledgement = get(blocked);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
+
+    fake.ack(0L);
+    assertThat(get(acknowledgement).offset()).isEqualTo(0L);
+  }
+
+  @Test
+  void zeroBufferAndZeroMaxBlockTimeFailWithoutConsumingOffset() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(0).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    assertThatThrownBy(() -> producer.send(newBody("a")))
+        .isInstanceOf(ProducerBufferExhaustedException.class)
+        .hasMessageContaining("backpressured");
+    assertThat(producer.lastSentOffset()).isEqualTo(-1L);
+    fake.assertNoRequest();
+  }
+
   @Test
   void trySendReportsBackpressureAndSignalsWhenCapacityReturns() throws Exception {
     Producer producer =

From ed4908134176d3d74d68d277c9d955e75e71eb80 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 12:52:45 +0200
Subject: [PATCH 09/17] blabla grpc client

---
 .../integration/GrpcIntegrationClient.java    | 33 ++++++++++++++++++
 .../integration/IntegrationClient.java        | 16 +++++++--
 .../integration/IntegrationClientImpl.java    | 34 +++++++++++++------
 .../integration/IntegrationClientTest.java    | 13 ++++++-
 4 files changed, 82 insertions(+), 14 deletions(-)
 create mode 100644 integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java

diff --git a/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java
new file mode 100644
index 00000000..9cf881fa
--- /dev/null
+++ b/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java
@@ -0,0 +1,33 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.integration;
+
+import io.grpc.Channel;
+import java.util.Objects;
+import org.jetbrains.annotations.ApiStatus;
+
+/**
+ * Internal bridge for first-party integrations that supply their own gRPC {@link Channel}.
+ *
+ * 
The supplied channel remains owned by the caller and is not shut down when the client is
+ * closed.
+ *
+ * @hidden
+ */
+@ApiStatus.Internal
+public final class GrpcIntegrationClient {
+
+  private GrpcIntegrationClient() {}
+
+  public static IntegrationClient.Builder builder(Channel channel) {
+    Objects.requireNonNull(channel, "channel");
+    return new IntegrationClient.Builder(
+        (authToken, integration) -> IntegrationClientImpl.create(channel, authToken, integration));
+  }
+}
diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
index 04779444..9c64aacf 100644
--- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
+++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
@@ -95,12 +95,22 @@ static Builder builder(String ingressUrl) {
 
   /** Builder for {@link IntegrationClient}. */
   final class Builder {
-    private final String target;
+    @FunctionalInterface
+    interface Factory {
+      IntegrationClient create(@Nullable String authToken, String integration);
+    }
+
+    private final Factory factory;
     private @Nullable String authToken;
     private String integration = Version.INTEGRATION;
 
     private Builder(String target) {
-      this.target = target;
+      this(
+          (authToken, integration) -> IntegrationClientImpl.create(target, authToken, integration));
+    }
+
+    Builder(Factory factory) {
+      this.factory = factory;
     }
 
     /** Bearer token sent as the {@code Authorization} header on the ingestion stream. */
@@ -119,7 +129,7 @@ public Builder integration(String name, String version) {
     }
 
     public IntegrationClient build() {
-      return IntegrationClientImpl.create(target, authToken, integration);
+      return factory.create(authToken, integration);
     }
   }
 }
diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
index c1a57ce9..ef96a717 100644
--- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
@@ -9,22 +9,25 @@
 package dev.restate.integration;
 
 import dev.restate.ingestion.v1.IngestionSvcGrpc;
+import io.grpc.Channel;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import org.jspecify.annotations.Nullable;
 
-/** {@link IntegrationClient} backed by a single gRPC {@link ManagedChannel} shared by producers. */
+/** {@link IntegrationClient} backed by a single gRPC {@link Channel} shared by producers. */
 final class IntegrationClientImpl implements IntegrationClient {
 
-  private final ManagedChannel channel;
+  private final @Nullable ManagedChannel ownedChannel;
   private final IngestionSvcGrpc.IngestionSvcStub stub;
   private final String integration;
 
   private IntegrationClientImpl(
-      ManagedChannel channel, IngestionSvcGrpc.IngestionSvcStub stub, String integration) {
-    this.channel = channel;
+      @Nullable ManagedChannel ownedChannel,
+      IngestionSvcGrpc.IngestionSvcStub stub,
+      String integration) {
+    this.ownedChannel = ownedChannel;
     this.stub = stub;
     this.integration = integration;
   }
@@ -40,16 +43,23 @@ static IntegrationClient create(String target, @Nullable String authToken, Strin
     }
     ManagedChannel channel = builder.build();
 
+    return create(channel, authToken, integration, channel);
+  }
+
+  static IntegrationClient create(Channel channel, @Nullable String authToken, String integration) {
+    return create(Objects.requireNonNull(channel, "channel"), authToken, integration, null);
+  }
+
+  private static IntegrationClient create(
+      Channel channel,
+      @Nullable String authToken,
+      String integration,
+      @Nullable ManagedChannel ownedChannel) {
     IngestionSvcGrpc.IngestionSvcStub stub = IngestionSvcGrpc.newStub(channel);
     if (authToken != null && !authToken.isBlank()) {
       stub = stub.withInterceptors(new AuthInterceptor(authToken));
     }
-    return new IntegrationClientImpl(channel, stub, integration);
-  }
-
-  /** Visible for testing: build a client over an already-created channel (e.g. gRPC in-process). */
-  static IntegrationClient forChannel(ManagedChannel channel, String integration) {
-    return new IntegrationClientImpl(channel, IngestionSvcGrpc.newStub(channel), integration);
+    return new IntegrationClientImpl(ownedChannel, stub, integration);
   }
 
   @Override
@@ -69,6 +79,10 @@ public ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOpt
 
   @Override
   public void close() {
+    ManagedChannel channel = ownedChannel;
+    if (channel == null) {
+      return;
+    }
     channel.shutdown();
     try {
       if (!channel.awaitTermination(5, TimeUnit.SECONDS)) {
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 8a3e4c67..9f989682 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -53,7 +53,7 @@ void setUp() throws IOException {
     fake = new FakeIngestionService();
     server = InProcessServerBuilder.forName(name).directExecutor().addService(fake).build().start();
     channel = InProcessChannelBuilder.forName(name).directExecutor().build();
-    client = IntegrationClientImpl.forChannel(channel, INTEGRATION);
+    client = GrpcIntegrationClient.builder(channel).integration("test-integration", "1.0").build();
   }
 
   @AfterEach
@@ -61,11 +61,22 @@ void tearDown() {
     if (client != null) {
       client.close();
     }
+    if (channel != null) {
+      channel.shutdownNow();
+    }
     if (server != null) {
       server.shutdownNow();
     }
   }
 
+  @Test
+  void grpcBridgeDoesNotCloseCallerOwnedChannel() {
+    client.close();
+
+    assertThat(channel.isShutdown()).isFalse();
+    client = null;
+  }
+
   @Test
   void producerSendsDisabledDedupHandshake() throws Exception {
     InvocationMetadata defaults = InvocationMetadata.create().setServiceName("Svc");

From 1c1849557da111286971e8cd9479dd2529e00118 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 12:59:45 +0200
Subject: [PATCH 10/17] fix some shit

---
 integration-client/build.gradle.kts           |  4 ++
 .../restate/integration/AbstractProducer.java | 39 ++++++++++++++++---
 .../integration/GrpcIntegrationClient.java    |  5 +--
 .../integration/IntegrationClient.java        | 25 ++++++------
 4 files changed, 52 insertions(+), 21 deletions(-)

diff --git a/integration-client/build.gradle.kts b/integration-client/build.gradle.kts
index bebbc1da..6682b752 100644
--- a/integration-client/build.gradle.kts
+++ b/integration-client/build.gradle.kts
@@ -30,8 +30,12 @@ dependencies {
   testImplementation(libs.junit.jupiter)
   testImplementation(libs.assertj)
   testImplementation(libs.protobuf.java)
+  testImplementation(project(":sdk-api"))
+  testImplementation(project(":sdk-serde-jackson"))
+  testImplementation(project(":sdk-testing"))
   // In-process transport to drive the client against a fake IngestionSvc in unit tests.
   testImplementation(libs.grpc.inprocess)
+  testRuntimeOnly(libs.log4j.core)
   testRuntimeOnly(libs.junit.platform.launcher)
 }
 
diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
index a4467b93..0179011d 100644
--- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
+++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
@@ -200,6 +200,7 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio
     PreparedSend prepared = prepare(offset, invocation);
     List> ready;
     CompletableFuture acknowledgement;
+    boolean directWrite;
     long waitStarted = System.nanoTime();
     synchronized (lock) {
       ensureOpenLocked();
@@ -223,9 +224,15 @@ final CompletableFuture doSend(long offset, InvocationImpl invocatio
         ensureOpenLocked();
       }
       acknowledgement = acceptLocked(prepared);
-      ready = drainLocked();
+      directWrite = bufferMemory == 0;
+      ready = directWrite ? List.of() : drainLocked();
+    }
+    if (directWrite) {
+      writeDirect(prepared.request());
+      drainAndWake();
+    } else {
+      completeReady(ready);
     }
-    completeReady(ready);
     return acknowledgement;
   }
 
@@ -234,11 +241,13 @@ final SendAttempt doTrySend(long offset, InvocationImpl invocation) {
     PreparedSend prepared = prepare(offset, invocation);
     List> ready;
     SendAttempt result;
+    boolean directWrite = false;
     synchronized (lock) {
       ensureOpenLocked();
       if (canAdmitLocked(prepared.bufferSize())) {
         result = new SendAttempt.Accepted(acceptLocked(prepared));
-        ready = drainLocked();
+        directWrite = bufferMemory == 0;
+        ready = directWrite ? List.of() : drainLocked();
       } else {
         CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
         AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future);
@@ -255,7 +264,12 @@ final SendAttempt doTrySend(long offset, InvocationImpl invocation) {
         ready = List.of();
       }
     }
-    completeReady(ready);
+    if (directWrite) {
+      writeDirect(prepared.request());
+      drainAndWake();
+    } else {
+      completeReady(ready);
+    }
     return result;
   }
 
@@ -291,8 +305,6 @@ private CompletableFuture acceptLocked(PreparedSend prepared) {
     ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed);
     if (bufferMemory == 0) {
       budget -= prepared.windowDebit();
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized")
-          .onNext(prepared.request());
     } else {
       bufferedSends.addLast(
           new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize()));
@@ -301,6 +313,21 @@ private CompletableFuture acceptLocked(PreparedSend prepared) {
     return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
   }
 
+  /** Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused. */
+  private void writeDirect(IngestionRequest request) {
+    try {
+      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(request);
+    } catch (RuntimeException e) {
+      IntegrationClientException cause =
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN,
+              "failed to write invocation to the ingestion stream",
+              e);
+      terminate(cause, false);
+      throw cause;
+    }
+  }
+
   /** Write as many queued records as transport and protocol flow control currently permit. */
   private List> drainLocked() {
     if (bufferMemory == 0) {
diff --git a/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java
index 9cf881fa..512da060 100644
--- a/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java
+++ b/integration-client/src/main/java/dev/restate/integration/GrpcIntegrationClient.java
@@ -9,7 +9,6 @@
 package dev.restate.integration;
 
 import io.grpc.Channel;
-import java.util.Objects;
 import org.jetbrains.annotations.ApiStatus;
 
 /**
@@ -26,8 +25,6 @@ public final class GrpcIntegrationClient {
   private GrpcIntegrationClient() {}
 
   public static IntegrationClient.Builder builder(Channel channel) {
-    Objects.requireNonNull(channel, "channel");
-    return new IntegrationClient.Builder(
-        (authToken, integration) -> IntegrationClientImpl.create(channel, authToken, integration));
+    return new IntegrationClient.Builder(channel);
   }
 }
diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
index 9c64aacf..557678c8 100644
--- a/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
+++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClient.java
@@ -8,6 +8,8 @@
 // https://github.com/restatedev/sdk-java/blob/main/LICENSE
 package dev.restate.integration;
 
+import io.grpc.Channel;
+import java.util.Objects;
 import org.jspecify.annotations.Nullable;
 
 /** Entry point for producing invocations to Restate ingress over the ingestion API. */
@@ -95,22 +97,19 @@ static Builder builder(String ingressUrl) {
 
   /** Builder for {@link IntegrationClient}. */
   final class Builder {
-    @FunctionalInterface
-    interface Factory {
-      IntegrationClient create(@Nullable String authToken, String integration);
-    }
-
-    private final Factory factory;
+    private final @Nullable String target;
+    private final @Nullable Channel channel;
     private @Nullable String authToken;
     private String integration = Version.INTEGRATION;
 
     private Builder(String target) {
-      this(
-          (authToken, integration) -> IntegrationClientImpl.create(target, authToken, integration));
+      this.target = target;
+      this.channel = null;
     }
 
-    Builder(Factory factory) {
-      this.factory = factory;
+    Builder(Channel channel) {
+      this.target = null;
+      this.channel = Objects.requireNonNull(channel, "channel");
     }
 
     /** Bearer token sent as the {@code Authorization} header on the ingestion stream. */
@@ -129,7 +128,11 @@ public Builder integration(String name, String version) {
     }
 
     public IntegrationClient build() {
-      return factory.create(authToken, integration);
+      if (channel != null) {
+        return IntegrationClientImpl.create(channel, authToken, integration);
+      }
+      return IntegrationClientImpl.create(
+          Objects.requireNonNull(target, "target"), authToken, integration);
     }
   }
 }

From 38a5d677d28d99fc5e11f93f79d9802abc306b22 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 14:09:44 +0200
Subject: [PATCH 11/17] fixes

---
 .../restate/integration/AbstractProducer.java |   7 +-
 .../IntegrationClientIntegrationTest.java     | 125 ++++++++++++++++++
 .../integration/IntegrationClientTest.java    | 105 ++++++++++++++-
 3 files changed, 234 insertions(+), 3 deletions(-)
 create mode 100644 integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java

diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
index 0179011d..4859fe8f 100644
--- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
+++ b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
@@ -313,10 +313,13 @@ private CompletableFuture acceptLocked(PreparedSend prepared) {
     return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
   }
 
-  /** Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused. */
+  /**
+   * Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused.
+   */
   private void writeDirect(IngestionRequest request) {
     try {
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(request);
+      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized")
+          .onNext(request);
     } catch (RuntimeException e) {
       IntegrationClientException cause =
           new IntegrationClientException(
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java
new file mode 100644
index 00000000..c504d70c
--- /dev/null
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java
@@ -0,0 +1,125 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.integration;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import dev.restate.client.Client;
+import dev.restate.sdk.Restate;
+import dev.restate.sdk.annotation.Handler;
+import dev.restate.sdk.annotation.Name;
+import dev.restate.sdk.annotation.VirtualObject;
+import dev.restate.sdk.common.StateKey;
+import dev.restate.sdk.testing.BindService;
+import dev.restate.sdk.testing.RestateClient;
+import dev.restate.sdk.testing.RestateTest;
+import dev.restate.sdk.testing.RestateURL;
+import java.nio.charset.StandardCharsets;
+import java.time.Duration;
+import java.util.UUID;
+import java.util.concurrent.TimeUnit;
+import org.junit.jupiter.api.Test;
+import org.junit.jupiter.api.Timeout;
+
+@RestateTest(containerImage = "ghcr.io/restatedev/restate:pr5026")
+@Timeout(value = 30)
+class IntegrationClientIntegrationTest {
+
+  private static final String SERVICE = "IntegrationClientCounter";
+  private static final byte[] ONE = "1".getBytes(StandardCharsets.UTF_8);
+
+  @BindService private final IntegrationClientCounter counter = new IntegrationClientCounterImpl();
+
+  @Test
+  void zeroBufferProducerDeliversAndAcknowledges(
+      @RestateURL String ingressUrl, @RestateClient Client ingressClient) throws Exception {
+    String key = UUID.randomUUID().toString();
+
+    try (IntegrationClient client = IntegrationClient.builder(ingressUrl).build();
+        Producer producer =
+            client.newProducer(
+                ProducerOptions.builder()
+                    .bufferMemory(0)
+                    .maxBlockTime(Duration.ofSeconds(10))
+                    .defaultMetadata(counterMetadata(key))
+                    .build())) {
+      SendResult result = producer.send(Invocation.create().setBody(ONE)).get(10, TimeUnit.SECONDS);
+
+      assertThat(result.offset()).isZero();
+      assertThat(producer.lastAcknowledgedOffset()).isZero();
+    }
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(1L);
+  }
+
+  @Test
+  void exactlyOnceProducerDeduplicatesAcrossStreams(
+      @RestateURL String ingressUrl, @RestateClient Client ingressClient) throws Exception {
+    String key = UUID.randomUUID().toString();
+    String producerId = "integration-test/" + UUID.randomUUID();
+    ProducerOptions options =
+        ProducerOptions.builder()
+            .bufferMemory(0)
+            .maxBlockTime(Duration.ofSeconds(10))
+            .defaultMetadata(counterMetadata(key))
+            .build();
+
+    sendExactlyOnce(ingressUrl, producerId, options);
+    sendExactlyOnce(ingressUrl, producerId, options);
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(1L);
+  }
+
+  private static void sendExactlyOnce(String ingressUrl, String producerId, ProducerOptions options)
+      throws Exception {
+    try (IntegrationClient client = IntegrationClient.builder(ingressUrl).build();
+        ExactlyOnceProducer producer = client.newExactlyOnceProducer(producerId, options)) {
+      SendResult result =
+          producer.send(0, Invocation.create().setBody(ONE)).get(10, TimeUnit.SECONDS);
+      assertThat(result.offset()).isZero();
+    }
+  }
+
+  private static InvocationMetadata counterMetadata(String key) {
+    return InvocationMetadata.create()
+        .setServiceName(SERVICE)
+        .setHandlerName("add")
+        .setKey(key)
+        .putHeader("content-type", "application/json");
+  }
+
+  @VirtualObject
+  @Name(SERVICE)
+  public interface IntegrationClientCounter {
+
+    @Handler
+    void add(long value);
+
+    @Handler
+    long get();
+  }
+
+  public static final class IntegrationClientCounterImpl implements IntegrationClientCounter {
+
+    private static final StateKey COUNT = StateKey.of("count", Long.class);
+
+    @Override
+    public void add(long value) {
+      Restate.State state = Restate.state();
+      state.set(COUNT, state.get(COUNT).orElse(0L) + value);
+    }
+
+    @Override
+    public long get() {
+      return Restate.state().get(COUNT).orElse(0L);
+    }
+  }
+}
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 9f989682..3f56b33a 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -18,7 +18,12 @@
 import dev.restate.ingestion.v1.IngestionResponse;
 import dev.restate.ingestion.v1.IngestionSvcGrpc;
 import dev.restate.ingestion.v1.WindowUpdate;
+import io.grpc.CallOptions;
+import io.grpc.Channel;
+import io.grpc.ClientCall;
+import io.grpc.ForwardingClientCall;
 import io.grpc.ManagedChannel;
+import io.grpc.MethodDescriptor;
 import io.grpc.Server;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
@@ -238,9 +243,16 @@ void zeroBufferTrySendWaitsForDirectWriteReadiness() throws Exception {
     assertThat(fake.take().getInvocation().getOffset()).isEqualTo(0L);
 
     // The first invocation exhausted the window, so another direct write is backpressured.
-    assertThat(producer.trySend(newBody("b"))).isInstanceOf(SendAttempt.Backpressured.class);
+    SendAttempt.Backpressured second = (SendAttempt.Backpressured) producer.trySend(newBody("b"));
     assertThat(producer.lastSentOffset()).isEqualTo(0L);
 
+    // Window updates must first repay the overshoot; readiness is signalled only once the budget
+    // becomes positive again.
+    fake.grantWindow(1);
+    assertThat(second.ready()).isNotDone();
+    fake.grantWindow(10_000);
+    get(second.ready());
+
     fake.ack(0L);
     assertThat(get(accepted.acknowledgement()).offset()).isEqualTo(0L);
   }
@@ -286,6 +298,66 @@ void zeroBufferAndZeroMaxBlockTimeFailWithoutConsumingOffset() throws Exception
     fake.assertNoRequest();
   }
 
+  @Test
+  void zeroBufferDirectWriteFailureTerminatesProducer() throws Exception {
+    client.close();
+    client =
+        GrpcIntegrationClient.builder(new FailingSecondWriteChannel(channel))
+            .integration("test-integration", "1.0")
+            .build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    fake.take(); // Start is the first write and succeeds.
+    fake.grantWindow(10_000);
+
+    assertThatThrownBy(() -> producer.send(newBody("a")))
+        .isInstanceOf(IntegrationClientException.class)
+        .hasMessageContaining("failed to write invocation");
+    assertThat(producer.lastSentOffset()).isZero();
+    assertThatThrownBy(() -> get(producer.flushAsync()))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
+    assertThatThrownBy(() -> producer.send(newBody("b")))
+        .isInstanceOf(IllegalStateException.class)
+        .hasCauseInstanceOf(IntegrationClientException.class);
+  }
+
+  @Test
+  void zeroBufferStreamErrorFailsReadinessWaiter() throws Exception {
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    fake.take(); // Start
+
+    SendAttempt.Backpressured backpressured =
+        (SendAttempt.Backpressured) producer.trySend(newBody("a"));
+    fake.error(ErrorKind.ERROR_KIND_GO_AWAY, "go away");
+
+    assertThatThrownBy(() -> get(backpressured.ready()))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class)
+        .extracting(t -> ((IntegrationClientException) t).getKind())
+        .isEqualTo(IntegrationClientException.Kind.GO_AWAY);
+  }
+
+  @Test
+  void exactlyOnceZeroBufferBackpressureDoesNotConsumeOffset() throws Exception {
+    ExactlyOnceProducer producer =
+        client.newExactlyOnceProducer(
+            "p1", ProducerOptions.builder().bufferMemory(0).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    SendAttempt.Backpressured backpressured =
+        (SendAttempt.Backpressured) producer.trySend(5, newBody("a"));
+    assertThatThrownBy(() -> producer.send(5, newBody("a")))
+        .isInstanceOf(ProducerBufferExhaustedException.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(-1L);
+
+    fake.grantWindow(10_000);
+    get(backpressured.ready());
+    assertThat(producer.trySend(5, newBody("a"))).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(5L);
+  }
+
   @Test
   void trySendReportsBackpressureAndSignalsWhenCapacityReturns() throws Exception {
     Producer producer =
@@ -611,4 +683,35 @@ void error(ErrorKind kind, String message, long lastCommitted) {
       responses.onCompleted();
     }
   }
+
+  private static final class FailingSecondWriteChannel extends Channel {
+
+    private final Channel delegate;
+
+    private FailingSecondWriteChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        private int writes;
+
+        @Override
+        public void sendMessage(RequestT message) {
+          if (++writes == 2) {
+            throw new IllegalStateException("simulated transport write failure");
+          }
+          super.sendMessage(message);
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
 }

From 56d3abe5af97771bb35dce3b83524ca8efb74041 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 16:29:20 +0200
Subject: [PATCH 12/17] more tests

---
 .../IntegrationClientIntegrationTest.java     | 140 ++++++++++++++----
 1 file changed, 114 insertions(+), 26 deletions(-)

diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java
index c504d70c..46600362 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientIntegrationTest.java
@@ -20,36 +20,35 @@
 import dev.restate.sdk.testing.RestateClient;
 import dev.restate.sdk.testing.RestateTest;
 import dev.restate.sdk.testing.RestateURL;
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.UUID;
 import java.util.concurrent.TimeUnit;
-import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 @RestateTest(containerImage = "ghcr.io/restatedev/restate:pr5026")
 @Timeout(value = 30)
 class IntegrationClientIntegrationTest {
 
   private static final String SERVICE = "IntegrationClientCounter";
-  private static final byte[] ONE = "1".getBytes(StandardCharsets.UTF_8);
 
   @BindService private final IntegrationClientCounter counter = new IntegrationClientCounterImpl();
 
-  @Test
-  void zeroBufferProducerDeliversAndAcknowledges(
-      @RestateURL String ingressUrl, @RestateClient Client ingressClient) throws Exception {
+  @BufferMemoryTest
+  void producerDeliversAndAcknowledges(
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient)
+      throws Exception {
     String key = UUID.randomUUID().toString();
 
     try (IntegrationClient client = IntegrationClient.builder(ingressUrl).build();
-        Producer producer =
-            client.newProducer(
-                ProducerOptions.builder()
-                    .bufferMemory(0)
-                    .maxBlockTime(Duration.ofSeconds(10))
-                    .defaultMetadata(counterMetadata(key))
-                    .build())) {
-      SendResult result = producer.send(Invocation.create().setBody(ONE)).get(10, TimeUnit.SECONDS);
+        Producer producer = client.newProducer(producerOptions(key, bufferMemory))) {
+      SendResult result = producer.send(invocation(1)).get(10, TimeUnit.SECONDS);
 
       assertThat(result.offset()).isZero();
       assertThat(producer.lastAcknowledgedOffset()).isZero();
@@ -59,35 +58,116 @@ void zeroBufferProducerDeliversAndAcknowledges(
         .isEqualTo(1L);
   }
 
-  @Test
+  @BufferMemoryTest
   void exactlyOnceProducerDeduplicatesAcrossStreams(
-      @RestateURL String ingressUrl, @RestateClient Client ingressClient) throws Exception {
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient) {
     String key = UUID.randomUUID().toString();
     String producerId = "integration-test/" + UUID.randomUUID();
-    ProducerOptions options =
-        ProducerOptions.builder()
-            .bufferMemory(0)
-            .maxBlockTime(Duration.ofSeconds(10))
-            .defaultMetadata(counterMetadata(key))
-            .build();
+    ProducerOptions options = producerOptions(key, bufferMemory);
 
-    sendExactlyOnce(ingressUrl, producerId, options);
-    sendExactlyOnce(ingressUrl, producerId, options);
+    sendExactlyOnce(ingressUrl, producerId, options, record(0, 1));
+    sendExactlyOnce(ingressUrl, producerId, options, record(0, 100));
 
     assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
         .isEqualTo(1L);
   }
 
-  private static void sendExactlyOnce(String ingressUrl, String producerId, ProducerOptions options)
+  @BufferMemoryTest
+  void exactlyOnceProducerReplaysCommittedPrefixAndAcceptsNewOffsets(
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient) {
+    String key = UUID.randomUUID().toString();
+    String producerId = "integration-test/" + UUID.randomUUID();
+    ProducerOptions options = producerOptions(key, bufferMemory);
+
+    sendExactlyOnce(ingressUrl, producerId, options, record(0, 1), record(1, 10));
+    sendExactlyOnce(
+        ingressUrl, producerId, options, record(0, 100), record(1, 1_000), record(2, 10_000));
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(10_011L);
+  }
+
+  @BufferMemoryTest
+  void exactlyOnceProducerDropsOffsetsBelowCommittedWatermark(
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient) {
+    String key = UUID.randomUUID().toString();
+    String producerId = "integration-test/" + UUID.randomUUID();
+    ProducerOptions options = producerOptions(key, bufferMemory);
+
+    sendExactlyOnce(ingressUrl, producerId, options, record(10, 1));
+    sendExactlyOnce(
+        ingressUrl, producerId, options, record(0, 100), record(9, 1_000), record(11, 10));
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(11L);
+  }
+
+  @BufferMemoryTest
+  void exactlyOnceDeduplicationIsScopedByProducerId(
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient) {
+    String key = UUID.randomUUID().toString();
+    ProducerOptions options = producerOptions(key, bufferMemory);
+
+    sendExactlyOnce(ingressUrl, "integration-test/" + UUID.randomUUID(), options, record(0, 1));
+    sendExactlyOnce(ingressUrl, "integration-test/" + UUID.randomUUID(), options, record(0, 10));
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(11L);
+  }
+
+  @BufferMemoryTest
+  void atLeastOnceProducerDoesNotDeduplicateAcrossStreams(
+      long bufferMemory, @RestateURL String ingressUrl, @RestateClient Client ingressClient)
       throws Exception {
+    String key = UUID.randomUUID().toString();
+    ProducerOptions options = producerOptions(key, bufferMemory);
+
+    sendAtLeastOnce(ingressUrl, options, 1);
+    sendAtLeastOnce(ingressUrl, options, 10);
+
+    assertThat(ingressClient.virtualObject(IntegrationClientCounter.class, key).get())
+        .isEqualTo(11L);
+  }
+
+  private static void sendExactlyOnce(
+      String ingressUrl, String producerId, ProducerOptions options, TestRecord... records) {
     try (IntegrationClient client = IntegrationClient.builder(ingressUrl).build();
         ExactlyOnceProducer producer = client.newExactlyOnceProducer(producerId, options)) {
-      SendResult result =
-          producer.send(0, Invocation.create().setBody(ONE)).get(10, TimeUnit.SECONDS);
+      for (TestRecord record : records) {
+        producer.send(record.offset(), invocation(record.value()));
+      }
+
+      long lastOffset = records[records.length - 1].offset();
+      assertThat(producer.flush()).isEqualTo(lastOffset);
+      assertThat(producer.lastAcknowledgedOffset()).isEqualTo(lastOffset);
+    }
+  }
+
+  private static void sendAtLeastOnce(String ingressUrl, ProducerOptions options, long value)
+      throws Exception {
+    try (IntegrationClient client = IntegrationClient.builder(ingressUrl).build();
+        Producer producer = client.newProducer(options)) {
+      SendResult result = producer.send(invocation(value)).get(10, TimeUnit.SECONDS);
       assertThat(result.offset()).isZero();
     }
   }
 
+  private static TestRecord record(long offset, long value) {
+    return new TestRecord(offset, value);
+  }
+
+  private static Invocation invocation(long value) {
+    return Invocation.create().setBody(Long.toString(value).getBytes(StandardCharsets.UTF_8));
+  }
+
+  private static ProducerOptions producerOptions(String key, long bufferMemory) {
+    return ProducerOptions.builder()
+        .bufferMemory(bufferMemory)
+        .maxBlockTime(Duration.ofSeconds(10))
+        .defaultMetadata(counterMetadata(key))
+        .build();
+  }
+
   private static InvocationMetadata counterMetadata(String key) {
     return InvocationMetadata.create()
         .setServiceName(SERVICE)
@@ -96,6 +176,14 @@ private static InvocationMetadata counterMetadata(String key) {
         .putHeader("content-type", "application/json");
   }
 
+  private record TestRecord(long offset, long value) {}
+
+  @Target(ElementType.METHOD)
+  @Retention(RetentionPolicy.RUNTIME)
+  @ParameterizedTest(name = "{displayName}: bufferMemory={0}")
+  @ValueSource(longs = {0L, ProducerOptions.DEFAULT_BUFFER_MEMORY})
+  private @interface BufferMemoryTest {}
+
   @VirtualObject
   @Name(SERVICE)
   public interface IntegrationClientCounter {

From 8147274c6d14356ec4eb5df82bd2527d55b6a361 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 17:08:25 +0200
Subject: [PATCH 13/17] Reducing code

---
 .../restate/integration/AbstractProducer.java | 600 ----------------
 .../restate/integration/AuthInterceptor.java  |  43 --
 .../integration/ExactlyOnceProducerImpl.java  |  55 --
 .../restate/integration/IngressEndpoint.java  |   5 +-
 .../integration/IntegrationClientImpl.java    |  11 +-
 .../integration/InvocationMetadataImpl.java   |  24 -
 .../dev/restate/integration/ProducerImpl.java | 645 +++++++++++++++++-
 .../integration/IntegrationClientTest.java    |  58 +-
 8 files changed, 710 insertions(+), 731 deletions(-)
 delete mode 100644 integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
 delete mode 100644 integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java
 delete mode 100644 integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java

diff --git a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java b/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
deleted file mode 100644
index 4859fe8f..00000000
--- a/integration-client/src/main/java/dev/restate/integration/AbstractProducer.java
+++ /dev/null
@@ -1,600 +0,0 @@
-// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
-//
-// This file is part of the Restate Java SDK,
-// which is released under the MIT license.
-//
-// You can find a copy of the license in file LICENSE in the root
-// directory of this repository or package, or at
-// https://github.com/restatedev/sdk-java/blob/main/LICENSE
-package dev.restate.integration;
-
-import dev.restate.ingestion.v1.DeduplicationMode;
-import dev.restate.ingestion.v1.ErrorKind;
-import dev.restate.ingestion.v1.IngestionRequest;
-import dev.restate.ingestion.v1.IngestionResponse;
-import dev.restate.ingestion.v1.IngestionStart;
-import dev.restate.ingestion.v1.IngestionSvcGrpc;
-import io.grpc.stub.ClientCallStreamObserver;
-import io.grpc.stub.ClientResponseObserver;
-import java.time.Duration;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
-import java.util.ConcurrentModificationException;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.TreeMap;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicReference;
-import org.jspecify.annotations.Nullable;
-
-/**
- * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and
- * the module docs for the concurrency contract.
- *
- * 
When buffering is enabled, accepted records wait in a byte-bounded queue until Restate's
- * send-window has credit ({@code budget}) and the transport is writable ({@code
- * callObserver.isReady()}). With buffering disabled, records are accepted only when they can be
- * handed directly to gRPC. Once handed off, only their acknowledgement futures remain until the
- * commit watermark passes their offsets.
- *
- * 
Two-tier concurrency:
- *
- * 
    
- *   
  • A fail-fast, KafkaConsumer-style guard ({@link #acquire()}/{@link #release()}) rejects
- *       concurrent use from multiple threads; sequential hand-off between threads is fine.
- *   
  • A single monitor ({@link #lock}) guards the small set of fields genuinely shared between
- *       the caller thread and gRPC's callback threads. Futures are always completed outside the
- *       monitor.
- * 

- */
-abstract class AbstractProducer implements ProducerBase {
-
-  private final Object lock = new Object();
-
-  // Set once, synchronously, in beforeStart() before the constructor sends the Start frame.
-  private volatile @Nullable ClientCallStreamObserver callObserver;
-
-  // ---- fail-fast single-thread guard ----
-  private final AtomicReference<@Nullable Thread> owner = new AtomicReference<>();
-  private int reentrancy;
-
-  // ---- state guarded by `lock` ----
-  private long budget = 0; // remaining Restate send window, in bytes; may go one message negative
-  private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet
-  private final ArrayDeque bufferedSends = new ArrayDeque<>();
-  private long bufferedBytes = 0;
-  private final List admissionWaiters = new ArrayList<>();
-  private final TreeMap>> ackWaiters = new TreeMap<>();
-  private boolean closed = false;
-  private @Nullable IntegrationClientException failure;
-
-  private final long bufferMemory;
-  private final Duration maxBlockTime;
-  private final long maxBlockNanos;
-
-  // Written only by the (guarded) caller thread; never touched by gRPC callbacks.
-  long lastSent = -1;
-
-  AbstractProducer(
-      IngestionSvcGrpc.IngestionSvcStub stub,
-      String producerId,
-      DeduplicationMode deduplicationMode,
-      ProducerOptions options,
-      String integration) {
-    this.bufferMemory = options.bufferMemory();
-    this.maxBlockTime = options.maxBlockTime();
-    this.maxBlockNanos = toNanosSaturated(maxBlockTime);
-    // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready
-    // handler.
-    stub.ingest(new ResponseObserver());
-    // Mandatory Start handshake: the first frame on the stream (not flow-controlled).
-    IngestionRequest start =
-        IngestionRequest.newBuilder()
-            .setStart(
-                IngestionStart.newBuilder()
-                    .setProducerId(producerId)
-                    .setIntegration(integration)
-                    .setDeduplicationMode(deduplicationMode)
-                    .setDefaults(options.toDefaults()))
-            .build();
-    synchronized (lock) {
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(start);
-    }
-  }
-
-  // ---- ProducerBase ----
-
-  @Override
-  public long lastSentOffset() {
-    acquire();
-    try {
-      return lastSent;
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public long lastAcknowledgedOffset() {
-    acquire();
-    try {
-      synchronized (lock) {
-        return lastCommitted;
-      }
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture waitAcknowledged(long offset) {
-    acquire();
-    try {
-      return registerAckWaiter(offset);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public long flush() {
-    acquire();
-    try {
-      return awaitFlush(registerAckWaiter(lastSent));
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture flushAsync() {
-    acquire();
-    try {
-      return registerAckWaiter(lastSent);
-    } finally {
-      release();
-    }
-  }
-
-  /**
-   * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark
-   * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are
-   * reentrant, so this is safe to call while already holding {@code lock}).
-   */
-  private CompletableFuture registerAckWaiter(long offset) {
-    synchronized (lock) {
-      if (closed) {
-        return CompletableFuture.failedFuture(
-            Objects.requireNonNull(failure, "closed producer has no failure"));
-      }
-      if (offset <= lastCommitted) {
-        return CompletableFuture.completedFuture(lastCommitted);
-      }
-      CompletableFuture f = new CompletableFuture<>();
-      ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(f);
-      return f;
-    }
-  }
-
-  @Override
-  public void close() {
-    acquire();
-    try {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN, "producer closed"),
-          true);
-    } finally {
-      release();
-    }
-  }
-
-  // ---- send path, shared by the subclasses (caller holds the guard) ----
-
-  /** Admit a record, blocking up to the configured maximum while the producer is backpressured. */
-  final CompletableFuture doSend(long offset, InvocationImpl invocation)
-      throws ProducerBufferExhaustedException {
-    PreparedSend prepared = prepare(offset, invocation);
-    List> ready;
-    CompletableFuture acknowledgement;
-    boolean directWrite;
-    long waitStarted = System.nanoTime();
-    synchronized (lock) {
-      ensureOpenLocked();
-      while (!canAdmitLocked(prepared.bufferSize())) {
-        if (maxBlockNanos == 0) {
-          throw admissionTimeout();
-        }
-        long remaining = maxBlockNanos - (System.nanoTime() - waitStarted);
-        if (remaining <= 0) {
-          throw admissionTimeout();
-        }
-        try {
-          long millis = remaining / 1_000_000;
-          int nanos = (int) (remaining % 1_000_000);
-          lock.wait(millis, nanos);
-        } catch (InterruptedException e) {
-          Thread.currentThread().interrupt();
-          throw new ProducerBufferExhaustedException(
-              "interrupted while waiting for producer admission", e);
-        }
-        ensureOpenLocked();
-      }
-      acknowledgement = acceptLocked(prepared);
-      directWrite = bufferMemory == 0;
-      ready = directWrite ? List.of() : drainLocked();
-    }
-    if (directWrite) {
-      writeDirect(prepared.request());
-      drainAndWake();
-    } else {
-      completeReady(ready);
-    }
-    return acknowledgement;
-  }
-
-  /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
-  final SendAttempt doTrySend(long offset, InvocationImpl invocation) {
-    PreparedSend prepared = prepare(offset, invocation);
-    List> ready;
-    SendAttempt result;
-    boolean directWrite = false;
-    synchronized (lock) {
-      ensureOpenLocked();
-      if (canAdmitLocked(prepared.bufferSize())) {
-        result = new SendAttempt.Accepted(acceptLocked(prepared));
-        directWrite = bufferMemory == 0;
-        ready = directWrite ? List.of() : drainLocked();
-      } else {
-        CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
-        AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future);
-        admissionWaiters.add(waiter);
-        future.whenComplete(
-            (ignored, failure) -> {
-              if (future.isCancelled()) {
-                synchronized (lock) {
-                  admissionWaiters.remove(waiter);
-                }
-              }
-            });
-        result = new SendAttempt.Backpressured(future);
-        ready = List.of();
-      }
-    }
-    if (directWrite) {
-      writeDirect(prepared.request());
-      drainAndWake();
-    } else {
-      completeReady(ready);
-    }
-    return result;
-  }
-
-  // ---- internals (all `*Locked` methods require `lock`) ----
-
-  private PreparedSend prepare(long offset, InvocationImpl invocation) {
-    IngestionRequest request =
-        IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build();
-    long bufferSize = request.getSerializedSize();
-    if (bufferMemory > 0 && bufferSize > bufferMemory) {
-      throw new IllegalArgumentException(
-          "serialized invocation requires "
-              + bufferSize
-              + " bytes, exceeding bufferMemory "
-              + bufferMemory);
-    }
-    return new PreparedSend(
-        offset, request, request.getInvocation().getSerializedSize(), bufferSize);
-  }
-
-  private boolean canAdmitLocked(long requiredBytes) {
-    if (bufferMemory == 0) {
-      ClientCallStreamObserver observer =
-          Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-      return budget > 0 && observer.isReady();
-    }
-    return requiredBytes <= bufferMemory - bufferedBytes;
-  }
-
-  private CompletableFuture acceptLocked(PreparedSend prepared) {
-    lastSent = prepared.offset();
-    CompletableFuture committed = new CompletableFuture<>();
-    ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed);
-    if (bufferMemory == 0) {
-      budget -= prepared.windowDebit();
-    } else {
-      bufferedSends.addLast(
-          new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize()));
-      bufferedBytes += prepared.bufferSize();
-    }
-    return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
-  }
-
-  /**
-   * Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused.
-   */
-  private void writeDirect(IngestionRequest request) {
-    try {
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized")
-          .onNext(request);
-    } catch (RuntimeException e) {
-      IntegrationClientException cause =
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN,
-              "failed to write invocation to the ingestion stream",
-              e);
-      terminate(cause, false);
-      throw cause;
-    }
-  }
-
-  /** Write as many queued records as transport and protocol flow control currently permit. */
-  private List> drainLocked() {
-    if (bufferMemory == 0) {
-      if (!canAdmitLocked(0)) {
-        return List.of();
-      }
-      lock.notifyAll();
-      List> ready = new ArrayList<>(admissionWaiters.size());
-      for (AdmissionWaiter waiter : admissionWaiters) {
-        ready.add(waiter.future());
-      }
-      admissionWaiters.clear();
-      return ready;
-    }
-
-    boolean freedCapacity = false;
-    ClientCallStreamObserver observer =
-        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-    while (!closed && budget > 0 && observer.isReady() && !bufferedSends.isEmpty()) {
-      BufferedSend send = bufferedSends.removeFirst();
-      bufferedBytes -= send.bufferSize();
-      budget -= send.windowDebit();
-      freedCapacity = true;
-      observer.onNext(send.request());
-    }
-
-    if (!freedCapacity) {
-      return List.of();
-    }
-
-    lock.notifyAll();
-    long available = bufferMemory - bufferedBytes;
-    List> ready = new ArrayList<>();
-    for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) {
-      AdmissionWaiter waiter = it.next();
-      if (waiter.requiredBytes() <= available) {
-        ready.add(waiter.future());
-        it.remove();
-      }
-    }
-    return ready;
-  }
-
-  private void drainAndWake() {
-    List> ready;
-    synchronized (lock) {
-      ready = drainLocked();
-    }
-    completeReady(ready);
-  }
-
-  private static void completeReady(List> ready) {
-    for (CompletableFuture<@Nullable Void> future : ready) {
-      future.complete(null);
-    }
-  }
-
-  private void ensureOpenLocked() {
-    if (closed) {
-      throw new IllegalStateException("producer is closed", failure);
-    }
-  }
-
-  private ProducerBufferExhaustedException admissionTimeout() {
-    String condition =
-        bufferMemory == 0 ? "producer remained backpressured" : "producer buffer remained full";
-    return new ProducerBufferExhaustedException(condition + " for " + maxBlockTime);
-  }
-
-  private static long awaitFlush(CompletableFuture flush) {
-    try {
-      return flush.get();
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      throw new IntegrationClientException(
-          IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e);
-    } catch (ExecutionException e) {
-      @Nullable Throwable cause = e.getCause();
-      if (cause instanceof RuntimeException runtimeException) {
-        throw runtimeException;
-      }
-      if (cause instanceof Error error) {
-        throw error;
-      }
-      if (cause == null) {
-        throw new IntegrationClientException(
-            IntegrationClientException.Kind.UNKNOWN, "producer flush failed");
-      }
-      throw new IntegrationClientException(
-          IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause);
-    }
-  }
-
-  private static long toNanosSaturated(Duration duration) {
-    try {
-      return duration.toNanos();
-    } catch (ArithmeticException ignored) {
-      return Long.MAX_VALUE;
-    }
-  }
-
-  private void onResponse(IngestionResponse resp) {
-    @Nullable List> acksToComplete = null;
-    long watermark = -1;
-    boolean drain = false;
-    @Nullable IntegrationClientException err = null;
-    synchronized (lock) {
-      if (closed) {
-        return;
-      }
-      if (resp.hasLastCommitted() && resp.getLastCommitted() > lastCommitted) {
-        lastCommitted = resp.getLastCommitted();
-        watermark = lastCommitted;
-        if (!ackWaiters.isEmpty()) {
-          acksToComplete = new ArrayList<>();
-          Map>> head = ackWaiters.headMap(watermark, true);
-          for (List> waiters : head.values()) {
-            acksToComplete.addAll(waiters);
-          }
-          head.clear();
-        }
-      }
-      if (resp.hasWindowUpdate()) {
-        // increment_bytes is a uint32; read it as unsigned.
-        budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes());
-        drain = true;
-      } else if (resp.hasError()) {
-        err = mapError(resp.getError());
-      }
-    }
-    if (err != null) {
-      terminate(err, false);
-    } else if (drain) {
-      drainAndWake();
-    }
-    if (acksToComplete != null) {
-      for (CompletableFuture f : acksToComplete) {
-        f.complete(watermark);
-      }
-    }
-  }
-
-  /** Mark the producer terminally closed and fail every pending future with {@code cause}. */
-  private void terminate(IntegrationClientException cause, boolean halfClose) {
-    List> capacity;
-    List> acks = new ArrayList<>();
-    synchronized (lock) {
-      if (closed) {
-        return;
-      }
-      closed = true;
-      failure = cause;
-      capacity = new ArrayList<>(admissionWaiters.size());
-      for (AdmissionWaiter waiter : admissionWaiters) {
-        capacity.add(waiter.future());
-      }
-      admissionWaiters.clear();
-      bufferedSends.clear();
-      bufferedBytes = 0;
-      lock.notifyAll();
-      for (List> waiters : ackWaiters.values()) {
-        acks.addAll(waiters);
-      }
-      ackWaiters.clear();
-    }
-    if (halfClose) {
-      @Nullable ClientCallStreamObserver obs = callObserver;
-      if (obs != null) {
-        try {
-          obs.onCompleted();
-        } catch (RuntimeException ignored) {
-          // Already torn down transport-side; nothing to half-close.
-        }
-      }
-    }
-    for (CompletableFuture<@Nullable Void> f : capacity) {
-      f.completeExceptionally(cause);
-    }
-    for (CompletableFuture f : acks) {
-      f.completeExceptionally(cause);
-    }
-  }
-
-  private static IntegrationClientException mapError(dev.restate.ingestion.v1.Error error) {
-    String detail =
-        error.hasInvocationOffset()
-            ? "[offset=" + error.getInvocationOffset() + "] " + error.getMessage()
-            : error.getMessage();
-    return new IntegrationClientException(mapKind(error.getKind()), detail);
-  }
-
-  private static IntegrationClientException.Kind mapKind(ErrorKind kind) {
-    switch (kind) {
-      case ERROR_KIND_SHUTTING_DOWN:
-        return IntegrationClientException.Kind.SHUTTING_DOWN;
-      case ERROR_KIND_GO_AWAY:
-        return IntegrationClientException.Kind.GO_AWAY;
-      case ERROR_KIND_NOT_FOUND:
-        return IntegrationClientException.Kind.NOT_FOUND;
-      case ERROR_KIND_BAD_REQUEST:
-        return IntegrationClientException.Kind.BAD_REQUEST;
-      default:
-        return IntegrationClientException.Kind.UNKNOWN;
-    }
-  }
-
-  // ---- fail-fast guard ----
-
-  final void acquire() {
-    Thread current = Thread.currentThread();
-    if (owner.get() == current) {
-      reentrancy++;
-      return;
-    }
-    if (!owner.compareAndSet(null, current)) {
-      throw new ConcurrentModificationException("Producer is not safe for multi-threaded access");
-    }
-    reentrancy = 1;
-  }
-
-  final void release() {
-    if (--reentrancy == 0) {
-      owner.set(null);
-    }
-  }
-
-  private final class ResponseObserver
-      implements ClientResponseObserver {
-    @Override
-    public void beforeStart(ClientCallStreamObserver requestStream) {
-      callObserver = requestStream;
-      requestStream.setOnReadyHandler(AbstractProducer.this::drainAndWake);
-    }
-
-    @Override
-    public void onNext(IngestionResponse value) {
-      onResponse(value);
-    }
-
-    @Override
-    public void onError(Throwable t) {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN,
-              "ingestion stream failed: " + t.getMessage(),
-              t),
-          false);
-    }
-
-    @Override
-    public void onCompleted() {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN, "ingestion stream closed by server"),
-          false);
-    }
-  }
-
-  private record PreparedSend(
-      long offset, IngestionRequest request, long windowDebit, long bufferSize) {}
-
-  private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {}
-
-  private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
-
-  private record SendResultImpl(long offset) implements SendResult {}
-}
diff --git a/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java b/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java
deleted file mode 100644
index 6ab9cafe..00000000
--- a/integration-client/src/main/java/dev/restate/integration/AuthInterceptor.java
+++ /dev/null
@@ -1,43 +0,0 @@
-// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
-//
-// This file is part of the Restate Java SDK,
-// which is released under the MIT license.
-//
-// You can find a copy of the license in file LICENSE in the root
-// directory of this repository or package, or at
-// https://github.com/restatedev/sdk-java/blob/main/LICENSE
-package dev.restate.integration;
-
-import io.grpc.CallOptions;
-import io.grpc.Channel;
-import io.grpc.ClientCall;
-import io.grpc.ClientInterceptor;
-import io.grpc.ForwardingClientCall;
-import io.grpc.Metadata;
-import io.grpc.MethodDescriptor;
-
-/** Adds an {@code Authorization: Bearer } header to every call. */
-final class AuthInterceptor implements ClientInterceptor {
-
-  private static final Metadata.Key AUTHORIZATION =
-      Metadata.Key.of("Authorization", Metadata.ASCII_STRING_MARSHALLER);
-
-  private final String bearer;
-
-  AuthInterceptor(String token) {
-    this.bearer = "Bearer " + token;
-  }
-
-  @Override
-  public  ClientCall interceptCall(
-      MethodDescriptor method, CallOptions callOptions, Channel next) {
-    return new ForwardingClientCall.SimpleForwardingClientCall<>(
-        next.newCall(method, callOptions)) {
-      @Override
-      public void start(Listener responseListener, Metadata headers) {
-        headers.put(AUTHORIZATION, bearer);
-        super.start(responseListener, headers);
-      }
-    };
-  }
-}
diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java
deleted file mode 100644
index 03b90f1c..00000000
--- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducerImpl.java
+++ /dev/null
@@ -1,55 +0,0 @@
-// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
-//
-// This file is part of the Restate Java SDK,
-// which is released under the MIT license.
-//
-// You can find a copy of the license in file LICENSE in the root
-// directory of this repository or package, or at
-// https://github.com/restatedev/sdk-java/blob/main/LICENSE
-package dev.restate.integration;
-
-import dev.restate.ingestion.v1.DeduplicationMode;
-import dev.restate.ingestion.v1.IngestionSvcGrpc;
-import java.util.concurrent.CompletableFuture;
-
-/** Exactly-once {@link ExactlyOnceProducer}: caller-supplied, strictly-increasing offsets. */
-final class ExactlyOnceProducerImpl extends AbstractProducer implements ExactlyOnceProducer {
-
-  ExactlyOnceProducerImpl(
-      IngestionSvcGrpc.IngestionSvcStub stub,
-      String producerId,
-      ProducerOptions options,
-      String integration) {
-    super(stub, producerId, DeduplicationMode.OFFSET_BASED, options, integration);
-  }
-
-  @Override
-  public SendAttempt trySend(long offset, Invocation invocation) {
-    acquire();
-    try {
-      checkOffset(offset);
-      return doTrySend(offset, (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture send(long offset, Invocation invocation)
-      throws ProducerBufferExhaustedException {
-    acquire();
-    try {
-      checkOffset(offset);
-      return doSend(offset, (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  private void checkOffset(long offset) {
-    if (offset <= lastSent) {
-      throw new IllegalArgumentException(
-          "offset must be strictly increasing; last sent " + lastSent + ", got " + offset);
-    }
-  }
-}
diff --git a/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java
index edc66734..401fc4f6 100644
--- a/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java
+++ b/integration-client/src/main/java/dev/restate/integration/IngressEndpoint.java
@@ -9,7 +9,6 @@
 package dev.restate.integration;
 
 import java.net.URI;
-import org.jspecify.annotations.Nullable;
 
 /** Where to reach the Restate ingestion gRPC endpoint, parsed from an http(s) URL. */
 final class IngressEndpoint {
@@ -39,7 +38,7 @@ static IngressEndpoint parse(String raw) {
       throw new IllegalArgumentException("ingress url is not a valid URL: '" + raw + "'", e);
     }
     boolean tls;
-    @Nullable String scheme = uri.getScheme() == null ? null : uri.getScheme().toLowerCase();
+    String scheme = uri.getScheme() == null ? null : uri.getScheme().toLowerCase();
     if ("https".equals(scheme)) {
       tls = true;
     } else if ("http".equals(scheme)) {
@@ -52,7 +51,7 @@ static IngressEndpoint parse(String raw) {
               + raw
               + "'");
     }
-    @Nullable String host = uri.getHost();
+    String host = uri.getHost();
     if (host == null) {
       throw new IllegalArgumentException("ingress url has no host: '" + raw + "'");
     }
diff --git a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
index ef96a717..a2f69e9c 100644
--- a/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/IntegrationClientImpl.java
@@ -12,6 +12,8 @@
 import io.grpc.Channel;
 import io.grpc.ManagedChannel;
 import io.grpc.ManagedChannelBuilder;
+import io.grpc.Metadata;
+import io.grpc.stub.MetadataUtils;
 import java.util.Objects;
 import java.util.concurrent.TimeUnit;
 import org.jspecify.annotations.Nullable;
@@ -19,6 +21,9 @@
 /** {@link IntegrationClient} backed by a single gRPC {@link Channel} shared by producers. */
 final class IntegrationClientImpl implements IntegrationClient {
 
+  private static final Metadata.Key AUTHORIZATION =
+      Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
+
   private final @Nullable ManagedChannel ownedChannel;
   private final IngestionSvcGrpc.IngestionSvcStub stub;
   private final String integration;
@@ -57,7 +62,9 @@ private static IntegrationClient create(
       @Nullable ManagedChannel ownedChannel) {
     IngestionSvcGrpc.IngestionSvcStub stub = IngestionSvcGrpc.newStub(channel);
     if (authToken != null && !authToken.isBlank()) {
-      stub = stub.withInterceptors(new AuthInterceptor(authToken));
+      Metadata headers = new Metadata();
+      headers.put(AUTHORIZATION, "Bearer " + authToken);
+      stub = stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(headers));
     }
     return new IntegrationClientImpl(ownedChannel, stub, integration);
   }
@@ -73,7 +80,7 @@ public ExactlyOnceProducer newExactlyOnceProducer(String producerId, ProducerOpt
       throw new IllegalArgumentException(
           "producerId must be non-empty for an exactly-once producer");
     }
-    return new ExactlyOnceProducerImpl(
+    return new ProducerImpl(
         stub, producerId, Objects.requireNonNull(options, "options"), integration);
   }
 
diff --git a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java
index 9d3294ab..9a52ab59 100644
--- a/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/InvocationMetadataImpl.java
@@ -23,30 +23,6 @@ sealed class InvocationMetadataImpl implements InvocationMetadata permits Invoca
 
   final IngestionInvocation.Builder builder = IngestionInvocation.newBuilder();
 
-  static InvocationMetadataImpl fromDefaults(IngestionDefaults defaults) {
-    InvocationMetadataImpl metadata = new InvocationMetadataImpl();
-    if (defaults.hasService()) {
-      metadata.builder.setService(defaults.getService());
-    }
-    if (defaults.hasHandler()) {
-      metadata.builder.setHandler(defaults.getHandler());
-    }
-    if (defaults.hasKey()) {
-      metadata.builder.setKey(defaults.getKey());
-    }
-    if (defaults.hasScope()) {
-      metadata.builder.setScope(defaults.getScope());
-    }
-    if (defaults.hasLimitKey()) {
-      metadata.builder.setLimitKey(defaults.getLimitKey());
-    }
-    if (defaults.hasIdempotencyKey()) {
-      metadata.builder.setIdempotencyKey(defaults.getIdempotencyKey());
-    }
-    metadata.builder.putAllAdditionalHeaders(defaults.getHeadersMap());
-    return metadata;
-  }
-
   @Override
   public InvocationMetadata setServiceName(@Nullable String serviceName) {
     if (serviceName == null) {
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
index 642be6bd..0f05bd7d 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
@@ -9,22 +9,126 @@
 package dev.restate.integration;
 
 import dev.restate.ingestion.v1.DeduplicationMode;
+import dev.restate.ingestion.v1.ErrorKind;
+import dev.restate.ingestion.v1.IngestionRequest;
+import dev.restate.ingestion.v1.IngestionResponse;
+import dev.restate.ingestion.v1.IngestionStart;
 import dev.restate.ingestion.v1.IngestionSvcGrpc;
+import io.grpc.stub.ClientCallStreamObserver;
+import io.grpc.stub.ClientResponseObserver;
+import java.time.Duration;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.ConcurrentModificationException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.atomic.AtomicReference;
+import org.jspecify.annotations.Nullable;
 
-/** At-least-once {@link Producer}: dedup disabled, client-assigned monotonic offsets. */
-final class ProducerImpl extends AbstractProducer implements Producer {
+/**
+ * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and
+ * the module docs for the concurrency contract.
+ *
+ * 
When buffering is enabled, accepted records wait in a byte-bounded queue until Restate's
+ * send-window has credit ({@code budget}) and the transport is writable ({@code
+ * callObserver.isReady()}). With buffering disabled, records are accepted only when they can be
+ * handed directly to gRPC. Once handed off, only their acknowledgement futures remain until the
+ * commit watermark passes their offsets.
+ *
+ * 
Two-tier concurrency:
+ *
+ * 
    
+ *   
  • A fail-fast, KafkaConsumer-style guard ({@link #acquire()}/{@link #release()}) rejects
+ *       concurrent use from multiple threads; sequential hand-off between threads is fine.
+ *   
  • A single monitor ({@link #lock}) guards the small set of fields genuinely shared between
+ *       the caller thread and gRPC's callback threads. Futures are always completed outside the
+ *       monitor.
+ * 

+ */
+final class ProducerImpl implements Producer, ExactlyOnceProducer {
+
+  private final Object lock = new Object();
+
+  // Set once, synchronously, in beforeStart() before the constructor sends the Start frame.
+  private volatile @Nullable ClientCallStreamObserver callObserver;
+
+  // ---- fail-fast single-thread guard ----
+  private final AtomicReference<@Nullable Thread> owner = new AtomicReference<>();
+  private int reentrancy;
+
+  // ---- state guarded by `lock` ----
+  private long budget = 0; // remaining Restate send window, in bytes; may go one message negative
+  private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet
+  private final ArrayDeque bufferedSends = new ArrayDeque<>();
+  private long bufferedBytes = 0;
+  private final List admissionWaiters = new ArrayList<>();
+  private final TreeMap>> ackWaiters = new TreeMap<>();
+  private boolean closed = false;
+  private @Nullable IntegrationClientException failure;
+
+  private final long bufferMemory;
+  private final Duration maxBlockTime;
+  private final long maxBlockNanos;
+  private final boolean exactlyOnce;
+
+  // Written only by the (guarded) caller thread; never touched by gRPC callbacks.
+  private long lastSent = -1;
 
   ProducerImpl(
       IngestionSvcGrpc.IngestionSvcStub stub, ProducerOptions options, String integration) {
-    super(stub, "", DeduplicationMode.DISABLED, options, integration);
+    this(stub, "", DeduplicationMode.DISABLED, false, options, integration);
+  }
+
+  ProducerImpl(
+      IngestionSvcGrpc.IngestionSvcStub stub,
+      String producerId,
+      ProducerOptions options,
+      String integration) {
+    this(stub, producerId, DeduplicationMode.OFFSET_BASED, true, options, integration);
+  }
+
+  private ProducerImpl(
+      IngestionSvcGrpc.IngestionSvcStub stub,
+      String producerId,
+      DeduplicationMode deduplicationMode,
+      boolean exactlyOnce,
+      ProducerOptions options,
+      String integration) {
+    this.bufferMemory = options.bufferMemory();
+    this.maxBlockTime = options.maxBlockTime();
+    this.maxBlockNanos = toNanosSaturated(maxBlockTime);
+    this.exactlyOnce = exactlyOnce;
+    // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready
+    // handler.
+    stub.ingest(new ResponseObserver());
+    // Mandatory Start handshake: the first frame on the stream (not flow-controlled).
+    IngestionRequest start =
+        IngestionRequest.newBuilder()
+            .setStart(
+                IngestionStart.newBuilder()
+                    .setProducerId(producerId)
+                    .setIntegration(integration)
+                    .setDeduplicationMode(deduplicationMode)
+                    .setDefaults(options.toDefaults()))
+            .build();
+    synchronized (lock) {
+      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(start);
+    }
   }
 
+  // ---- Producer / ExactlyOnceProducer ----
+
   @Override
   public CompletableFuture send(Invocation invocation)
       throws ProducerBufferExhaustedException {
     acquire();
     try {
+      checkMode(false);
       return doSend(lastSent + 1, (InvocationImpl) invocation);
     } finally {
       release();
@@ -35,9 +139,544 @@ public CompletableFuture send(Invocation invocation)
   public SendAttempt trySend(Invocation invocation) {
     acquire();
     try {
+      checkMode(false);
       return doTrySend(lastSent + 1, (InvocationImpl) invocation);
     } finally {
       release();
     }
   }
+
+  @Override
+  public CompletableFuture send(long offset, Invocation invocation)
+      throws ProducerBufferExhaustedException {
+    acquire();
+    try {
+      checkMode(true);
+      checkOffset(offset);
+      return doSend(offset, (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public SendAttempt trySend(long offset, Invocation invocation) {
+    acquire();
+    try {
+      checkMode(true);
+      checkOffset(offset);
+      return doTrySend(offset, (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  private void checkOffset(long offset) {
+    if (offset <= lastSent) {
+      throw new IllegalArgumentException(
+          "offset must be strictly increasing; last sent " + lastSent + ", got " + offset);
+    }
+  }
+
+  private void checkMode(boolean exactlyOnceExpected) {
+    if (exactlyOnce != exactlyOnceExpected) {
+      throw new IllegalStateException(
+          exactlyOnce
+              ? "exactly-once producers require explicit offsets"
+              : "at-least-once producers assign offsets automatically");
+    }
+  }
+
+  // ---- ProducerBase ----
+
+  @Override
+  public long lastSentOffset() {
+    acquire();
+    try {
+      return lastSent;
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public long lastAcknowledgedOffset() {
+    acquire();
+    try {
+      synchronized (lock) {
+        return lastCommitted;
+      }
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture waitAcknowledged(long offset) {
+    acquire();
+    try {
+      return registerAckWaiter(offset);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public long flush() {
+    acquire();
+    try {
+      return awaitFlush(registerAckWaiter(lastSent));
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture flushAsync() {
+    acquire();
+    try {
+      return registerAckWaiter(lastSent);
+    } finally {
+      release();
+    }
+  }
+
+  /**
+   * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark
+   * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are
+   * reentrant, so this is safe to call while already holding {@code lock}).
+   */
+  private CompletableFuture registerAckWaiter(long offset) {
+    synchronized (lock) {
+      if (closed) {
+        return CompletableFuture.failedFuture(
+            Objects.requireNonNull(failure, "closed producer has no failure"));
+      }
+      if (offset <= lastCommitted) {
+        return CompletableFuture.completedFuture(lastCommitted);
+      }
+      CompletableFuture f = new CompletableFuture<>();
+      ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(f);
+      return f;
+    }
+  }
+
+  @Override
+  public void close() {
+    acquire();
+    try {
+      terminate(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN, "producer closed"),
+          true);
+    } finally {
+      release();
+    }
+  }
+
+  // ---- send path, shared by both producer modes (caller holds the guard) ----
+
+  /** Admit a record, blocking up to the configured maximum while the producer is backpressured. */
+  private CompletableFuture doSend(long offset, InvocationImpl invocation)
+      throws ProducerBufferExhaustedException {
+    PreparedSend prepared = prepare(offset, invocation);
+    List> ready;
+    CompletableFuture acknowledgement;
+    boolean directWrite;
+    long waitStarted = System.nanoTime();
+    synchronized (lock) {
+      ensureOpenLocked();
+      while (!canAdmitLocked(prepared.bufferSize())) {
+        if (maxBlockNanos == 0) {
+          throw admissionTimeout();
+        }
+        long remaining = maxBlockNanos - (System.nanoTime() - waitStarted);
+        if (remaining <= 0) {
+          throw admissionTimeout();
+        }
+        try {
+          long millis = remaining / 1_000_000;
+          int nanos = (int) (remaining % 1_000_000);
+          lock.wait(millis, nanos);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          throw new ProducerBufferExhaustedException(
+              "interrupted while waiting for producer admission", e);
+        }
+        ensureOpenLocked();
+      }
+      acknowledgement = acceptLocked(prepared);
+      directWrite = bufferMemory == 0;
+      ready = directWrite ? List.of() : drainLocked();
+    }
+    if (directWrite) {
+      writeDirect(prepared.request());
+      drainAndWake();
+    } else {
+      completeReady(ready);
+    }
+    return acknowledgement;
+  }
+
+  /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
+  private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
+    PreparedSend prepared = prepare(offset, invocation);
+    List> ready;
+    SendAttempt result;
+    boolean directWrite = false;
+    synchronized (lock) {
+      ensureOpenLocked();
+      if (canAdmitLocked(prepared.bufferSize())) {
+        result = new SendAttempt.Accepted(acceptLocked(prepared));
+        directWrite = bufferMemory == 0;
+        ready = directWrite ? List.of() : drainLocked();
+      } else {
+        CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
+        AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future);
+        admissionWaiters.add(waiter);
+        future.whenComplete(
+            (ignored, failure) -> {
+              if (future.isCancelled()) {
+                synchronized (lock) {
+                  admissionWaiters.remove(waiter);
+                }
+              }
+            });
+        result = new SendAttempt.Backpressured(future);
+        ready = List.of();
+      }
+    }
+    if (directWrite) {
+      writeDirect(prepared.request());
+      drainAndWake();
+    } else {
+      completeReady(ready);
+    }
+    return result;
+  }
+
+  // ---- internals (all `*Locked` methods require `lock`) ----
+
+  private PreparedSend prepare(long offset, InvocationImpl invocation) {
+    IngestionRequest request =
+        IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build();
+    long bufferSize = request.getSerializedSize();
+    if (bufferMemory > 0 && bufferSize > bufferMemory) {
+      throw new IllegalArgumentException(
+          "serialized invocation requires "
+              + bufferSize
+              + " bytes, exceeding bufferMemory "
+              + bufferMemory);
+    }
+    return new PreparedSend(
+        offset, request, request.getInvocation().getSerializedSize(), bufferSize);
+  }
+
+  private boolean canAdmitLocked(long requiredBytes) {
+    if (bufferMemory == 0) {
+      ClientCallStreamObserver observer =
+          Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+      return budget > 0 && observer.isReady();
+    }
+    return requiredBytes <= bufferMemory - bufferedBytes;
+  }
+
+  private CompletableFuture acceptLocked(PreparedSend prepared) {
+    lastSent = prepared.offset();
+    CompletableFuture committed = new CompletableFuture<>();
+    ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed);
+    if (bufferMemory == 0) {
+      budget -= prepared.windowDebit();
+    } else {
+      bufferedSends.addLast(
+          new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize()));
+      bufferedBytes += prepared.bufferSize();
+    }
+    return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
+  }
+
+  /**
+   * Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused.
+   */
+  private void writeDirect(IngestionRequest request) {
+    try {
+      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized")
+          .onNext(request);
+    } catch (RuntimeException e) {
+      IntegrationClientException cause =
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN,
+              "failed to write invocation to the ingestion stream",
+              e);
+      terminate(cause, false);
+      throw cause;
+    }
+  }
+
+  /** Write as many queued records as transport and protocol flow control currently permit. */
+  private List> drainLocked() {
+    if (bufferMemory == 0) {
+      if (!canAdmitLocked(0)) {
+        return List.of();
+      }
+      lock.notifyAll();
+      List> ready = new ArrayList<>(admissionWaiters.size());
+      for (AdmissionWaiter waiter : admissionWaiters) {
+        ready.add(waiter.future());
+      }
+      admissionWaiters.clear();
+      return ready;
+    }
+
+    boolean freedCapacity = false;
+    ClientCallStreamObserver observer =
+        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+    while (!closed && budget > 0 && observer.isReady() && !bufferedSends.isEmpty()) {
+      BufferedSend send = bufferedSends.removeFirst();
+      bufferedBytes -= send.bufferSize();
+      budget -= send.windowDebit();
+      freedCapacity = true;
+      observer.onNext(send.request());
+    }
+
+    if (!freedCapacity) {
+      return List.of();
+    }
+
+    lock.notifyAll();
+    long available = bufferMemory - bufferedBytes;
+    List> ready = new ArrayList<>();
+    for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) {
+      AdmissionWaiter waiter = it.next();
+      if (waiter.requiredBytes() <= available) {
+        ready.add(waiter.future());
+        it.remove();
+      }
+    }
+    return ready;
+  }
+
+  private void drainAndWake() {
+    List> ready;
+    synchronized (lock) {
+      ready = drainLocked();
+    }
+    completeReady(ready);
+  }
+
+  private static void completeReady(List> ready) {
+    for (CompletableFuture<@Nullable Void> future : ready) {
+      future.complete(null);
+    }
+  }
+
+  private void ensureOpenLocked() {
+    if (closed) {
+      throw new IllegalStateException("producer is closed", failure);
+    }
+  }
+
+  private ProducerBufferExhaustedException admissionTimeout() {
+    String condition =
+        bufferMemory == 0 ? "producer remained backpressured" : "producer buffer remained full";
+    return new ProducerBufferExhaustedException(condition + " for " + maxBlockTime);
+  }
+
+  private static long awaitFlush(CompletableFuture flush) {
+    try {
+      return flush.get();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e);
+    } catch (ExecutionException e) {
+      @Nullable Throwable cause = e.getCause();
+      if (cause instanceof RuntimeException runtimeException) {
+        throw runtimeException;
+      }
+      if (cause instanceof Error error) {
+        throw error;
+      }
+      if (cause == null) {
+        throw new IntegrationClientException(
+            IntegrationClientException.Kind.UNKNOWN, "producer flush failed");
+      }
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause);
+    }
+  }
+
+  private static long toNanosSaturated(Duration duration) {
+    try {
+      return duration.toNanos();
+    } catch (ArithmeticException ignored) {
+      return Long.MAX_VALUE;
+    }
+  }
+
+  private void onResponse(IngestionResponse resp) {
+    @Nullable List> acksToComplete = null;
+    long watermark = -1;
+    boolean drain = false;
+    @Nullable IntegrationClientException err = null;
+    synchronized (lock) {
+      if (closed) {
+        return;
+      }
+      if (resp.hasLastCommitted() && resp.getLastCommitted() > lastCommitted) {
+        lastCommitted = resp.getLastCommitted();
+        watermark = lastCommitted;
+        if (!ackWaiters.isEmpty()) {
+          acksToComplete = new ArrayList<>();
+          Map>> head = ackWaiters.headMap(watermark, true);
+          for (List> waiters : head.values()) {
+            acksToComplete.addAll(waiters);
+          }
+          head.clear();
+        }
+      }
+      if (resp.hasWindowUpdate()) {
+        // increment_bytes is a uint32; read it as unsigned.
+        budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes());
+        drain = true;
+      } else if (resp.hasError()) {
+        err = mapError(resp.getError());
+      }
+    }
+    if (err != null) {
+      terminate(err, false);
+    } else if (drain) {
+      drainAndWake();
+    }
+    if (acksToComplete != null) {
+      for (CompletableFuture f : acksToComplete) {
+        f.complete(watermark);
+      }
+    }
+  }
+
+  /** Mark the producer terminally closed and fail every pending future with {@code cause}. */
+  private void terminate(IntegrationClientException cause, boolean halfClose) {
+    List> capacity;
+    List> acks = new ArrayList<>();
+    synchronized (lock) {
+      if (closed) {
+        return;
+      }
+      closed = true;
+      failure = cause;
+      capacity = new ArrayList<>(admissionWaiters.size());
+      for (AdmissionWaiter waiter : admissionWaiters) {
+        capacity.add(waiter.future());
+      }
+      admissionWaiters.clear();
+      bufferedSends.clear();
+      bufferedBytes = 0;
+      lock.notifyAll();
+      for (List> waiters : ackWaiters.values()) {
+        acks.addAll(waiters);
+      }
+      ackWaiters.clear();
+    }
+    if (halfClose) {
+      @Nullable ClientCallStreamObserver obs = callObserver;
+      if (obs != null) {
+        try {
+          obs.onCompleted();
+        } catch (RuntimeException ignored) {
+          // Already torn down transport-side; nothing to half-close.
+        }
+      }
+    }
+    for (CompletableFuture<@Nullable Void> f : capacity) {
+      f.completeExceptionally(cause);
+    }
+    for (CompletableFuture f : acks) {
+      f.completeExceptionally(cause);
+    }
+  }
+
+  private static IntegrationClientException mapError(dev.restate.ingestion.v1.Error error) {
+    String detail =
+        error.hasInvocationOffset()
+            ? "[offset=" + error.getInvocationOffset() + "] " + error.getMessage()
+            : error.getMessage();
+    return new IntegrationClientException(mapKind(error.getKind()), detail);
+  }
+
+  private static IntegrationClientException.Kind mapKind(ErrorKind kind) {
+    switch (kind) {
+      case ERROR_KIND_SHUTTING_DOWN:
+        return IntegrationClientException.Kind.SHUTTING_DOWN;
+      case ERROR_KIND_GO_AWAY:
+        return IntegrationClientException.Kind.GO_AWAY;
+      case ERROR_KIND_NOT_FOUND:
+        return IntegrationClientException.Kind.NOT_FOUND;
+      case ERROR_KIND_BAD_REQUEST:
+        return IntegrationClientException.Kind.BAD_REQUEST;
+      default:
+        return IntegrationClientException.Kind.UNKNOWN;
+    }
+  }
+
+  // ---- fail-fast guard ----
+
+  private void acquire() {
+    Thread current = Thread.currentThread();
+    if (owner.get() == current) {
+      reentrancy++;
+      return;
+    }
+    if (!owner.compareAndSet(null, current)) {
+      throw new ConcurrentModificationException("Producer is not safe for multi-threaded access");
+    }
+    reentrancy = 1;
+  }
+
+  private void release() {
+    if (--reentrancy == 0) {
+      owner.set(null);
+    }
+  }
+
+  private final class ResponseObserver
+      implements ClientResponseObserver {
+    @Override
+    public void beforeStart(ClientCallStreamObserver requestStream) {
+      callObserver = requestStream;
+      requestStream.setOnReadyHandler(ProducerImpl.this::drainAndWake);
+    }
+
+    @Override
+    public void onNext(IngestionResponse value) {
+      onResponse(value);
+    }
+
+    @Override
+    public void onError(Throwable t) {
+      terminate(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN,
+              "ingestion stream failed: " + t.getMessage(),
+              t),
+          false);
+    }
+
+    @Override
+    public void onCompleted() {
+      terminate(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN, "ingestion stream closed by server"),
+          false);
+    }
+  }
+
+  private record PreparedSend(
+      long offset, IngestionRequest request, long windowDebit, long bufferSize) {}
+
+  private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {}
+
+  private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
+
+  private record SendResultImpl(long offset) implements SendResult {}
 }
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 3f56b33a..709fb4dd 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -23,8 +23,13 @@
 import io.grpc.ClientCall;
 import io.grpc.ForwardingClientCall;
 import io.grpc.ManagedChannel;
+import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
 import io.grpc.Server;
+import io.grpc.ServerCall;
+import io.grpc.ServerCallHandler;
+import io.grpc.ServerInterceptor;
+import io.grpc.ServerInterceptors;
 import io.grpc.inprocess.InProcessChannelBuilder;
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.stub.StreamObserver;
@@ -38,25 +43,48 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.ValueSource;
 
 /** Drives the producer client against an in-process fake {@code IngestionSvc}. */
 class IntegrationClientTest {
 
   private static final String INTEGRATION = "test-integration/1.0";
+  private static final Metadata.Key AUTHORIZATION =
+      Metadata.Key.of("authorization", Metadata.ASCII_STRING_MARSHALLER);
 
   private Server server;
   private ManagedChannel channel;
   private FakeIngestionService fake;
   private IntegrationClient client;
+  private final AtomicReference authorization = new AtomicReference<>();
 
   @BeforeEach
   void setUp() throws IOException {
     String name = InProcessServerBuilder.generateName();
     fake = new FakeIngestionService();
-    server = InProcessServerBuilder.forName(name).directExecutor().addService(fake).build().start();
+    server =
+        InProcessServerBuilder.forName(name)
+            .directExecutor()
+            .addService(
+                ServerInterceptors.intercept(
+                    fake,
+                    new ServerInterceptor() {
+                      @Override
+                      public  ServerCall.Listener interceptCall(
+                          ServerCall call,
+                          Metadata headers,
+                          ServerCallHandler next) {
+                        authorization.set(headers.get(AUTHORIZATION));
+                        return next.startCall(call, headers);
+                      }
+                    }))
+            .build()
+            .start();
     channel = InProcessChannelBuilder.forName(name).directExecutor().build();
     client = GrpcIntegrationClient.builder(channel).integration("test-integration", "1.0").build();
   }
@@ -82,6 +110,18 @@ void grpcBridgeDoesNotCloseCallerOwnedChannel() {
     client = null;
   }
 
+  @ParameterizedTest(name = "authToken={0}")
+  @ValueSource(strings = {"secret-token", "", " "})
+  void authTokenIsAttachedAsBearerMetadata(String authToken) throws Exception {
+    client.close();
+    client = GrpcIntegrationClient.builder(channel).authToken(authToken).build();
+
+    client.newProducer();
+    fake.take(); // Start
+
+    assertThat(authorization.get()).isEqualTo(authToken.isBlank() ? null : "Bearer " + authToken);
+  }
+
   @Test
   void producerSendsDisabledDedupHandshake() throws Exception {
     InvocationMetadata defaults = InvocationMetadata.create().setServiceName("Svc");
@@ -153,6 +193,22 @@ void exactlyOnceProducerRequiresProducerId() {
         .isInstanceOf(IllegalArgumentException.class);
   }
 
+  @ParameterizedTest(name = "exactlyOnce={0}")
+  @ValueSource(booleans = {false, true})
+  void producerRejectsMethodsFromTheOtherMode(boolean exactlyOnce) throws Exception {
+    Object producer =
+        exactlyOnce ? client.newExactlyOnceProducer("producer-1") : client.newProducer();
+    fake.take(); // Start
+
+    Runnable wrongSend =
+        exactlyOnce
+            ? () -> ((Producer) producer).send(newBody("a"))
+            : () -> ((ExactlyOnceProducer) producer).send(0L, newBody("a"));
+
+    assertThatThrownBy(wrongSend::run).isInstanceOf(IllegalStateException.class);
+    fake.assertNoRequest();
+  }
+
   @Test
   void producerAssignsMonotonicOffsetsAndFuturesCompleteOnCommit() throws Exception {
     Producer producer = client.newProducer();

From 90edfa8a183d68d31c324173cf0e12257f8991d3 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Mon, 24 Aug 2026 20:36:26 +0200
Subject: [PATCH 14/17] whatever changes

---
 .../integration/ExactlyOnceProducer.java      |   2 +
 .../dev/restate/integration/Producer.java     |   2 +
 .../dev/restate/integration/ProducerBase.java |   7 +
 .../dev/restate/integration/ProducerImpl.java | 439 ++++++++++++------
 .../integration/IntegrationClientTest.java    | 291 +++++++++++-
 5 files changed, 582 insertions(+), 159 deletions(-)

diff --git a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
index 47acfdcd..2b2aeeb5 100644
--- a/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
+++ b/integration-client/src/main/java/dev/restate/integration/ExactlyOnceProducer.java
@@ -94,6 +94,8 @@ public interface ExactlyOnceProducer extends ProducerBase {
    *     SendResult} carrying {@code offset}
    * @throws ProducerBufferExhaustedException if the producer cannot admit the invocation before the
    *     configured maximum blocking time elapses, or the thread is interrupted while waiting
+   * @throws IllegalStateException if a reentrant producer callback invokes this method when it
+   *     would block
    * @throws IllegalArgumentException if {@code offset} is not strictly greater than {@link
    *     #lastSentOffset()}, or buffering is enabled and the serialized invocation is larger than
    *     {@link ProducerOptions#bufferMemory()}
diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java
index 55eeebf0..3e220c82 100644
--- a/integration-client/src/main/java/dev/restate/integration/Producer.java
+++ b/integration-client/src/main/java/dev/restate/integration/Producer.java
@@ -79,6 +79,8 @@ public interface Producer extends ProducerBase {
    * @return a future completing, once the invocation is durably committed by Restate.
    * @throws ProducerBufferExhaustedException if the producer cannot admit the invocation before the
    *     configured maximum blocking time elapses, or the thread is interrupted while waiting
+   * @throws IllegalStateException if a reentrant producer callback invokes this method when it
+   *     would block
    * @throws IllegalArgumentException if buffering is enabled and the serialized invocation is
    *     larger than {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
index 03db1994..3457974d 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerBase.java
@@ -13,6 +13,11 @@
 /**
  * Common producer offsets, acknowledgement, flushing, and lifecycle operations.
  *
+ * 
Producer futures can complete inline on a transport callback. A synchronous continuation must
+ * not invoke an operation that would block, such as {@link #flush()}; offload the continuation to
+ * an executor instead. A reentrant call that would block is rejected with {@link
+ * IllegalStateException} rather than deadlocking the transport callback lane.
+ *
  * @see Producer
  * @see ExactlyOnceProducer
  */
@@ -58,6 +63,8 @@ public interface ProducerBase extends AutoCloseable {
    * @return the highest durably committed offset
    * @throws IntegrationClientException if the producer fails before all invocations are
    *     acknowledged
+   * @throws IllegalStateException if a reentrant producer callback invokes this method when it
+   *     would block
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
index 0f05bd7d..25eaddfb 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
@@ -27,7 +27,7 @@
 import java.util.TreeMap;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
-import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import org.jspecify.annotations.Nullable;
 
 /**
@@ -52,24 +52,27 @@
  */
 final class ProducerImpl implements Producer, ExactlyOnceProducer {
 
+  // Blocking from an inline transport/future callback can deadlock gRPC's serialized callback lane.
+  private static final ThreadLocal INLINE_CALLBACK = new ThreadLocal<>();
+
   private final Object lock = new Object();
 
   // Set once, synchronously, in beforeStart() before the constructor sends the Start frame.
   private volatile @Nullable ClientCallStreamObserver callObserver;
 
-  // ---- fail-fast single-thread guard ----
-  private final AtomicReference<@Nullable Thread> owner = new AtomicReference<>();
-  private int reentrancy;
+  // Reentrant because completing a future can synchronously call back into this producer.
+  private final ReentrantLock usageGuard = new ReentrantLock();
 
   // ---- state guarded by `lock` ----
   private long budget = 0; // remaining Restate send window, in bytes; may go one message negative
   private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet
-  private final ArrayDeque bufferedSends = new ArrayDeque<>();
+  private final ArrayDeque pendingWrites = new ArrayDeque<>();
   private long bufferedBytes = 0;
   private final List admissionWaiters = new ArrayList<>();
-  private final TreeMap>> ackWaiters = new TreeMap<>();
-  private boolean closed = false;
-  private @Nullable IntegrationClientException failure;
+  private final TreeMap> ackWaiters = new TreeMap<>();
+  private @Nullable IntegrationClientException terminalFailure;
+  private boolean writeInProgress = false;
+  private boolean halfClosePending = false;
 
   private final long bufferMemory;
   private final Duration maxBlockTime;
@@ -81,7 +84,7 @@ final class ProducerImpl implements Producer, ExactlyOnceProducer {
 
   ProducerImpl(
       IngestionSvcGrpc.IngestionSvcStub stub, ProducerOptions options, String integration) {
-    this(stub, "", DeduplicationMode.DISABLED, false, options, integration);
+    this(stub, "", DeduplicationMode.DISABLED, options, integration);
   }
 
   ProducerImpl(
@@ -89,20 +92,19 @@ final class ProducerImpl implements Producer, ExactlyOnceProducer {
       String producerId,
       ProducerOptions options,
       String integration) {
-    this(stub, producerId, DeduplicationMode.OFFSET_BASED, true, options, integration);
+    this(stub, producerId, DeduplicationMode.OFFSET_BASED, options, integration);
   }
 
   private ProducerImpl(
       IngestionSvcGrpc.IngestionSvcStub stub,
       String producerId,
       DeduplicationMode deduplicationMode,
-      boolean exactlyOnce,
       ProducerOptions options,
       String integration) {
     this.bufferMemory = options.bufferMemory();
     this.maxBlockTime = options.maxBlockTime();
     this.maxBlockNanos = toNanosSaturated(maxBlockTime);
-    this.exactlyOnce = exactlyOnce;
+    this.exactlyOnce = deduplicationMode == DeduplicationMode.OFFSET_BASED;
     // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready
     // handler.
     stub.ingest(new ResponseObserver());
@@ -116,9 +118,7 @@ private ProducerImpl(
                     .setDeduplicationMode(deduplicationMode)
                     .setDefaults(options.toDefaults()))
             .build();
-    synchronized (lock) {
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized").onNext(start);
-    }
+    writeToTransport(start);
   }
 
   // ---- Producer / ExactlyOnceProducer ----
@@ -225,7 +225,11 @@ public CompletableFuture waitAcknowledged(long offset) {
   public long flush() {
     acquire();
     try {
-      return awaitFlush(registerAckWaiter(lastSent));
+      CompletableFuture flush = registerAckWaiter(lastSent);
+      if (cannotBlockInline() && !flush.isDone()) {
+        throw new IllegalStateException("cannot block in a reentrant producer call");
+      }
+      return awaitFlush(flush);
     } finally {
       release();
     }
@@ -248,17 +252,19 @@ public CompletableFuture flushAsync() {
    */
   private CompletableFuture registerAckWaiter(long offset) {
     synchronized (lock) {
-      if (closed) {
-        return CompletableFuture.failedFuture(
-            Objects.requireNonNull(failure, "closed producer has no failure"));
+      if (terminalFailure != null) {
+        return CompletableFuture.failedFuture(terminalFailure);
       }
-      if (offset <= lastCommitted) {
-        return CompletableFuture.completedFuture(lastCommitted);
-      }
-      CompletableFuture f = new CompletableFuture<>();
-      ackWaiters.computeIfAbsent(offset, k -> new ArrayList<>()).add(f);
-      return f;
+      return ackBarrierLocked(offset).copy();
+    }
+  }
+
+  /** Returns the internal completion barrier shared by all waiters for {@code offset}. */
+  private CompletableFuture ackBarrierLocked(long offset) {
+    if (offset <= lastCommitted) {
+      return CompletableFuture.completedFuture(lastCommitted);
     }
+    return ackWaiters.computeIfAbsent(offset, ignored -> new CompletableFuture<>());
   }
 
   @Override
@@ -280,9 +286,7 @@ public void close() {
   private CompletableFuture doSend(long offset, InvocationImpl invocation)
       throws ProducerBufferExhaustedException {
     PreparedSend prepared = prepare(offset, invocation);
-    List> ready;
     CompletableFuture acknowledgement;
-    boolean directWrite;
     long waitStarted = System.nanoTime();
     synchronized (lock) {
       ensureOpenLocked();
@@ -290,6 +294,9 @@ private CompletableFuture doSend(long offset, InvocationImpl invocat
         if (maxBlockNanos == 0) {
           throw admissionTimeout();
         }
+        if (cannotBlockInline()) {
+          throw new IllegalStateException("cannot block in a reentrant producer call");
+        }
         long remaining = maxBlockNanos - (System.nanoTime() - waitStarted);
         if (remaining <= 0) {
           throw admissionTimeout();
@@ -306,30 +313,21 @@ private CompletableFuture doSend(long offset, InvocationImpl invocat
         ensureOpenLocked();
       }
       acknowledgement = acceptLocked(prepared);
-      directWrite = bufferMemory == 0;
-      ready = directWrite ? List.of() : drainLocked();
-    }
-    if (directWrite) {
-      writeDirect(prepared.request());
-      drainAndWake();
-    } else {
-      completeReady(ready);
     }
+    dispatchAccepted(prepared);
     return acknowledgement;
   }
 
   /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
   private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
     PreparedSend prepared = prepare(offset, invocation);
-    List> ready;
     SendAttempt result;
-    boolean directWrite = false;
+    boolean accepted = false;
     synchronized (lock) {
       ensureOpenLocked();
       if (canAdmitLocked(prepared.bufferSize())) {
         result = new SendAttempt.Accepted(acceptLocked(prepared));
-        directWrite = bufferMemory == 0;
-        ready = directWrite ? List.of() : drainLocked();
+        accepted = true;
       } else {
         CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
         AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future);
@@ -343,14 +341,10 @@ private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
               }
             });
         result = new SendAttempt.Backpressured(future);
-        ready = List.of();
       }
     }
-    if (directWrite) {
-      writeDirect(prepared.request());
-      drainAndWake();
-    } else {
-      completeReady(ready);
+    if (accepted) {
+      dispatchAccepted(prepared);
     }
     return result;
   }
@@ -376,73 +370,176 @@ private boolean canAdmitLocked(long requiredBytes) {
     if (bufferMemory == 0) {
       ClientCallStreamObserver observer =
           Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-      return budget > 0 && observer.isReady();
+      return !writeInProgress && budget > 0 && observer.isReady();
     }
     return requiredBytes <= bufferMemory - bufferedBytes;
   }
 
   private CompletableFuture acceptLocked(PreparedSend prepared) {
     lastSent = prepared.offset();
-    CompletableFuture committed = new CompletableFuture<>();
-    ackWaiters.computeIfAbsent(prepared.offset(), ignored -> new ArrayList<>()).add(committed);
+    CompletableFuture committed = ackBarrierLocked(prepared.offset());
     if (bufferMemory == 0) {
       budget -= prepared.windowDebit();
+      writeInProgress = true;
     } else {
-      bufferedSends.addLast(
-          new BufferedSend(prepared.request(), prepared.windowDebit(), prepared.bufferSize()));
+      pendingWrites.addLast(prepared);
       bufferedBytes += prepared.bufferSize();
     }
     return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
   }
 
-  /**
-   * Hand a zero-buffer invocation directly to gRPC, failing the producer if the write is refused.
-   */
-  private void writeDirect(IngestionRequest request) {
+  private void dispatchAccepted(PreparedSend prepared) {
+    if (bufferMemory == 0) {
+      writeDirect(prepared.request());
+      wakeDirectAdmission();
+    } else {
+      drainBuffered();
+    }
+  }
+
+  /** Writes one request without holding {@link #lock}, terminating the producer on failure. */
+  private void writeToTransport(IngestionRequest request) {
+    ClientCallStreamObserver observer =
+        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
     try {
-      Objects.requireNonNull(callObserver, "gRPC request stream was not initialized")
-          .onNext(request);
+      runInlineCallbacks(() -> observer.onNext(request));
     } catch (RuntimeException e) {
-      IntegrationClientException cause =
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN,
-              "failed to write invocation to the ingestion stream",
-              e);
-      terminate(cause, false);
+      IntegrationClientException cause = transportWriteFailure(e);
+      failTransportWrite(observer, cause);
       throw cause;
+    } catch (Error e) {
+      IntegrationClientException cause = transportWriteFailure(e);
+      failTransportWrite(observer, cause);
+      throw e;
     }
   }
 
-  /** Write as many queued records as transport and protocol flow control currently permit. */
-  private List> drainLocked() {
-    if (bufferMemory == 0) {
-      if (!canAdmitLocked(0)) {
-        return List.of();
+  private static IntegrationClientException transportWriteFailure(Throwable cause) {
+    return new IntegrationClientException(
+        IntegrationClientException.Kind.UNKNOWN,
+        "failed to write invocation to the ingestion stream",
+        cause);
+  }
+
+  private void failTransportWrite(
+      ClientCallStreamObserver observer, IntegrationClientException cause) {
+    @Nullable Termination termination;
+    synchronized (lock) {
+      termination = beginTerminationLocked(cause, false);
+    }
+    cancelTransport(observer, cause);
+    if (termination != null) {
+      finishTermination(termination);
+    }
+  }
+
+  private static void cancelTransport(
+      ClientCallStreamObserver observer, IntegrationClientException cause) {
+    try {
+      observer.cancel(cause.getMessage(), cause);
+    } catch (RuntimeException ignored) {
+      // The failed write may already have torn down the call.
+    }
+  }
+
+  private void writeDirect(IngestionRequest request) {
+    boolean failed = true;
+    try {
+      writeToTransport(request);
+      failed = false;
+    } finally {
+      finishWrite(failed);
+    }
+  }
+
+  private void drainFromCallback() {
+    try {
+      if (bufferMemory == 0) {
+        wakeDirectAdmission();
+      } else {
+        drainBuffered();
+      }
+    } catch (IntegrationClientException ignored) {
+      // writeToTransport already made the failure terminal and failed pending futures.
+    }
+  }
+
+  private void wakeDirectAdmission() {
+    List> ready;
+    synchronized (lock) {
+      if (terminalFailure != null || !canAdmitLocked(0)) {
+        return;
       }
       lock.notifyAll();
-      List> ready = new ArrayList<>(admissionWaiters.size());
+      ready = new ArrayList<>(admissionWaiters.size());
       for (AdmissionWaiter waiter : admissionWaiters) {
         ready.add(waiter.future());
       }
       admissionWaiters.clear();
-      return ready;
     }
+    completeReady(ready);
+  }
 
-    boolean freedCapacity = false;
-    ClientCallStreamObserver observer =
-        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-    while (!closed && budget > 0 && observer.isReady() && !bufferedSends.isEmpty()) {
-      BufferedSend send = bufferedSends.removeFirst();
-      bufferedBytes -= send.bufferSize();
-      budget -= send.windowDebit();
-      freedCapacity = true;
-      observer.onNext(send.request());
+  /** Serializes buffered writes while invoking gRPC only outside the state monitor. */
+  private void drainBuffered() {
+    while (true) {
+      PreparedSend send;
+      synchronized (lock) {
+        if (terminalFailure != null || writeInProgress) {
+          return;
+        }
+        ClientCallStreamObserver observer =
+            Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+        if (budget <= 0 || !observer.isReady() || pendingWrites.isEmpty()) {
+          return;
+        }
+        writeInProgress = true;
+        send = pendingWrites.getFirst();
+        budget -= send.windowDebit();
+      }
+
+      try {
+        writeToTransport(send.request());
+      } catch (RuntimeException | Error e) {
+        finishWrite(true);
+        throw e;
+      }
+
+      List> ready;
+      boolean halfClose;
+      synchronized (lock) {
+        if (pendingWrites.peekFirst() == send) {
+          pendingWrites.removeFirst();
+          bufferedBytes -= send.bufferSize();
+          ready = takeCapacityWaitersLocked();
+        } else {
+          // A synchronous terminal callback cleared the queue during the write.
+          ready = List.of();
+        }
+        writeInProgress = false;
+        halfClose = halfClosePending;
+        halfClosePending = false;
+      }
+      if (halfClose) {
+        completeRequestStream();
+      }
+      completeReady(ready);
     }
+  }
 
-    if (!freedCapacity) {
-      return List.of();
+  private void finishWrite(boolean failed) {
+    boolean halfClose;
+    synchronized (lock) {
+      writeInProgress = false;
+      halfClose = halfClosePending && !failed;
+      halfClosePending = false;
     }
+    if (halfClose) {
+      completeRequestStream();
+    }
+  }
 
+  private List> takeCapacityWaitersLocked() {
     lock.notifyAll();
     long available = bufferMemory - bufferedBytes;
     List> ready = new ArrayList<>();
@@ -456,23 +553,36 @@ private void writeDirect(IngestionRequest request) {
     return ready;
   }
 
-  private void drainAndWake() {
-    List> ready;
-    synchronized (lock) {
-      ready = drainLocked();
+  private boolean cannotBlockInline() {
+    return usageGuard.getHoldCount() > 1 || INLINE_CALLBACK.get() != null;
+  }
+
+  private static void runInlineCallbacks(Runnable action) {
+    boolean alreadyInline = INLINE_CALLBACK.get() != null;
+    if (!alreadyInline) {
+      INLINE_CALLBACK.set(true);
+    }
+    try {
+      action.run();
+    } finally {
+      if (!alreadyInline) {
+        INLINE_CALLBACK.remove();
+      }
     }
-    completeReady(ready);
   }
 
   private static void completeReady(List> ready) {
-    for (CompletableFuture<@Nullable Void> future : ready) {
-      future.complete(null);
-    }
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture<@Nullable Void> future : ready) {
+            future.complete(null);
+          }
+        });
   }
 
   private void ensureOpenLocked() {
-    if (closed) {
-      throw new IllegalStateException("producer is closed", failure);
+    if (terminalFailure != null) {
+      throw new IllegalStateException("producer is closed", terminalFailure);
     }
   }
 
@@ -515,84 +625,109 @@ private static long toNanosSaturated(Duration duration) {
   }
 
   private void onResponse(IngestionResponse resp) {
-    @Nullable List> acksToComplete = null;
+    List> acksToComplete = List.of();
     long watermark = -1;
     boolean drain = false;
-    @Nullable IntegrationClientException err = null;
+    @Nullable Termination termination = null;
     synchronized (lock) {
-      if (closed) {
+      if (terminalFailure != null) {
         return;
       }
       if (resp.hasLastCommitted() && resp.getLastCommitted() > lastCommitted) {
         lastCommitted = resp.getLastCommitted();
         watermark = lastCommitted;
         if (!ackWaiters.isEmpty()) {
-          acksToComplete = new ArrayList<>();
-          Map>> head = ackWaiters.headMap(watermark, true);
-          for (List> waiters : head.values()) {
-            acksToComplete.addAll(waiters);
-          }
+          Map> head = ackWaiters.headMap(watermark, true);
+          acksToComplete = new ArrayList<>(head.values());
           head.clear();
         }
       }
-      if (resp.hasWindowUpdate()) {
+      if (resp.hasError()) {
+        IntegrationClientException cause = mapError(resp.getError());
+        termination = beginTerminationLocked(cause, false);
+      } else if (resp.hasWindowUpdate()) {
         // increment_bytes is a uint32; read it as unsigned.
         budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes());
         drain = true;
-      } else if (resp.hasError()) {
-        err = mapError(resp.getError());
       }
     }
-    if (err != null) {
-      terminate(err, false);
-    } else if (drain) {
-      drainAndWake();
+    if (termination != null) {
+      finishTermination(termination);
     }
-    if (acksToComplete != null) {
-      for (CompletableFuture f : acksToComplete) {
-        f.complete(watermark);
-      }
+    List> completedAcks = acksToComplete;
+    long committed = watermark;
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture future : completedAcks) {
+            future.complete(committed);
+          }
+        });
+    if (drain) {
+      drainFromCallback();
     }
   }
 
   /** Mark the producer terminally closed and fail every pending future with {@code cause}. */
   private void terminate(IntegrationClientException cause, boolean halfClose) {
-    List> capacity;
-    List> acks = new ArrayList<>();
+    @Nullable Termination termination;
     synchronized (lock) {
-      if (closed) {
-        return;
-      }
-      closed = true;
-      failure = cause;
-      capacity = new ArrayList<>(admissionWaiters.size());
-      for (AdmissionWaiter waiter : admissionWaiters) {
-        capacity.add(waiter.future());
-      }
-      admissionWaiters.clear();
-      bufferedSends.clear();
-      bufferedBytes = 0;
-      lock.notifyAll();
-      for (List> waiters : ackWaiters.values()) {
-        acks.addAll(waiters);
-      }
-      ackWaiters.clear();
+      termination = beginTerminationLocked(cause, halfClose);
     }
-    if (halfClose) {
-      @Nullable ClientCallStreamObserver obs = callObserver;
-      if (obs != null) {
-        try {
-          obs.onCompleted();
-        } catch (RuntimeException ignored) {
-          // Already torn down transport-side; nothing to half-close.
-        }
-      }
+    if (termination != null) {
+      finishTermination(termination);
     }
-    for (CompletableFuture<@Nullable Void> f : capacity) {
-      f.completeExceptionally(cause);
+  }
+
+  private @Nullable Termination beginTerminationLocked(
+      IntegrationClientException cause, boolean halfClose) {
+    if (terminalFailure != null) {
+      return null;
+    }
+    terminalFailure = cause;
+
+    List> capacity = new ArrayList<>(admissionWaiters.size());
+    for (AdmissionWaiter waiter : admissionWaiters) {
+      capacity.add(waiter.future());
     }
-    for (CompletableFuture f : acks) {
-      f.completeExceptionally(cause);
+    admissionWaiters.clear();
+    pendingWrites.clear();
+    bufferedBytes = 0;
+    lock.notifyAll();
+
+    List> acks = new ArrayList<>(ackWaiters.values());
+    ackWaiters.clear();
+
+    boolean completeStream = halfClose && !writeInProgress;
+    if (halfClose && writeInProgress) {
+      halfClosePending = true;
+    }
+    return new Termination(cause, capacity, acks, completeStream);
+  }
+
+  private void finishTermination(Termination termination) {
+    if (termination.completeStream()) {
+      completeRequestStream();
+    }
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture<@Nullable Void> future : termination.capacity()) {
+            future.completeExceptionally(termination.cause());
+          }
+          for (CompletableFuture future : termination.acknowledgements()) {
+            future.completeExceptionally(termination.cause());
+          }
+        });
+  }
+
+  private void completeRequestStream() {
+    @Nullable ClientCallStreamObserver observer = callObserver;
+    if (observer == null) {
+      return;
+    }
+    try {
+      observer.onCompleted();
+    } catch (RuntimeException ignored) {
+      // Already torn down transport-side; nothing to half-close.
     }
   }
 
@@ -622,21 +757,13 @@ private static IntegrationClientException.Kind mapKind(ErrorKind kind) {
   // ---- fail-fast guard ----
 
   private void acquire() {
-    Thread current = Thread.currentThread();
-    if (owner.get() == current) {
-      reentrancy++;
-      return;
-    }
-    if (!owner.compareAndSet(null, current)) {
+    if (!usageGuard.tryLock()) {
       throw new ConcurrentModificationException("Producer is not safe for multi-threaded access");
     }
-    reentrancy = 1;
   }
 
   private void release() {
-    if (--reentrancy == 0) {
-      owner.set(null);
-    }
+    usageGuard.unlock();
   }
 
   private final class ResponseObserver
@@ -644,7 +771,7 @@ private final class ResponseObserver
     @Override
     public void beforeStart(ClientCallStreamObserver requestStream) {
       callObserver = requestStream;
-      requestStream.setOnReadyHandler(ProducerImpl.this::drainAndWake);
+      requestStream.setOnReadyHandler(ProducerImpl.this::drainFromCallback);
     }
 
     @Override
@@ -674,9 +801,13 @@ public void onCompleted() {
   private record PreparedSend(
       long offset, IngestionRequest request, long windowDebit, long bufferSize) {}
 
-  private record BufferedSend(IngestionRequest request, long windowDebit, long bufferSize) {}
-
   private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
 
+  private record Termination(
+      IntegrationClientException cause,
+      List> capacity,
+      List> acknowledgements,
+      boolean completeStream) {}
+
   private record SendResultImpl(long offset) implements SendResult {}
 }
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 709fb4dd..183b7178 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -36,6 +36,7 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.ConcurrentModificationException;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
@@ -240,11 +241,16 @@ void lastAcknowledgedOffsetRemainsAvailableAfterFailure() throws Exception {
     Producer producer = client.newProducer();
     fake.take(); // Start
     fake.grantWindow(10_000);
-    producer.send(newBody("a"));
-    producer.send(newBody("b"));
+    CompletableFuture committed = producer.send(newBody("a"));
+    CompletableFuture rejected = producer.send(newBody("b"));
 
     fake.error(ErrorKind.ERROR_KIND_BAD_REQUEST, "nope", 0L);
 
+    assertThat(get(committed).offset()).isZero();
+    assertThatThrownBy(() -> get(rejected))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
     assertThat(producer.lastAcknowledgedOffset()).isEqualTo(0L);
   }
 
@@ -354,14 +360,16 @@ void zeroBufferAndZeroMaxBlockTimeFailWithoutConsumingOffset() throws Exception
     fake.assertNoRequest();
   }
 
-  @Test
-  void zeroBufferDirectWriteFailureTerminatesProducer() throws Exception {
+  @ParameterizedTest(name = "bufferMemory={0}")
+  @ValueSource(longs = {0L, 128L})
+  void writeFailureTerminatesProducer(long bufferMemory) throws Exception {
     client.close();
     client =
         GrpcIntegrationClient.builder(new FailingSecondWriteChannel(channel))
             .integration("test-integration", "1.0")
             .build();
-    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    Producer producer =
+        client.newProducer(ProducerOptions.builder().bufferMemory(bufferMemory).build());
     fake.take(); // Start is the first write and succeeds.
     fake.grantWindow(10_000);
 
@@ -378,6 +386,124 @@ void zeroBufferDirectWriteFailureTerminatesProducer() throws Exception {
         .hasCauseInstanceOf(IntegrationClientException.class);
   }
 
+  @Test
+  void callbackDrivenWriteFailureTerminatesProducer() throws Exception {
+    client.close();
+    client =
+        GrpcIntegrationClient.builder(new FailingSecondWriteChannel(channel))
+            .integration("test-integration", "1.0")
+            .build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(128).build());
+    fake.take(); // Start is the first write and succeeds.
+
+    CompletableFuture acknowledgement = producer.send(newBody("a"));
+    fake.grantWindow(10_000);
+
+    assertThatThrownBy(() -> get(acknowledgement))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
+    assertThatThrownBy(() -> producer.send(newBody("b")))
+        .isInstanceOf(IllegalStateException.class)
+        .hasCauseInstanceOf(IntegrationClientException.class);
+  }
+
+  @ParameterizedTest(name = "bufferMemory={0}")
+  @ValueSource(longs = {0L, 128L})
+  void closeDuringWriteDefersHalfClose(long bufferMemory) throws Exception {
+    client.close();
+    DuringInvocationWriteChannel duringWrite = new DuringInvocationWriteChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(duringWrite).integration("test-integration", "1.0").build();
+    Producer producer =
+        client.newProducer(ProducerOptions.builder().bufferMemory(bufferMemory).build());
+    duringWrite.runDuringInvocation(producer::close);
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+
+    CompletableFuture acknowledgement = producer.send(newBody("a"));
+
+    assertThatThrownBy(() -> get(acknowledgement))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
+    assertThat(duringWrite.halfCloseCount()).isOne();
+    assertThat(duringWrite.halfClosedDuringWrite()).isFalse();
+  }
+
+  @ParameterizedTest(name = "bufferMemory={0}")
+  @ValueSource(longs = {0L, ProducerOptions.DEFAULT_BUFFER_MEMORY})
+  void replayBelowKnownWatermarkIsAlreadyAcknowledged(long bufferMemory) throws Exception {
+    ExactlyOnceProducer producer =
+        client.newExactlyOnceProducer(
+            "p1", ProducerOptions.builder().bufferMemory(bufferMemory).build());
+    fake.take(); // Start
+    fake.ack(10L);
+    fake.grantWindow(10_000);
+
+    CompletableFuture replay = producer.send(5L, newBody("replay"));
+
+    assertThat(get(replay).offset()).isEqualTo(5L);
+    assertThat(producer.lastAcknowledgedOffset()).isEqualTo(10L);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(5L);
+  }
+
+  @Test
+  void reentrantTrySendPreservesTransportOrder() throws Exception {
+    client.close();
+    DuringInvocationWriteChannel duringWrite = new DuringInvocationWriteChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(duringWrite).integration("test-integration", "1.0").build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(128).build());
+    AtomicReference reentrantAttempt = new AtomicReference<>();
+    duringWrite.runDuringInvocation(
+        () -> reentrantAttempt.set(producer.trySend(newBody("second"))));
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+
+    producer.send(newBody("first"));
+
+    assertThat(reentrantAttempt.get()).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+    assertThat(fake.take().getInvocation().getOffset()).isOne();
+  }
+
+  @ParameterizedTest(name = "bufferMemory={0}")
+  @ValueSource(longs = {0L, 128L})
+  void reentrantBlockingSendIsRejected(long bufferMemory) throws Exception {
+    client.close();
+    DuringInvocationWriteChannel duringWrite = new DuringInvocationWriteChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(duringWrite).integration("test-integration", "1.0").build();
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder()
+                .bufferMemory(bufferMemory)
+                .maxBlockTime(Duration.ofMillis(100))
+                .build());
+    AtomicReference reentrantFailure = new AtomicReference<>();
+    duringWrite.runDuringInvocation(
+        () -> {
+          try {
+            producer.send(newBody("b".repeat(80)));
+          } catch (Throwable t) {
+            reentrantFailure.set(t);
+          }
+        });
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+
+    CompletableFuture first = producer.send(newBody("a".repeat(80)));
+
+    assertThat(reentrantFailure.get())
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("reentrant");
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+    fake.assertNoRequest();
+    fake.ack(0L);
+    assertThat(get(first).offset()).isZero();
+  }
+
   @Test
   void zeroBufferStreamErrorFailsReadinessWaiter() throws Exception {
     Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
@@ -510,6 +636,23 @@ void waitAcknowledgedCompletesAtWatermark() throws Exception {
     assertThat(get(acked)).isEqualTo(1L);
   }
 
+  @Test
+  void cancellingAcknowledgementViewsDoesNotCancelSharedBarrier() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+
+    CompletableFuture send = producer.send(newBody("a"));
+    CompletableFuture wait = producer.waitAcknowledged(0L);
+    CompletableFuture flush = producer.flushAsync();
+    assertThat(send.cancel(false)).isTrue();
+    assertThat(wait.cancel(false)).isTrue();
+
+    fake.ack(0L);
+
+    assertThat(get(flush)).isZero();
+  }
+
   @Test
   void flushAsyncCompletesWhenEverythingSentIsCommitted() throws Exception {
     Producer producer = client.newProducer();
@@ -554,6 +697,68 @@ void flushBlocksUntilEverythingSentIsCommitted() throws Exception {
     assertThat(get(flushed)).isEqualTo(0L);
   }
 
+  @Test
+  void concurrentUseFailsFastAndSequentialThreadHandoffWorks() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    producer.send(newBody("a"));
+    fake.take();
+
+    CountDownLatch flushing = new CountDownLatch(1);
+    AtomicReference result = new AtomicReference<>();
+    AtomicReference failure = new AtomicReference<>();
+    Thread flusher =
+        new Thread(
+            () -> {
+              flushing.countDown();
+              try {
+                result.set(producer.flush());
+              } catch (Throwable t) {
+                failure.set(t);
+              }
+            });
+    flusher.setDaemon(true);
+    flusher.start();
+    assertThat(flushing.await(5, TimeUnit.SECONDS)).isTrue();
+
+    try {
+      awaitState(flusher, Thread.State.WAITING);
+      assertThatThrownBy(producer::lastSentOffset)
+          .isInstanceOf(ConcurrentModificationException.class);
+    } finally {
+      fake.ack(0L);
+      flusher.join(TimeUnit.SECONDS.toMillis(5));
+    }
+
+    assertThat(flusher.isAlive()).isFalse();
+    assertThat(failure.get()).isNull();
+    assertThat(result.get()).isZero();
+    assertThat(producer.lastSentOffset()).isZero();
+  }
+
+  @Test
+  void blockingFlushFromInlineAcknowledgementCallbackIsRejected() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    CompletableFuture first = producer.send(newBody("a"));
+    CompletableFuture second = producer.send(newBody("b"));
+    fake.take();
+    fake.take();
+    CompletableFuture continuation = first.thenRun(producer::flush);
+
+    get(CompletableFuture.runAsync(() -> fake.ack(0L)));
+
+    assertThatThrownBy(() -> get(continuation))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IllegalStateException.class)
+        .hasMessageContaining("reentrant");
+    fake.ack(1L);
+    assertThat(get(second).offset()).isOne();
+  }
+
   @Test
   void streamErrorFailsPendingFuturesFast() throws Exception {
     Producer producer = client.newProducer();
@@ -673,6 +878,21 @@ private static  T get(CompletableFuture f)
     return f.get(5, TimeUnit.SECONDS);
   }
 
+  private static void awaitState(Thread thread, Thread.State expected) throws InterruptedException {
+    long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+    while (System.nanoTime() < deadline) {
+      if (thread.getState() == expected) {
+        return;
+      }
+      if (!thread.isAlive()) {
+        throw new AssertionError("thread terminated before reaching " + expected);
+      }
+      Thread.sleep(1);
+    }
+    throw new AssertionError(
+        "thread did not reach " + expected + "; current state is " + thread.getState());
+  }
+
   /** Fake service capturing requests and scripting responses. */
   private static final class FakeIngestionService extends IngestionSvcGrpc.IngestionSvcImplBase {
 
@@ -770,4 +990,65 @@ public String authority() {
       return delegate.authority();
     }
   }
+
+  private static final class DuringInvocationWriteChannel extends Channel {
+
+    private final Channel delegate;
+    private volatile Runnable duringInvocation = () -> {};
+    private volatile boolean insideInvocationWrite;
+    private volatile boolean halfClosedDuringWrite;
+    private volatile int halfCloseCount;
+
+    private DuringInvocationWriteChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    void runDuringInvocation(Runnable action) {
+      this.duringInvocation = action;
+    }
+
+    boolean halfClosedDuringWrite() {
+      return halfClosedDuringWrite;
+    }
+
+    int halfCloseCount() {
+      return halfCloseCount;
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        private int writes;
+
+        @Override
+        public void sendMessage(RequestT message) {
+          if (++writes != 2) {
+            super.sendMessage(message);
+            return;
+          }
+          insideInvocationWrite = true;
+          try {
+            duringInvocation.run();
+            super.sendMessage(message);
+          } finally {
+            insideInvocationWrite = false;
+          }
+        }
+
+        @Override
+        public void halfClose() {
+          halfCloseCount++;
+          halfClosedDuringWrite |= insideInvocationWrite;
+          super.halfClose();
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
 }

From 1b63705d2139048b1fd11211f871a5191c2cbc5e Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Tue, 25 Aug 2026 11:48:49 +0200
Subject: [PATCH 15/17] stuff

---
 .../dev/restate/integration/Producer.java     |  10 +-
 .../dev/restate/integration/ProducerImpl.java | 243 +++++++++--------
 .../restate/integration/ProducerOptions.java  |  11 +-
 .../integration/IntegrationClientTest.java    | 244 ++++++++++++++++++
 4 files changed, 391 insertions(+), 117 deletions(-)

diff --git a/integration-client/src/main/java/dev/restate/integration/Producer.java b/integration-client/src/main/java/dev/restate/integration/Producer.java
index 3e220c82..5639dfb6 100644
--- a/integration-client/src/main/java/dev/restate/integration/Producer.java
+++ b/integration-client/src/main/java/dev/restate/integration/Producer.java
@@ -56,8 +56,9 @@
  * }
  * }

  *
- * 
The client assigns monotonically increasing offsets. Producer-level deduplication is disabled;
- * set an idempotency key on an invocation when handler-level deduplication is required.
+ * 
The client assigns monotonically increasing offsets from {@code 0} through {@link
+ * Long#MAX_VALUE}. Producer-level deduplication is disabled; set an idempotency key on an
+ * invocation when handler-level deduplication is required.
  *
  * 
A producer is not thread-safe and fails fast with {@link
  * java.util.ConcurrentModificationException} if used from more than one thread at once.
@@ -79,8 +80,8 @@ public interface Producer extends ProducerBase {
    * @return a future completing, once the invocation is durably committed by Restate.
    * @throws ProducerBufferExhaustedException if the producer cannot admit the invocation before the
    *     configured maximum blocking time elapses, or the thread is interrupted while waiting
-   * @throws IllegalStateException if a reentrant producer callback invokes this method when it
-   *     would block
+   * @throws IllegalStateException if the producer has exhausted the {@code long} offset range, or a
+   *     reentrant producer callback invokes this method when it would block
    * @throws IllegalArgumentException if buffering is enabled and the serialized invocation is
    *     larger than {@link ProducerOptions#bufferMemory()}
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
@@ -99,6 +100,7 @@ public interface Producer extends ProducerBase {
    * @return the admission result
    * @throws IllegalArgumentException if buffering is enabled and the serialized invocation is
    *     larger than {@link ProducerOptions#bufferMemory()}
+   * @throws IllegalStateException if the producer has exhausted the {@code long} offset range
    * @throws java.util.ConcurrentModificationException if the producer is used concurrently from
    *     another thread
    */
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
index 25eaddfb..7f0b2599 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
@@ -71,7 +71,8 @@ final class ProducerImpl implements Producer, ExactlyOnceProducer {
   private final List admissionWaiters = new ArrayList<>();
   private final TreeMap> ackWaiters = new TreeMap<>();
   private @Nullable IntegrationClientException terminalFailure;
-  private boolean writeInProgress = false;
+  // Exactly one thread at a time may call the non-thread-safe outbound observer.
+  private boolean draining = false;
   private boolean halfClosePending = false;
 
   private final long bufferMemory;
@@ -129,7 +130,7 @@ public CompletableFuture send(Invocation invocation)
     acquire();
     try {
       checkMode(false);
-      return doSend(lastSent + 1, (InvocationImpl) invocation);
+      return doSend(nextOffset(), (InvocationImpl) invocation);
     } finally {
       release();
     }
@@ -140,7 +141,7 @@ public SendAttempt trySend(Invocation invocation) {
     acquire();
     try {
       checkMode(false);
-      return doTrySend(lastSent + 1, (InvocationImpl) invocation);
+      return doTrySend(nextOffset(), (InvocationImpl) invocation);
     } finally {
       release();
     }
@@ -178,6 +179,13 @@ private void checkOffset(long offset) {
     }
   }
 
+  private long nextOffset() {
+    if (lastSent == Long.MAX_VALUE) {
+      throw new IllegalStateException("producer offset sequence is exhausted");
+    }
+    return lastSent + 1;
+  }
+
   private void checkMode(boolean exactlyOnceExpected) {
     if (exactlyOnce != exactlyOnceExpected) {
       throw new IllegalStateException(
@@ -252,9 +260,6 @@ public CompletableFuture flushAsync() {
    */
   private CompletableFuture registerAckWaiter(long offset) {
     synchronized (lock) {
-      if (terminalFailure != null) {
-        return CompletableFuture.failedFuture(terminalFailure);
-      }
       return ackBarrierLocked(offset).copy();
     }
   }
@@ -264,6 +269,9 @@ private CompletableFuture ackBarrierLocked(long offset) {
     if (offset <= lastCommitted) {
       return CompletableFuture.completedFuture(lastCommitted);
     }
+    if (terminalFailure != null) {
+      return CompletableFuture.failedFuture(terminalFailure);
+    }
     return ackWaiters.computeIfAbsent(offset, ignored -> new CompletableFuture<>());
   }
 
@@ -286,11 +294,11 @@ public void close() {
   private CompletableFuture doSend(long offset, InvocationImpl invocation)
       throws ProducerBufferExhaustedException {
     PreparedSend prepared = prepare(offset, invocation);
-    CompletableFuture acknowledgement;
+    AcceptedSend accepted;
     long waitStarted = System.nanoTime();
     synchronized (lock) {
       ensureOpenLocked();
-      while (!canAdmitLocked(prepared.bufferSize())) {
+      while (!canAdmitLocked(prepared.size())) {
         if (maxBlockNanos == 0) {
           throw admissionTimeout();
         }
@@ -312,25 +320,28 @@ private CompletableFuture doSend(long offset, InvocationImpl invocat
         }
         ensureOpenLocked();
       }
-      acknowledgement = acceptLocked(prepared);
+      accepted = acceptLocked(prepared);
     }
-    dispatchAccepted(prepared);
-    return acknowledgement;
+    drain(accepted.claimedWrite());
+    return accepted.acknowledgement();
   }
 
   /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
   private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
     PreparedSend prepared = prepare(offset, invocation);
     SendAttempt result;
+    @Nullable PreparedSend claimedWrite = null;
     boolean accepted = false;
     synchronized (lock) {
       ensureOpenLocked();
-      if (canAdmitLocked(prepared.bufferSize())) {
-        result = new SendAttempt.Accepted(acceptLocked(prepared));
+      if (canAdmitLocked(prepared.size())) {
+        AcceptedSend acceptedSend = acceptLocked(prepared);
+        result = new SendAttempt.Accepted(acceptedSend.acknowledgement());
+        claimedWrite = acceptedSend.claimedWrite();
         accepted = true;
       } else {
         CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
-        AdmissionWaiter waiter = new AdmissionWaiter(prepared.bufferSize(), future);
+        AdmissionWaiter waiter = new AdmissionWaiter(prepared.size(), future);
         admissionWaiters.add(waiter);
         future.whenComplete(
             (ignored, failure) -> {
@@ -344,7 +355,7 @@ private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
       }
     }
     if (accepted) {
-      dispatchAccepted(prepared);
+      drain(claimedWrite);
     }
     return result;
   }
@@ -354,47 +365,41 @@ private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
   private PreparedSend prepare(long offset, InvocationImpl invocation) {
     IngestionRequest request =
         IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build();
-    long bufferSize = request.getSerializedSize();
-    if (bufferMemory > 0 && bufferSize > bufferMemory) {
+    long size = request.getInvocation().getSerializedSize();
+    if (bufferMemory > 0 && size > bufferMemory) {
       throw new IllegalArgumentException(
           "serialized invocation requires "
-              + bufferSize
+              + size
               + " bytes, exceeding bufferMemory "
               + bufferMemory);
     }
-    return new PreparedSend(
-        offset, request, request.getInvocation().getSerializedSize(), bufferSize);
+    return new PreparedSend(offset, request, size);
   }
 
   private boolean canAdmitLocked(long requiredBytes) {
     if (bufferMemory == 0) {
       ClientCallStreamObserver observer =
           Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-      return !writeInProgress && budget > 0 && observer.isReady();
+      return !draining && pendingWrites.isEmpty() && budget > 0 && observer.isReady();
     }
     return requiredBytes <= bufferMemory - bufferedBytes;
   }
 
-  private CompletableFuture acceptLocked(PreparedSend prepared) {
+  private AcceptedSend acceptLocked(PreparedSend prepared) {
     lastSent = prepared.offset();
     CompletableFuture committed = ackBarrierLocked(prepared.offset());
+    pendingWrites.addLast(prepared);
+    @Nullable PreparedSend claimedWrite = null;
     if (bufferMemory == 0) {
-      budget -= prepared.windowDebit();
-      writeInProgress = true;
-    } else {
-      pendingWrites.addLast(prepared);
-      bufferedBytes += prepared.bufferSize();
-    }
-    return committed.thenApply(ignored -> new SendResultImpl(prepared.offset()));
-  }
-
-  private void dispatchAccepted(PreparedSend prepared) {
-    if (bufferMemory == 0) {
-      writeDirect(prepared.request());
-      wakeDirectAdmission();
+      // Direct admission reserves the observed readiness for this caller. Do not re-check it.
+      draining = true;
+      budget -= prepared.size();
+      claimedWrite = prepared;
     } else {
-      drainBuffered();
+      bufferedBytes += prepared.size();
     }
+    return new AcceptedSend(
+        committed.thenApply(ignored -> new SendResultImpl(prepared.offset())), claimedWrite);
   }
 
   /** Writes one request without holding {@link #lock}, terminating the producer on failure. */
@@ -426,6 +431,9 @@ private void failTransportWrite(
     @Nullable Termination termination;
     synchronized (lock) {
       termination = beginTerminationLocked(cause, false);
+      draining = false;
+      // A failed write is cancelled, never followed by a deferred half-close.
+      halfClosePending = false;
     }
     cancelTransport(observer, cause);
     if (termination != null) {
@@ -442,105 +450,114 @@ private static void cancelTransport(
     }
   }
 
-  private void writeDirect(IngestionRequest request) {
-    boolean failed = true;
-    try {
-      writeToTransport(request);
-      failed = false;
-    } finally {
-      finishWrite(failed);
-    }
-  }
-
   private void drainFromCallback() {
     try {
-      if (bufferMemory == 0) {
-        wakeDirectAdmission();
-      } else {
-        drainBuffered();
-      }
+      drain(null);
     } catch (IntegrationClientException ignored) {
       // writeToTransport already made the failure terminal and failed pending futures.
     }
   }
 
-  private void wakeDirectAdmission() {
-    List> ready;
-    synchronized (lock) {
-      if (terminalFailure != null || !canAdmitLocked(0)) {
-        return;
-      }
-      lock.notifyAll();
-      ready = new ArrayList<>(admissionWaiters.size());
-      for (AdmissionWaiter waiter : admissionWaiters) {
-        ready.add(waiter.future());
-      }
-      admissionWaiters.clear();
-    }
-    completeReady(ready);
-  }
-
-  /** Serializes buffered writes while invoking gRPC only outside the state monitor. */
-  private void drainBuffered() {
-    while (true) {
-      PreparedSend send;
+  /**
+   * Hands accepted invocations to gRPC in FIFO order.
+   *
+   * 
The queue head remains present while {@code onNext} runs, and {@link #draining} gives that
+   * caller exclusive use of the outbound observer. This lets synchronous callbacks enqueue more
+   * buffered writes, fail the producer, or request a deferred half-close without overlapping gRPC
+   * calls. The observer and user futures are always invoked outside {@link #lock}.
+   */
+  private void drain(@Nullable PreparedSend send) {
+    if (send == null) {
+      List> ready;
       synchronized (lock) {
-        if (terminalFailure != null || writeInProgress) {
+        if (terminalFailure != null || draining) {
           return;
         }
-        ClientCallStreamObserver observer =
-            Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-        if (budget <= 0 || !observer.isReady() || pendingWrites.isEmpty()) {
-          return;
+        send = nextWriteLocked();
+        if (send != null) {
+          draining = true;
+          ready = List.of();
+        } else {
+          ready = takeAdmissionWaitersLocked();
         }
-        writeInProgress = true;
-        send = pendingWrites.getFirst();
-        budget -= send.windowDebit();
       }
-
-      try {
-        writeToTransport(send.request());
-      } catch (RuntimeException | Error e) {
-        finishWrite(true);
-        throw e;
+      if (send == null) {
+        completeReady(ready);
+        return;
       }
+    }
+
+    while (true) {
+      writeToTransport(send.request());
 
       List> ready;
-      boolean halfClose;
       synchronized (lock) {
         if (pendingWrites.peekFirst() == send) {
           pendingWrites.removeFirst();
-          bufferedBytes -= send.bufferSize();
-          ready = takeCapacityWaitersLocked();
+          if (bufferMemory > 0) {
+            bufferedBytes -= send.size();
+          }
+          lock.notifyAll();
+        }
+        ready =
+            terminalFailure == null && bufferMemory > 0 ? takeAdmissionWaitersLocked() : List.of();
+      }
+      completeReady(ready);
+
+      boolean halfClose;
+      @Nullable PreparedSend next;
+      synchronized (lock) {
+        // A synchronous transport or readiness callback may have changed the queue.
+        next = nextWriteLocked();
+        if (next == null) {
+          draining = false;
+          halfClose = halfClosePending;
+          halfClosePending = false;
+          ready = terminalFailure == null ? takeAdmissionWaitersLocked() : List.of();
         } else {
-          // A synchronous terminal callback cleared the queue during the write.
+          halfClose = false;
           ready = List.of();
         }
-        writeInProgress = false;
-        halfClose = halfClosePending;
-        halfClosePending = false;
       }
       if (halfClose) {
         completeRequestStream();
       }
       completeReady(ready);
+      if (next == null) {
+        return;
+      }
+      send = next;
     }
   }
 
-  private void finishWrite(boolean failed) {
-    boolean halfClose;
-    synchronized (lock) {
-      writeInProgress = false;
-      halfClose = halfClosePending && !failed;
-      halfClosePending = false;
+  private @Nullable PreparedSend nextWriteLocked() {
+    if (terminalFailure != null || pendingWrites.isEmpty() || budget <= 0) {
+      return null;
     }
-    if (halfClose) {
-      completeRequestStream();
+    ClientCallStreamObserver observer =
+        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+    if (!observer.isReady()) {
+      return null;
     }
+    PreparedSend send = pendingWrites.getFirst();
+    budget -= send.size();
+    return send;
   }
 
-  private List> takeCapacityWaitersLocked() {
-    lock.notifyAll();
+  private List> takeAdmissionWaitersLocked() {
+    if (bufferMemory == 0) {
+      if (!canAdmitLocked(0)) {
+        return List.of();
+      }
+      lock.notifyAll();
+      List> ready = new ArrayList<>(admissionWaiters.size());
+      for (AdmissionWaiter waiter : admissionWaiters) {
+        ready.add(waiter.future());
+      }
+      admissionWaiters.clear();
+      return ready;
+    }
+
     long available = bufferMemory - bufferedBytes;
     List> ready = new ArrayList<>();
     for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) {
@@ -571,11 +588,19 @@ private static void runInlineCallbacks(Runnable action) {
     }
   }
 
-  private static void completeReady(List> ready) {
+  private void completeReady(List> ready) {
     runInlineCallbacks(
         () -> {
           for (CompletableFuture<@Nullable Void> future : ready) {
-            future.complete(null);
+            @Nullable IntegrationClientException failure;
+            synchronized (lock) {
+              failure = terminalFailure;
+            }
+            if (failure == null) {
+              future.complete(null);
+            } else {
+              future.completeExceptionally(failure);
+            }
           }
         });
   }
@@ -697,8 +722,8 @@ private void terminate(IntegrationClientException cause, boolean halfClose) {
     List> acks = new ArrayList<>(ackWaiters.values());
     ackWaiters.clear();
 
-    boolean completeStream = halfClose && !writeInProgress;
-    if (halfClose && writeInProgress) {
+    boolean completeStream = halfClose && !draining;
+    if (halfClose && draining) {
       halfClosePending = true;
     }
     return new Termination(cause, capacity, acks, completeStream);
@@ -798,8 +823,10 @@ public void onCompleted() {
     }
   }
 
-  private record PreparedSend(
-      long offset, IngestionRequest request, long windowDebit, long bufferSize) {}
+  private record PreparedSend(long offset, IngestionRequest request, long size) {}
+
+  private record AcceptedSend(
+      CompletableFuture acknowledgement, @Nullable PreparedSend claimedWrite) {}
 
   private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
 
diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
index 17607bea..2784ee56 100644
--- a/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerOptions.java
@@ -48,9 +48,10 @@ public static Builder builder() {
   }
 
   /**
-   * Maximum serialized bytes retained while invocations wait to be handed to the transport. A value
-   * of zero disables local buffering: {@code send} waits until the invocation can be handed
-   * directly to the transport, and {@code trySend} reports backpressure until that is possible.
+   * Maximum combined serialized size of accepted invocations retained while they wait to be handed
+   * to the transport. A value of zero disables local buffering: {@code send} waits until the
+   * invocation can be handed directly to the transport, and {@code trySend} reports backpressure
+   * until that is possible.
    *
    * @return the local buffer limit in bytes
    */
@@ -82,8 +83,8 @@ public static final class Builder {
     private Builder() {}
 
     /**
-     * Sets the maximum serialized bytes retained while invocations wait to be handed to the
-     * transport. Set this to zero to disable local buffering.
+     * Sets the maximum combined serialized size of accepted invocations retained while they wait to
+     * be handed to the transport. Set this to zero to disable local buffering.
      *
      * @param bytes a non-negative byte count
      */
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 183b7178..6de12640 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -44,6 +44,7 @@
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -252,6 +253,22 @@ void lastAcknowledgedOffsetRemainsAvailableAfterFailure() throws Exception {
         .cause()
         .isInstanceOf(IntegrationClientException.class);
     assertThat(producer.lastAcknowledgedOffset()).isEqualTo(0L);
+    assertThat(get(producer.waitAcknowledged(0L))).isZero();
+    assertThatThrownBy(() -> get(producer.flushAsync())).isInstanceOf(ExecutionException.class);
+  }
+
+  @Test
+  void fullyAcknowledgedFlushRemainsAvailableAfterFailure() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    CompletableFuture sent = producer.send(newBody("a"));
+
+    fake.error(ErrorKind.ERROR_KIND_BAD_REQUEST, "stream failed after commit", 0L);
+
+    assertThat(get(sent).offset()).isZero();
+    assertThat(get(producer.flushAsync())).isZero();
+    assertThat(producer.flush()).isZero();
   }
 
   @Test
@@ -360,6 +377,30 @@ void zeroBufferAndZeroMaxBlockTimeFailWithoutConsumingOffset() throws Exception
     fake.assertNoRequest();
   }
 
+  @ParameterizedTest(name = "trySend={0}")
+  @ValueSource(booleans = {false, true})
+  void zeroBufferAdmissionReservesObservedReadiness(boolean trySend) throws Exception {
+    client.close();
+    OneShotReadyChannel oneShotReady = new OneShotReadyChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(oneShotReady).integration("test-integration", "1.0").build();
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(0).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    oneShotReady.allowOneReadyCheck();
+
+    CompletableFuture acknowledgement =
+        trySend
+            ? ((SendAttempt.Accepted) producer.trySend(newBody("a"))).acknowledgement()
+            : producer.send(newBody("a"));
+
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+    fake.ack(0L);
+    assertThat(get(acknowledgement).offset()).isZero();
+  }
+
   @ParameterizedTest(name = "bufferMemory={0}")
   @ValueSource(longs = {0L, 128L})
   void writeFailureTerminatesProducer(long bufferMemory) throws Exception {
@@ -468,6 +509,34 @@ void reentrantTrySendPreservesTransportOrder() throws Exception {
     assertThat(fake.take().getInvocation().getOffset()).isOne();
   }
 
+  @Test
+  void concurrentBufferedSendDoesNotOverlapTransportWrites() throws Exception {
+    client.close();
+    GatedWriteChannel gatedWrite = new GatedWriteChannel(channel, 2);
+    client =
+        GrpcIntegrationClient.builder(gatedWrite).integration("test-integration", "1.0").build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(1_024).build());
+    fake.take(); // Start
+    CompletableFuture first = producer.send(newBody("first"));
+
+    CompletableFuture granting = CompletableFuture.runAsync(() -> fake.grantWindow(10_000));
+    assertThat(gatedWrite.awaitEntered()).isTrue();
+    try {
+      CompletableFuture second = producer.send(newBody("second"));
+      assertThat(gatedWrite.maxActiveWrites()).isOne();
+
+      gatedWrite.release();
+      get(granting);
+      assertThat(fake.take().getInvocation().getOffset()).isZero();
+      assertThat(fake.take().getInvocation().getOffset()).isOne();
+      fake.ack(1L);
+      assertThat(get(first).offset()).isZero();
+      assertThat(get(second).offset()).isOne();
+    } finally {
+      gatedWrite.release();
+    }
+  }
+
   @ParameterizedTest(name = "bufferMemory={0}")
   @ValueSource(longs = {0L, 128L})
   void reentrantBlockingSendIsRejected(long bufferMemory) throws Exception {
@@ -566,6 +635,38 @@ void trySendReportsBackpressureAndSignalsWhenCapacityReturns() throws Exception
     assertThat(fake.take().getInvocation().getOffset()).isEqualTo(1L);
   }
 
+  @Test
+  void bufferedReadinessIsSignalledAfterEachHandoff() throws Exception {
+    client.close();
+    GatedWriteChannel gatedWrite = new GatedWriteChannel(channel, 3);
+    client =
+        GrpcIntegrationClient.builder(gatedWrite).integration("test-integration", "1.0").build();
+    Invocation invocation = newBody("a".repeat(80));
+    long nonZeroOffsetSize =
+        ((InvocationImpl) invocation).toProtoInvocation(1L).getSerializedSize();
+    Producer producer =
+        client.newProducer(ProducerOptions.builder().bufferMemory(2 * nonZeroOffsetSize).build());
+    fake.take(); // Start
+
+    assertThat(producer.trySend(invocation)).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(producer.trySend(invocation)).isInstanceOf(SendAttempt.Accepted.class);
+    SendAttempt.Backpressured third = (SendAttempt.Backpressured) producer.trySend(invocation);
+
+    CompletableFuture granting = CompletableFuture.runAsync(() -> fake.grantWindow(10_000));
+    assertThat(gatedWrite.awaitEntered()).isTrue();
+    try {
+      assertThat(third.ready()).isDone();
+    } finally {
+      gatedWrite.release();
+    }
+    get(granting);
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+    assertThat(fake.take().getInvocation().getOffset()).isOne();
+
+    assertThat(producer.trySend(invocation)).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(fake.take().getInvocation().getOffset()).isEqualTo(2L);
+  }
+
   @Test
   void sendBlocksUntilBufferCapacityReturns() throws Exception {
     Producer producer =
@@ -618,6 +719,23 @@ void oversizedInvocationIsRejectedWithoutConsumingOffset() throws Exception {
     assertThat(producer.lastSentOffset()).isEqualTo(-1L);
   }
 
+  @Test
+  void invocationExactlyMatchingBufferLimitIsAccepted() throws Exception {
+    Invocation invocation = newBody("a".repeat(100));
+    long serializedSize = ((InvocationImpl) invocation).toProtoInvocation(0L).getSerializedSize();
+    Producer producer =
+        client.newProducer(ProducerOptions.builder().bufferMemory(serializedSize).build());
+    fake.take(); // Start
+
+    CompletableFuture acknowledgement = producer.send(invocation);
+
+    assertThat(producer.lastSentOffset()).isZero();
+    fake.grantWindow(10_000);
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+    fake.ack(0L);
+    assertThat(get(acknowledgement).offset()).isZero();
+  }
+
   @Test
   void waitAcknowledgedCompletesAtWatermark() throws Exception {
     Producer producer = client.newProducer();
@@ -801,6 +919,34 @@ void streamErrorFailsBufferedAcknowledgementsAndBackpressureWaiters() throws Exc
     assertThatThrownBy(() -> get(backpressured.ready())).isInstanceOf(ExecutionException.class);
   }
 
+  @Test
+  void closeDoesNotFlushAndFailsPendingWork() throws Exception {
+    Producer producer =
+        client.newProducer(
+            ProducerOptions.builder().bufferMemory(128).maxBlockTime(Duration.ZERO).build());
+    fake.take(); // Start
+
+    SendAttempt.Accepted accepted =
+        (SendAttempt.Accepted) producer.trySend(newBody("a".repeat(80)));
+    SendAttempt.Backpressured backpressured =
+        (SendAttempt.Backpressured) producer.trySend(newBody("b".repeat(80)));
+
+    producer.close();
+    producer.close(); // Idempotent.
+
+    assertThatThrownBy(() -> get(accepted.acknowledgement()))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
+    assertThatThrownBy(() -> get(backpressured.ready()))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class);
+    assertThat(producer.lastSentOffset()).isZero();
+    assertThat(producer.lastAcknowledgedOffset()).isEqualTo(-1L);
+    fake.assertNoRequest();
+  }
+
   @Test
   void exactlyOnceRejectsNonIncreasingOffsets() throws Exception {
     ExactlyOnceProducer producer = client.newExactlyOnceProducer("p1");
@@ -991,6 +1137,104 @@ public String authority() {
     }
   }
 
+  private static final class OneShotReadyChannel extends Channel {
+
+    private final Channel delegate;
+    private final AtomicInteger readyMode = new AtomicInteger();
+
+    private OneShotReadyChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    void allowOneReadyCheck() {
+      readyMode.set(1);
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        @Override
+        public boolean isReady() {
+          int mode = readyMode.get();
+          if (mode == 0) {
+            return super.isReady();
+          }
+          if (readyMode.compareAndSet(1, 2)) {
+            return true;
+          }
+          return false;
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
+
+  private static final class GatedWriteChannel extends Channel {
+
+    private final Channel delegate;
+    private final int gatedWrite;
+    private final AtomicInteger writes = new AtomicInteger();
+    private final AtomicInteger activeWrites = new AtomicInteger();
+    private final AtomicInteger maxActiveWrites = new AtomicInteger();
+    private final CountDownLatch entered = new CountDownLatch(1);
+    private final CountDownLatch released = new CountDownLatch(1);
+
+    private GatedWriteChannel(Channel delegate, int gatedWrite) {
+      this.delegate = delegate;
+      this.gatedWrite = gatedWrite;
+    }
+
+    boolean awaitEntered() throws InterruptedException {
+      return entered.await(5, TimeUnit.SECONDS);
+    }
+
+    void release() {
+      released.countDown();
+    }
+
+    int maxActiveWrites() {
+      return maxActiveWrites.get();
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        @Override
+        public void sendMessage(RequestT message) {
+          int active = activeWrites.incrementAndGet();
+          maxActiveWrites.accumulateAndGet(active, Math::max);
+          try {
+            if (writes.incrementAndGet() == gatedWrite) {
+              entered.countDown();
+              if (!released.await(5, TimeUnit.SECONDS)) {
+                throw new AssertionError("timed out waiting to release transport write");
+              }
+            }
+            super.sendMessage(message);
+          } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new AssertionError("interrupted while gating transport write", e);
+          } finally {
+            activeWrites.decrementAndGet();
+          }
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
+
   private static final class DuringInvocationWriteChannel extends Channel {
 
     private final Channel delegate;

From dc58c146a4ebc8c4679141a4580970dc36bd6e47 Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Tue, 25 Aug 2026 11:59:29 +0200
Subject: [PATCH 16/17] remove old impl

---
 .../dev/restate/integration/ProducerImpl.java | 840 ------------------
 1 file changed, 840 deletions(-)
 delete mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerImpl.java

diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
deleted file mode 100644
index 7f0b2599..00000000
--- a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
+++ /dev/null
@@ -1,840 +0,0 @@
-// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
-//
-// This file is part of the Restate Java SDK,
-// which is released under the MIT license.
-//
-// You can find a copy of the license in file LICENSE in the root
-// directory of this repository or package, or at
-// https://github.com/restatedev/sdk-java/blob/main/LICENSE
-package dev.restate.integration;
-
-import dev.restate.ingestion.v1.DeduplicationMode;
-import dev.restate.ingestion.v1.ErrorKind;
-import dev.restate.ingestion.v1.IngestionRequest;
-import dev.restate.ingestion.v1.IngestionResponse;
-import dev.restate.ingestion.v1.IngestionStart;
-import dev.restate.ingestion.v1.IngestionSvcGrpc;
-import io.grpc.stub.ClientCallStreamObserver;
-import io.grpc.stub.ClientResponseObserver;
-import java.time.Duration;
-import java.util.ArrayDeque;
-import java.util.ArrayList;
-import java.util.ConcurrentModificationException;
-import java.util.Iterator;
-import java.util.List;
-import java.util.Map;
-import java.util.Objects;
-import java.util.TreeMap;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.locks.ReentrantLock;
-import org.jspecify.annotations.Nullable;
-
-/**
- * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and
- * the module docs for the concurrency contract.
- *
- * 
When buffering is enabled, accepted records wait in a byte-bounded queue until Restate's
- * send-window has credit ({@code budget}) and the transport is writable ({@code
- * callObserver.isReady()}). With buffering disabled, records are accepted only when they can be
- * handed directly to gRPC. Once handed off, only their acknowledgement futures remain until the
- * commit watermark passes their offsets.
- *
- * 
Two-tier concurrency:
- *
- * 
    
- *   
  • A fail-fast, KafkaConsumer-style guard ({@link #acquire()}/{@link #release()}) rejects
- *       concurrent use from multiple threads; sequential hand-off between threads is fine.
- *   
  • A single monitor ({@link #lock}) guards the small set of fields genuinely shared between
- *       the caller thread and gRPC's callback threads. Futures are always completed outside the
- *       monitor.
- * 

- */
-final class ProducerImpl implements Producer, ExactlyOnceProducer {
-
-  // Blocking from an inline transport/future callback can deadlock gRPC's serialized callback lane.
-  private static final ThreadLocal INLINE_CALLBACK = new ThreadLocal<>();
-
-  private final Object lock = new Object();
-
-  // Set once, synchronously, in beforeStart() before the constructor sends the Start frame.
-  private volatile @Nullable ClientCallStreamObserver callObserver;
-
-  // Reentrant because completing a future can synchronously call back into this producer.
-  private final ReentrantLock usageGuard = new ReentrantLock();
-
-  // ---- state guarded by `lock` ----
-  private long budget = 0; // remaining Restate send window, in bytes; may go one message negative
-  private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet
-  private final ArrayDeque pendingWrites = new ArrayDeque<>();
-  private long bufferedBytes = 0;
-  private final List admissionWaiters = new ArrayList<>();
-  private final TreeMap> ackWaiters = new TreeMap<>();
-  private @Nullable IntegrationClientException terminalFailure;
-  // Exactly one thread at a time may call the non-thread-safe outbound observer.
-  private boolean draining = false;
-  private boolean halfClosePending = false;
-
-  private final long bufferMemory;
-  private final Duration maxBlockTime;
-  private final long maxBlockNanos;
-  private final boolean exactlyOnce;
-
-  // Written only by the (guarded) caller thread; never touched by gRPC callbacks.
-  private long lastSent = -1;
-
-  ProducerImpl(
-      IngestionSvcGrpc.IngestionSvcStub stub, ProducerOptions options, String integration) {
-    this(stub, "", DeduplicationMode.DISABLED, options, integration);
-  }
-
-  ProducerImpl(
-      IngestionSvcGrpc.IngestionSvcStub stub,
-      String producerId,
-      ProducerOptions options,
-      String integration) {
-    this(stub, producerId, DeduplicationMode.OFFSET_BASED, options, integration);
-  }
-
-  private ProducerImpl(
-      IngestionSvcGrpc.IngestionSvcStub stub,
-      String producerId,
-      DeduplicationMode deduplicationMode,
-      ProducerOptions options,
-      String integration) {
-    this.bufferMemory = options.bufferMemory();
-    this.maxBlockTime = options.maxBlockTime();
-    this.maxBlockNanos = toNanosSaturated(maxBlockTime);
-    this.exactlyOnce = deduplicationMode == DeduplicationMode.OFFSET_BASED;
-    // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready
-    // handler.
-    stub.ingest(new ResponseObserver());
-    // Mandatory Start handshake: the first frame on the stream (not flow-controlled).
-    IngestionRequest start =
-        IngestionRequest.newBuilder()
-            .setStart(
-                IngestionStart.newBuilder()
-                    .setProducerId(producerId)
-                    .setIntegration(integration)
-                    .setDeduplicationMode(deduplicationMode)
-                    .setDefaults(options.toDefaults()))
-            .build();
-    writeToTransport(start);
-  }
-
-  // ---- Producer / ExactlyOnceProducer ----
-
-  @Override
-  public CompletableFuture send(Invocation invocation)
-      throws ProducerBufferExhaustedException {
-    acquire();
-    try {
-      checkMode(false);
-      return doSend(nextOffset(), (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public SendAttempt trySend(Invocation invocation) {
-    acquire();
-    try {
-      checkMode(false);
-      return doTrySend(nextOffset(), (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture send(long offset, Invocation invocation)
-      throws ProducerBufferExhaustedException {
-    acquire();
-    try {
-      checkMode(true);
-      checkOffset(offset);
-      return doSend(offset, (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public SendAttempt trySend(long offset, Invocation invocation) {
-    acquire();
-    try {
-      checkMode(true);
-      checkOffset(offset);
-      return doTrySend(offset, (InvocationImpl) invocation);
-    } finally {
-      release();
-    }
-  }
-
-  private void checkOffset(long offset) {
-    if (offset <= lastSent) {
-      throw new IllegalArgumentException(
-          "offset must be strictly increasing; last sent " + lastSent + ", got " + offset);
-    }
-  }
-
-  private long nextOffset() {
-    if (lastSent == Long.MAX_VALUE) {
-      throw new IllegalStateException("producer offset sequence is exhausted");
-    }
-    return lastSent + 1;
-  }
-
-  private void checkMode(boolean exactlyOnceExpected) {
-    if (exactlyOnce != exactlyOnceExpected) {
-      throw new IllegalStateException(
-          exactlyOnce
-              ? "exactly-once producers require explicit offsets"
-              : "at-least-once producers assign offsets automatically");
-    }
-  }
-
-  // ---- ProducerBase ----
-
-  @Override
-  public long lastSentOffset() {
-    acquire();
-    try {
-      return lastSent;
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public long lastAcknowledgedOffset() {
-    acquire();
-    try {
-      synchronized (lock) {
-        return lastCommitted;
-      }
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture waitAcknowledged(long offset) {
-    acquire();
-    try {
-      return registerAckWaiter(offset);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public long flush() {
-    acquire();
-    try {
-      CompletableFuture flush = registerAckWaiter(lastSent);
-      if (cannotBlockInline() && !flush.isDone()) {
-        throw new IllegalStateException("cannot block in a reentrant producer call");
-      }
-      return awaitFlush(flush);
-    } finally {
-      release();
-    }
-  }
-
-  @Override
-  public CompletableFuture flushAsync() {
-    acquire();
-    try {
-      return registerAckWaiter(lastSent);
-    } finally {
-      release();
-    }
-  }
-
-  /**
-   * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark
-   * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are
-   * reentrant, so this is safe to call while already holding {@code lock}).
-   */
-  private CompletableFuture registerAckWaiter(long offset) {
-    synchronized (lock) {
-      return ackBarrierLocked(offset).copy();
-    }
-  }
-
-  /** Returns the internal completion barrier shared by all waiters for {@code offset}. */
-  private CompletableFuture ackBarrierLocked(long offset) {
-    if (offset <= lastCommitted) {
-      return CompletableFuture.completedFuture(lastCommitted);
-    }
-    if (terminalFailure != null) {
-      return CompletableFuture.failedFuture(terminalFailure);
-    }
-    return ackWaiters.computeIfAbsent(offset, ignored -> new CompletableFuture<>());
-  }
-
-  @Override
-  public void close() {
-    acquire();
-    try {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN, "producer closed"),
-          true);
-    } finally {
-      release();
-    }
-  }
-
-  // ---- send path, shared by both producer modes (caller holds the guard) ----
-
-  /** Admit a record, blocking up to the configured maximum while the producer is backpressured. */
-  private CompletableFuture doSend(long offset, InvocationImpl invocation)
-      throws ProducerBufferExhaustedException {
-    PreparedSend prepared = prepare(offset, invocation);
-    AcceptedSend accepted;
-    long waitStarted = System.nanoTime();
-    synchronized (lock) {
-      ensureOpenLocked();
-      while (!canAdmitLocked(prepared.size())) {
-        if (maxBlockNanos == 0) {
-          throw admissionTimeout();
-        }
-        if (cannotBlockInline()) {
-          throw new IllegalStateException("cannot block in a reentrant producer call");
-        }
-        long remaining = maxBlockNanos - (System.nanoTime() - waitStarted);
-        if (remaining <= 0) {
-          throw admissionTimeout();
-        }
-        try {
-          long millis = remaining / 1_000_000;
-          int nanos = (int) (remaining % 1_000_000);
-          lock.wait(millis, nanos);
-        } catch (InterruptedException e) {
-          Thread.currentThread().interrupt();
-          throw new ProducerBufferExhaustedException(
-              "interrupted while waiting for producer admission", e);
-        }
-        ensureOpenLocked();
-      }
-      accepted = acceptLocked(prepared);
-    }
-    drain(accepted.claimedWrite());
-    return accepted.acknowledgement();
-  }
-
-  /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
-  private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
-    PreparedSend prepared = prepare(offset, invocation);
-    SendAttempt result;
-    @Nullable PreparedSend claimedWrite = null;
-    boolean accepted = false;
-    synchronized (lock) {
-      ensureOpenLocked();
-      if (canAdmitLocked(prepared.size())) {
-        AcceptedSend acceptedSend = acceptLocked(prepared);
-        result = new SendAttempt.Accepted(acceptedSend.acknowledgement());
-        claimedWrite = acceptedSend.claimedWrite();
-        accepted = true;
-      } else {
-        CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
-        AdmissionWaiter waiter = new AdmissionWaiter(prepared.size(), future);
-        admissionWaiters.add(waiter);
-        future.whenComplete(
-            (ignored, failure) -> {
-              if (future.isCancelled()) {
-                synchronized (lock) {
-                  admissionWaiters.remove(waiter);
-                }
-              }
-            });
-        result = new SendAttempt.Backpressured(future);
-      }
-    }
-    if (accepted) {
-      drain(claimedWrite);
-    }
-    return result;
-  }
-
-  // ---- internals (all `*Locked` methods require `lock`) ----
-
-  private PreparedSend prepare(long offset, InvocationImpl invocation) {
-    IngestionRequest request =
-        IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build();
-    long size = request.getInvocation().getSerializedSize();
-    if (bufferMemory > 0 && size > bufferMemory) {
-      throw new IllegalArgumentException(
-          "serialized invocation requires "
-              + size
-              + " bytes, exceeding bufferMemory "
-              + bufferMemory);
-    }
-    return new PreparedSend(offset, request, size);
-  }
-
-  private boolean canAdmitLocked(long requiredBytes) {
-    if (bufferMemory == 0) {
-      ClientCallStreamObserver observer =
-          Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-      return !draining && pendingWrites.isEmpty() && budget > 0 && observer.isReady();
-    }
-    return requiredBytes <= bufferMemory - bufferedBytes;
-  }
-
-  private AcceptedSend acceptLocked(PreparedSend prepared) {
-    lastSent = prepared.offset();
-    CompletableFuture committed = ackBarrierLocked(prepared.offset());
-    pendingWrites.addLast(prepared);
-    @Nullable PreparedSend claimedWrite = null;
-    if (bufferMemory == 0) {
-      // Direct admission reserves the observed readiness for this caller. Do not re-check it.
-      draining = true;
-      budget -= prepared.size();
-      claimedWrite = prepared;
-    } else {
-      bufferedBytes += prepared.size();
-    }
-    return new AcceptedSend(
-        committed.thenApply(ignored -> new SendResultImpl(prepared.offset())), claimedWrite);
-  }
-
-  /** Writes one request without holding {@link #lock}, terminating the producer on failure. */
-  private void writeToTransport(IngestionRequest request) {
-    ClientCallStreamObserver observer =
-        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-    try {
-      runInlineCallbacks(() -> observer.onNext(request));
-    } catch (RuntimeException e) {
-      IntegrationClientException cause = transportWriteFailure(e);
-      failTransportWrite(observer, cause);
-      throw cause;
-    } catch (Error e) {
-      IntegrationClientException cause = transportWriteFailure(e);
-      failTransportWrite(observer, cause);
-      throw e;
-    }
-  }
-
-  private static IntegrationClientException transportWriteFailure(Throwable cause) {
-    return new IntegrationClientException(
-        IntegrationClientException.Kind.UNKNOWN,
-        "failed to write invocation to the ingestion stream",
-        cause);
-  }
-
-  private void failTransportWrite(
-      ClientCallStreamObserver observer, IntegrationClientException cause) {
-    @Nullable Termination termination;
-    synchronized (lock) {
-      termination = beginTerminationLocked(cause, false);
-      draining = false;
-      // A failed write is cancelled, never followed by a deferred half-close.
-      halfClosePending = false;
-    }
-    cancelTransport(observer, cause);
-    if (termination != null) {
-      finishTermination(termination);
-    }
-  }
-
-  private static void cancelTransport(
-      ClientCallStreamObserver observer, IntegrationClientException cause) {
-    try {
-      observer.cancel(cause.getMessage(), cause);
-    } catch (RuntimeException ignored) {
-      // The failed write may already have torn down the call.
-    }
-  }
-
-  private void drainFromCallback() {
-    try {
-      drain(null);
-    } catch (IntegrationClientException ignored) {
-      // writeToTransport already made the failure terminal and failed pending futures.
-    }
-  }
-
-  /**
-   * Hands accepted invocations to gRPC in FIFO order.
-   *
-   * 
The queue head remains present while {@code onNext} runs, and {@link #draining} gives that
-   * caller exclusive use of the outbound observer. This lets synchronous callbacks enqueue more
-   * buffered writes, fail the producer, or request a deferred half-close without overlapping gRPC
-   * calls. The observer and user futures are always invoked outside {@link #lock}.
-   */
-  private void drain(@Nullable PreparedSend send) {
-    if (send == null) {
-      List> ready;
-      synchronized (lock) {
-        if (terminalFailure != null || draining) {
-          return;
-        }
-        send = nextWriteLocked();
-        if (send != null) {
-          draining = true;
-          ready = List.of();
-        } else {
-          ready = takeAdmissionWaitersLocked();
-        }
-      }
-      if (send == null) {
-        completeReady(ready);
-        return;
-      }
-    }
-
-    while (true) {
-      writeToTransport(send.request());
-
-      List> ready;
-      synchronized (lock) {
-        if (pendingWrites.peekFirst() == send) {
-          pendingWrites.removeFirst();
-          if (bufferMemory > 0) {
-            bufferedBytes -= send.size();
-          }
-          lock.notifyAll();
-        }
-        ready =
-            terminalFailure == null && bufferMemory > 0 ? takeAdmissionWaitersLocked() : List.of();
-      }
-      completeReady(ready);
-
-      boolean halfClose;
-      @Nullable PreparedSend next;
-      synchronized (lock) {
-        // A synchronous transport or readiness callback may have changed the queue.
-        next = nextWriteLocked();
-        if (next == null) {
-          draining = false;
-          halfClose = halfClosePending;
-          halfClosePending = false;
-          ready = terminalFailure == null ? takeAdmissionWaitersLocked() : List.of();
-        } else {
-          halfClose = false;
-          ready = List.of();
-        }
-      }
-      if (halfClose) {
-        completeRequestStream();
-      }
-      completeReady(ready);
-      if (next == null) {
-        return;
-      }
-      send = next;
-    }
-  }
-
-  private @Nullable PreparedSend nextWriteLocked() {
-    if (terminalFailure != null || pendingWrites.isEmpty() || budget <= 0) {
-      return null;
-    }
-    ClientCallStreamObserver observer =
-        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
-    if (!observer.isReady()) {
-      return null;
-    }
-    PreparedSend send = pendingWrites.getFirst();
-    budget -= send.size();
-    return send;
-  }
-
-  private List> takeAdmissionWaitersLocked() {
-    if (bufferMemory == 0) {
-      if (!canAdmitLocked(0)) {
-        return List.of();
-      }
-      lock.notifyAll();
-      List> ready = new ArrayList<>(admissionWaiters.size());
-      for (AdmissionWaiter waiter : admissionWaiters) {
-        ready.add(waiter.future());
-      }
-      admissionWaiters.clear();
-      return ready;
-    }
-
-    long available = bufferMemory - bufferedBytes;
-    List> ready = new ArrayList<>();
-    for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) {
-      AdmissionWaiter waiter = it.next();
-      if (waiter.requiredBytes() <= available) {
-        ready.add(waiter.future());
-        it.remove();
-      }
-    }
-    return ready;
-  }
-
-  private boolean cannotBlockInline() {
-    return usageGuard.getHoldCount() > 1 || INLINE_CALLBACK.get() != null;
-  }
-
-  private static void runInlineCallbacks(Runnable action) {
-    boolean alreadyInline = INLINE_CALLBACK.get() != null;
-    if (!alreadyInline) {
-      INLINE_CALLBACK.set(true);
-    }
-    try {
-      action.run();
-    } finally {
-      if (!alreadyInline) {
-        INLINE_CALLBACK.remove();
-      }
-    }
-  }
-
-  private void completeReady(List> ready) {
-    runInlineCallbacks(
-        () -> {
-          for (CompletableFuture<@Nullable Void> future : ready) {
-            @Nullable IntegrationClientException failure;
-            synchronized (lock) {
-              failure = terminalFailure;
-            }
-            if (failure == null) {
-              future.complete(null);
-            } else {
-              future.completeExceptionally(failure);
-            }
-          }
-        });
-  }
-
-  private void ensureOpenLocked() {
-    if (terminalFailure != null) {
-      throw new IllegalStateException("producer is closed", terminalFailure);
-    }
-  }
-
-  private ProducerBufferExhaustedException admissionTimeout() {
-    String condition =
-        bufferMemory == 0 ? "producer remained backpressured" : "producer buffer remained full";
-    return new ProducerBufferExhaustedException(condition + " for " + maxBlockTime);
-  }
-
-  private static long awaitFlush(CompletableFuture flush) {
-    try {
-      return flush.get();
-    } catch (InterruptedException e) {
-      Thread.currentThread().interrupt();
-      throw new IntegrationClientException(
-          IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e);
-    } catch (ExecutionException e) {
-      @Nullable Throwable cause = e.getCause();
-      if (cause instanceof RuntimeException runtimeException) {
-        throw runtimeException;
-      }
-      if (cause instanceof Error error) {
-        throw error;
-      }
-      if (cause == null) {
-        throw new IntegrationClientException(
-            IntegrationClientException.Kind.UNKNOWN, "producer flush failed");
-      }
-      throw new IntegrationClientException(
-          IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause);
-    }
-  }
-
-  private static long toNanosSaturated(Duration duration) {
-    try {
-      return duration.toNanos();
-    } catch (ArithmeticException ignored) {
-      return Long.MAX_VALUE;
-    }
-  }
-
-  private void onResponse(IngestionResponse resp) {
-    List> acksToComplete = List.of();
-    long watermark = -1;
-    boolean drain = false;
-    @Nullable Termination termination = null;
-    synchronized (lock) {
-      if (terminalFailure != null) {
-        return;
-      }
-      if (resp.hasLastCommitted() && resp.getLastCommitted() > lastCommitted) {
-        lastCommitted = resp.getLastCommitted();
-        watermark = lastCommitted;
-        if (!ackWaiters.isEmpty()) {
-          Map> head = ackWaiters.headMap(watermark, true);
-          acksToComplete = new ArrayList<>(head.values());
-          head.clear();
-        }
-      }
-      if (resp.hasError()) {
-        IntegrationClientException cause = mapError(resp.getError());
-        termination = beginTerminationLocked(cause, false);
-      } else if (resp.hasWindowUpdate()) {
-        // increment_bytes is a uint32; read it as unsigned.
-        budget += Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes());
-        drain = true;
-      }
-    }
-    if (termination != null) {
-      finishTermination(termination);
-    }
-    List> completedAcks = acksToComplete;
-    long committed = watermark;
-    runInlineCallbacks(
-        () -> {
-          for (CompletableFuture future : completedAcks) {
-            future.complete(committed);
-          }
-        });
-    if (drain) {
-      drainFromCallback();
-    }
-  }
-
-  /** Mark the producer terminally closed and fail every pending future with {@code cause}. */
-  private void terminate(IntegrationClientException cause, boolean halfClose) {
-    @Nullable Termination termination;
-    synchronized (lock) {
-      termination = beginTerminationLocked(cause, halfClose);
-    }
-    if (termination != null) {
-      finishTermination(termination);
-    }
-  }
-
-  private @Nullable Termination beginTerminationLocked(
-      IntegrationClientException cause, boolean halfClose) {
-    if (terminalFailure != null) {
-      return null;
-    }
-    terminalFailure = cause;
-
-    List> capacity = new ArrayList<>(admissionWaiters.size());
-    for (AdmissionWaiter waiter : admissionWaiters) {
-      capacity.add(waiter.future());
-    }
-    admissionWaiters.clear();
-    pendingWrites.clear();
-    bufferedBytes = 0;
-    lock.notifyAll();
-
-    List> acks = new ArrayList<>(ackWaiters.values());
-    ackWaiters.clear();
-
-    boolean completeStream = halfClose && !draining;
-    if (halfClose && draining) {
-      halfClosePending = true;
-    }
-    return new Termination(cause, capacity, acks, completeStream);
-  }
-
-  private void finishTermination(Termination termination) {
-    if (termination.completeStream()) {
-      completeRequestStream();
-    }
-    runInlineCallbacks(
-        () -> {
-          for (CompletableFuture<@Nullable Void> future : termination.capacity()) {
-            future.completeExceptionally(termination.cause());
-          }
-          for (CompletableFuture future : termination.acknowledgements()) {
-            future.completeExceptionally(termination.cause());
-          }
-        });
-  }
-
-  private void completeRequestStream() {
-    @Nullable ClientCallStreamObserver observer = callObserver;
-    if (observer == null) {
-      return;
-    }
-    try {
-      observer.onCompleted();
-    } catch (RuntimeException ignored) {
-      // Already torn down transport-side; nothing to half-close.
-    }
-  }
-
-  private static IntegrationClientException mapError(dev.restate.ingestion.v1.Error error) {
-    String detail =
-        error.hasInvocationOffset()
-            ? "[offset=" + error.getInvocationOffset() + "] " + error.getMessage()
-            : error.getMessage();
-    return new IntegrationClientException(mapKind(error.getKind()), detail);
-  }
-
-  private static IntegrationClientException.Kind mapKind(ErrorKind kind) {
-    switch (kind) {
-      case ERROR_KIND_SHUTTING_DOWN:
-        return IntegrationClientException.Kind.SHUTTING_DOWN;
-      case ERROR_KIND_GO_AWAY:
-        return IntegrationClientException.Kind.GO_AWAY;
-      case ERROR_KIND_NOT_FOUND:
-        return IntegrationClientException.Kind.NOT_FOUND;
-      case ERROR_KIND_BAD_REQUEST:
-        return IntegrationClientException.Kind.BAD_REQUEST;
-      default:
-        return IntegrationClientException.Kind.UNKNOWN;
-    }
-  }
-
-  // ---- fail-fast guard ----
-
-  private void acquire() {
-    if (!usageGuard.tryLock()) {
-      throw new ConcurrentModificationException("Producer is not safe for multi-threaded access");
-    }
-  }
-
-  private void release() {
-    usageGuard.unlock();
-  }
-
-  private final class ResponseObserver
-      implements ClientResponseObserver {
-    @Override
-    public void beforeStart(ClientCallStreamObserver requestStream) {
-      callObserver = requestStream;
-      requestStream.setOnReadyHandler(ProducerImpl.this::drainFromCallback);
-    }
-
-    @Override
-    public void onNext(IngestionResponse value) {
-      onResponse(value);
-    }
-
-    @Override
-    public void onError(Throwable t) {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN,
-              "ingestion stream failed: " + t.getMessage(),
-              t),
-          false);
-    }
-
-    @Override
-    public void onCompleted() {
-      terminate(
-          new IntegrationClientException(
-              IntegrationClientException.Kind.UNKNOWN, "ingestion stream closed by server"),
-          false);
-    }
-  }
-
-  private record PreparedSend(long offset, IngestionRequest request, long size) {}
-
-  private record AcceptedSend(
-      CompletableFuture acknowledgement, @Nullable PreparedSend claimedWrite) {}
-
-  private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
-
-  private record Termination(
-      IntegrationClientException cause,
-      List> capacity,
-      List> acknowledgements,
-      boolean completeStream) {}
-
-  private record SendResultImpl(long offset) implements SendResult {}
-}

From ecae87520e895cb067a9b289d7f4408d773a99ed Mon Sep 17 00:00:00 2001
From: slinkydeveloper 
Date: Tue, 25 Aug 2026 14:46:30 +0200
Subject: [PATCH 17/17] whtever

---
 .../dev/restate/integration/ProducerImpl.java | 1112 +++++++++++++++++
 .../integration/IntegrationClientTest.java    |  411 +++++-
 2 files changed, 1522 insertions(+), 1 deletion(-)
 create mode 100644 integration-client/src/main/java/dev/restate/integration/ProducerImpl.java

diff --git a/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
new file mode 100644
index 00000000..bd2b7770
--- /dev/null
+++ b/integration-client/src/main/java/dev/restate/integration/ProducerImpl.java
@@ -0,0 +1,1112 @@
+// Copyright (c) 2023 - Restate Software, Inc., Restate GmbH
+//
+// This file is part of the Restate Java SDK,
+// which is released under the MIT license.
+//
+// You can find a copy of the license in file LICENSE in the root
+// directory of this repository or package, or at
+// https://github.com/restatedev/sdk-java/blob/main/LICENSE
+package dev.restate.integration;
+
+import dev.restate.ingestion.v1.DeduplicationMode;
+import dev.restate.ingestion.v1.ErrorKind;
+import dev.restate.ingestion.v1.IngestionRequest;
+import dev.restate.ingestion.v1.IngestionResponse;
+import dev.restate.ingestion.v1.IngestionStart;
+import dev.restate.ingestion.v1.IngestionSvcGrpc;
+import io.grpc.stub.ClientCallStreamObserver;
+import io.grpc.stub.ClientResponseObserver;
+import java.time.Duration;
+import java.util.ArrayDeque;
+import java.util.ArrayList;
+import java.util.ConcurrentModificationException;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+import java.util.Objects;
+import java.util.TreeMap;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.locks.ReentrantLock;
+import org.jspecify.annotations.Nullable;
+
+/**
+ * Owns exactly one ingestion bidi stream and all its send-side state. See {@link ProducerBase} and
+ * the module docs for the concurrency contract.
+ *
+ * 
When buffering is enabled, accepted records wait in a byte-bounded queue until Restate's
+ * send-window has credit ({@code budget}) and the transport is writable ({@code
+ * callObserver.isReady()}). With buffering disabled, records are accepted only when they can be
+ * handed directly to gRPC. Once handed off, only their acknowledgement futures remain until the
+ * commit watermark passes their offsets.
+ *
+ * 
Two-tier concurrency:
+ *
+ * 
    
+ *   
  • A fail-fast, KafkaConsumer-style guard ({@link #acquire()}/{@link #release()}) rejects
+ *       concurrent use from multiple threads; sequential hand-off between threads is fine.
+ *   
  • A single monitor ({@link #lock}) guards the small set of fields genuinely shared between
+ *       the caller thread and gRPC's callback threads. Futures are always completed outside the
+ *       monitor.
+ * 

+ */
+final class ProducerImpl implements Producer, ExactlyOnceProducer {
+
+  // Blocking from an inline transport/future callback can deadlock gRPC's serialized callback lane.
+  private static final ThreadLocal INLINE_CALLBACK = new ThreadLocal<>();
+  private static final ReadinessObservation BUFFERED_READY = new ReadinessObservation(true, -1);
+
+  private final Object lock = new Object();
+  // Held across each actual outbound observer call. Terminal state is always recorded under
+  // `lock` before waiting for this gate, so a terminal callback cannot starve behind the drain.
+  private final ReentrantLock outboundLock = new ReentrantLock();
+
+  // Set once, synchronously, in beforeStart() before the constructor sends the Start frame.
+  private volatile @Nullable ClientCallStreamObserver callObserver;
+
+  // Reentrant because completing a future can synchronously call back into this producer.
+  private final ReentrantLock usageGuard = new ReentrantLock();
+
+  // ---- state guarded by `lock` ----
+  private long budget = 0; // remaining Restate send window, in bytes; may go one message negative
+  private long lastCommitted = -1; // ack watermark; -1 == nothing committed yet
+  private final ArrayDeque pendingWrites = new ArrayDeque<>();
+  private long bufferedBytes = 0;
+  private final List admissionWaiters = new ArrayList<>();
+  private final TreeMap> ackWaiters = new TreeMap<>();
+  private @Nullable IntegrationClientException terminalFailure;
+  // Exactly one thread at a time may call the non-thread-safe outbound observer.
+  private boolean draining = false;
+  private boolean outboundWriteActive = false;
+  private boolean outboundTerminated = false;
+  private OutboundAction pendingOutboundAction = OutboundAction.NONE;
+  private long readinessEpoch = 0;
+
+  private final long bufferMemory;
+  private final Duration maxBlockTime;
+  private final long maxBlockNanos;
+  private final boolean exactlyOnce;
+
+  // Written only by the (guarded) caller thread; never touched by gRPC callbacks.
+  private long lastSent = -1;
+
+  ProducerImpl(
+      IngestionSvcGrpc.IngestionSvcStub stub, ProducerOptions options, String integration) {
+    this(stub, "", DeduplicationMode.DISABLED, options, integration);
+  }
+
+  ProducerImpl(
+      IngestionSvcGrpc.IngestionSvcStub stub,
+      String producerId,
+      ProducerOptions options,
+      String integration) {
+    this(stub, producerId, DeduplicationMode.OFFSET_BASED, options, integration);
+  }
+
+  private ProducerImpl(
+      IngestionSvcGrpc.IngestionSvcStub stub,
+      String producerId,
+      DeduplicationMode deduplicationMode,
+      ProducerOptions options,
+      String integration) {
+    this.bufferMemory = options.bufferMemory();
+    this.maxBlockTime = options.maxBlockTime();
+    this.maxBlockNanos = toNanosSaturated(maxBlockTime);
+    this.exactlyOnce = deduplicationMode == DeduplicationMode.OFFSET_BASED;
+    // Opening the call invokes beforeStart() synchronously, wiring callObserver + the ready
+    // handler.
+    stub.ingest(new ResponseObserver());
+    // Mandatory Start handshake: the first frame on the stream (not flow-controlled).
+    IngestionRequest start =
+        IngestionRequest.newBuilder()
+            .setStart(
+                IngestionStart.newBuilder()
+                    .setProducerId(producerId)
+                    .setIntegration(integration)
+                    .setDeduplicationMode(deduplicationMode)
+                    .setDefaults(options.toDefaults()))
+            .build();
+    if (!writeToTransport(start)) {
+      synchronized (lock) {
+        throw new IllegalStateException(
+            "ingestion stream terminated before the producer Start frame", terminalFailure);
+      }
+    }
+  }
+
+  // ---- Producer / ExactlyOnceProducer ----
+
+  @Override
+  public CompletableFuture send(Invocation invocation)
+      throws ProducerBufferExhaustedException {
+    acquire();
+    try {
+      checkMode(false);
+      return doSend(nextOffset(), (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public SendAttempt trySend(Invocation invocation) {
+    acquire();
+    try {
+      checkMode(false);
+      return doTrySend(nextOffset(), (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture send(long offset, Invocation invocation)
+      throws ProducerBufferExhaustedException {
+    acquire();
+    try {
+      checkMode(true);
+      checkOffset(offset);
+      return doSend(offset, (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public SendAttempt trySend(long offset, Invocation invocation) {
+    acquire();
+    try {
+      checkMode(true);
+      checkOffset(offset);
+      return doTrySend(offset, (InvocationImpl) invocation);
+    } finally {
+      release();
+    }
+  }
+
+  private void checkOffset(long offset) {
+    if (offset <= lastSent) {
+      throw new IllegalArgumentException(
+          "offset must be strictly increasing; last sent " + lastSent + ", got " + offset);
+    }
+  }
+
+  private long nextOffset() {
+    if (lastSent == Long.MAX_VALUE) {
+      throw new IllegalStateException("producer offset sequence is exhausted");
+    }
+    return lastSent + 1;
+  }
+
+  private void checkMode(boolean exactlyOnceExpected) {
+    if (exactlyOnce != exactlyOnceExpected) {
+      throw new IllegalStateException(
+          exactlyOnce
+              ? "exactly-once producers require explicit offsets"
+              : "at-least-once producers assign offsets automatically");
+    }
+  }
+
+  // ---- ProducerBase ----
+
+  @Override
+  public long lastSentOffset() {
+    acquire();
+    try {
+      return lastSent;
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public long lastAcknowledgedOffset() {
+    acquire();
+    try {
+      synchronized (lock) {
+        return lastCommitted;
+      }
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture waitAcknowledged(long offset) {
+    acquire();
+    try {
+      return registerAckWaiter(offset);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public long flush() {
+    acquire();
+    try {
+      CompletableFuture flush = registerAckWaiter(lastSent);
+      if (cannotBlockInline() && !flush.isDone()) {
+        throw new IllegalStateException("cannot block in a reentrant producer call");
+      }
+      return awaitFlush(flush);
+    } finally {
+      release();
+    }
+  }
+
+  @Override
+  public CompletableFuture flushAsync() {
+    acquire();
+    try {
+      return registerAckWaiter(lastSent);
+    } finally {
+      release();
+    }
+  }
+
+  /**
+   * Register an ack waiter for {@code offset}. The returned future completes with the ack watermark
+   * once it reaches {@code offset}. Only touches {@code lock}-guarded state (Java monitors are
+   * reentrant, so this is safe to call while already holding {@code lock}).
+   */
+  private CompletableFuture registerAckWaiter(long offset) {
+    synchronized (lock) {
+      return ackBarrierLocked(offset).copy();
+    }
+  }
+
+  /** Returns the internal completion barrier shared by all waiters for {@code offset}. */
+  private CompletableFuture ackBarrierLocked(long offset) {
+    if (offset <= lastCommitted) {
+      return CompletableFuture.completedFuture(lastCommitted);
+    }
+    if (terminalFailure != null) {
+      return CompletableFuture.failedFuture(terminalFailure);
+    }
+    return ackWaiters.computeIfAbsent(offset, ignored -> new CompletableFuture<>());
+  }
+
+  @Override
+  public void close() {
+    acquire();
+    try {
+      terminate(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN, "producer closed"),
+          OutboundAction.HALF_CLOSE);
+    } finally {
+      release();
+    }
+  }
+
+  // ---- send path, shared by both producer modes (caller holds the guard) ----
+
+  /** Admit a record, blocking up to the configured maximum while the producer is backpressured. */
+  private CompletableFuture doSend(long offset, InvocationImpl invocation)
+      throws ProducerBufferExhaustedException {
+    PreparedSend prepared = prepare(offset, invocation);
+    AcceptedSend accepted;
+    long waitStarted = System.nanoTime();
+    while (true) {
+      ReadinessObservation readiness =
+          bufferMemory == 0 ? observeTransportReadiness() : BUFFERED_READY;
+      synchronized (lock) {
+        ensureOpenLocked();
+        if (bufferMemory == 0 && readiness.epoch() != readinessEpoch) {
+          continue;
+        }
+        if (canAdmitLocked(prepared.size(), readiness.ready())) {
+          accepted = acceptLocked(prepared);
+          break;
+        }
+        if (maxBlockNanos == 0) {
+          throw admissionTimeout();
+        }
+        if (cannotBlockInline()) {
+          throw new IllegalStateException("cannot block in a reentrant producer call");
+        }
+        long remaining = maxBlockNanos - (System.nanoTime() - waitStarted);
+        if (remaining <= 0) {
+          throw admissionTimeout();
+        }
+        try {
+          long millis = remaining / 1_000_000;
+          int nanos = (int) (remaining % 1_000_000);
+          lock.wait(millis, nanos);
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+          throw new ProducerBufferExhaustedException(
+              "interrupted while waiting for producer admission", e);
+        }
+        ensureOpenLocked();
+      }
+    }
+    drain(accepted.claimedWrite());
+    return accepted.acknowledgement();
+  }
+
+  /** Attempt to admit a record without blocking or consuming an offset under backpressure. */
+  private SendAttempt doTrySend(long offset, InvocationImpl invocation) {
+    PreparedSend prepared = prepare(offset, invocation);
+    SendAttempt result;
+    @Nullable PreparedSend claimedWrite = null;
+    boolean accepted = false;
+    while (true) {
+      ReadinessObservation readiness =
+          bufferMemory == 0 ? observeTransportReadiness() : BUFFERED_READY;
+      synchronized (lock) {
+        ensureOpenLocked();
+        if (bufferMemory == 0 && readiness.epoch() != readinessEpoch) {
+          continue;
+        }
+        if (canAdmitLocked(prepared.size(), readiness.ready())) {
+          AcceptedSend acceptedSend = acceptLocked(prepared);
+          result = new SendAttempt.Accepted(acceptedSend.acknowledgement());
+          claimedWrite = acceptedSend.claimedWrite();
+          accepted = true;
+        } else {
+          CompletableFuture<@Nullable Void> future = new CompletableFuture<>();
+          AdmissionWaiter waiter = new AdmissionWaiter(prepared.size(), future);
+          admissionWaiters.add(waiter);
+          future.whenComplete(
+              (ignored, failure) -> {
+                if (future.isCancelled()) {
+                  synchronized (lock) {
+                    admissionWaiters.remove(waiter);
+                  }
+                }
+              });
+          result = new SendAttempt.Backpressured(future);
+        }
+        break;
+      }
+    }
+    if (accepted) {
+      drain(claimedWrite);
+    }
+    return result;
+  }
+
+  // ---- internals (all `*Locked` methods require `lock`) ----
+
+  private PreparedSend prepare(long offset, InvocationImpl invocation) {
+    IngestionRequest request =
+        IngestionRequest.newBuilder().setInvocation(invocation.toProtoInvocation(offset)).build();
+    long size = request.getInvocation().getSerializedSize();
+    if (bufferMemory > 0 && size > bufferMemory) {
+      throw new IllegalArgumentException(
+          "serialized invocation requires "
+              + size
+              + " bytes, exceeding bufferMemory "
+              + bufferMemory);
+    }
+    return new PreparedSend(offset, request, size);
+  }
+
+  private boolean canAdmitLocked(long requiredBytes, boolean transportReady) {
+    if (bufferMemory == 0) {
+      return terminalFailure == null
+          && !draining
+          && pendingWrites.isEmpty()
+          && budget > 0
+          && transportReady;
+    }
+    return requiredBytes <= bufferMemory - bufferedBytes;
+  }
+
+  private ReadinessObservation observeTransportReadiness() {
+    long observedEpoch;
+    synchronized (lock) {
+      observedEpoch = readinessEpoch;
+      if (terminalFailure != null || outboundTerminated || outboundWriteActive) {
+        return new ReadinessObservation(false, observedEpoch);
+      }
+    }
+    if (!outboundLock.tryLock()) {
+      return new ReadinessObservation(false, observedEpoch);
+    }
+    boolean reservationAcquired = false;
+    boolean noReadinessSignalBeforeRelease = false;
+    long completedEpoch = observedEpoch;
+    boolean ready = false;
+    try {
+      ClientCallStreamObserver observer;
+      synchronized (lock) {
+        observedEpoch = readinessEpoch;
+        if (terminalFailure != null || outboundTerminated || outboundWriteActive) {
+          return new ReadinessObservation(false, observedEpoch);
+        }
+        outboundWriteActive = true;
+        reservationAcquired = true;
+        observer = Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+      }
+
+      try {
+        ready = observer.isReady();
+      } finally {
+        OutboundAction action;
+        synchronized (lock) {
+          noReadinessSignalBeforeRelease = readinessEpoch == observedEpoch;
+          outboundWriteActive = false;
+          action = claimPendingOutboundActionLocked();
+        }
+        performOutboundAction(action);
+      }
+    } finally {
+      outboundLock.unlock();
+      if (reservationAcquired) {
+        completedEpoch = publishOutboundAvailability(observedEpoch, noReadinessSignalBeforeRelease);
+      }
+    }
+    return new ReadinessObservation(ready, completedEpoch);
+  }
+
+  private AcceptedSend acceptLocked(PreparedSend prepared) {
+    lastSent = prepared.offset();
+    CompletableFuture committed = ackBarrierLocked(prepared.offset());
+    pendingWrites.addLast(prepared);
+    @Nullable PreparedSend claimedWrite = null;
+    if (bufferMemory == 0) {
+      // Direct admission reserves the observed readiness for this caller. Do not re-check it.
+      draining = true;
+      budget -= prepared.size();
+      claimedWrite = prepared;
+    } else {
+      bufferedBytes += prepared.size();
+    }
+    long acceptedOffset = prepared.offset();
+    return new AcceptedSend(
+        committed.thenApply(ignored -> new SendResultImpl(acceptedOffset)), claimedWrite);
+  }
+
+  /**
+   * Writes one request without holding {@link #lock}, terminating the producer on failure.
+   *
+   * @return whether the request was written; {@code false} means a terminal callback won before the
+   *     write began
+   */
+  private boolean writeToTransport(IngestionRequest request) {
+    ClientCallStreamObserver observer =
+        Objects.requireNonNull(callObserver, "gRPC request stream was not initialized");
+    @Nullable Termination termination = null;
+    @Nullable Throwable thrown = null;
+    @Nullable RuntimeException runtimeFailure = null;
+    boolean reservationAcquired = false;
+    outboundLock.lock();
+    try {
+      synchronized (lock) {
+        if (terminalFailure != null || outboundTerminated) {
+          return false;
+        }
+        outboundWriteActive = true;
+        reservationAcquired = true;
+      }
+
+      try {
+        runInlineCallbacks(() -> observer.onNext(request));
+      } catch (RuntimeException | Error t) {
+        thrown = t;
+      }
+
+      OutboundAction outboundAction;
+      synchronized (lock) {
+        if (thrown != null) {
+          termination = beginTerminationLocked(transportWriteFailure(thrown));
+          draining = false;
+          scheduleOutboundActionLocked(OutboundAction.CANCEL);
+          if (thrown instanceof RuntimeException) {
+            runtimeFailure = Objects.requireNonNull(terminalFailure);
+          }
+        }
+        outboundWriteActive = false;
+        outboundAction = claimPendingOutboundActionLocked();
+      }
+      performOutboundAction(outboundAction);
+    } finally {
+      outboundLock.unlock();
+      if (reservationAcquired) {
+        publishOutboundAvailability();
+      }
+    }
+
+    if (termination != null) {
+      finishTermination(termination);
+    }
+    if (thrown instanceof Error error) {
+      throw error;
+    }
+    if (runtimeFailure != null) {
+      throw runtimeFailure;
+    }
+    return true;
+  }
+
+  private static IntegrationClientException transportWriteFailure(Throwable cause) {
+    return new IntegrationClientException(
+        IntegrationClientException.Kind.UNKNOWN,
+        "failed to write invocation to the ingestion stream",
+        cause);
+  }
+
+  private static void cancelTransport(
+      ClientCallStreamObserver observer, IntegrationClientException cause) {
+    try {
+      observer.cancel(cause.getMessage(), cause);
+    } catch (RuntimeException ignored) {
+      // The failed write may already have torn down the call.
+    }
+  }
+
+  private void drainFromCallback() {
+    try {
+      drain(null);
+    } catch (IntegrationClientException ignored) {
+      // writeToTransport already made the failure terminal and failed pending futures.
+    }
+  }
+
+  private void onTransportReady() {
+    synchronized (lock) {
+      readinessEpoch++;
+      lock.notifyAll();
+    }
+    drainFromCallback();
+  }
+
+  /**
+   * Hands accepted invocations to gRPC in FIFO order.
+   *
+   * 
The queue head remains present while {@code onNext} runs, and {@link #draining} gives that
+   * caller exclusive use of the outbound observer. This lets synchronous callbacks enqueue more
+   * buffered writes, fail the producer, or request a deferred half-close without overlapping gRPC
+   * calls. The observer and user futures are always invoked outside {@link #lock}.
+   */
+  private void drain(@Nullable PreparedSend send) {
+    if (send == null) {
+      List> ready;
+      synchronized (lock) {
+        if (terminalFailure != null || draining) {
+          return;
+        }
+        // Claim drain ownership before sampling readiness so a concurrent sender cannot create a
+        // second sampler and then lose its wake-up to this one.
+        draining = true;
+      }
+      while (true) {
+        ReadinessObservation readiness = observeTransportReadiness();
+        synchronized (lock) {
+          if (terminalFailure != null) {
+            draining = false;
+            return;
+          }
+          if (readiness.epoch() != readinessEpoch) {
+            continue;
+          }
+          send = nextWriteLocked(readiness.ready());
+          if (send == null) {
+            draining = false;
+            ready = takeAdmissionWaitersLocked(readiness.ready());
+          } else {
+            ready = List.of();
+          }
+          break;
+        }
+      }
+      if (send == null) {
+        completeReady(ready);
+        return;
+      }
+    }
+
+    while (true) {
+      if (!writeToTransport(send.request())) {
+        synchronized (lock) {
+          draining = false;
+          lock.notifyAll();
+        }
+        return;
+      }
+
+      List> ready;
+      synchronized (lock) {
+        if (pendingWrites.peekFirst() == send) {
+          pendingWrites.removeFirst();
+          if (bufferMemory > 0) {
+            bufferedBytes -= send.size();
+          }
+          lock.notifyAll();
+        }
+        ready =
+            terminalFailure == null && bufferMemory > 0
+                ? takeAdmissionWaitersLocked(false)
+                : List.of();
+      }
+      completeReady(ready);
+
+      @Nullable PreparedSend next;
+      while (true) {
+        ReadinessObservation readiness = observeTransportReadiness();
+        synchronized (lock) {
+          if (readiness.epoch() != readinessEpoch) {
+            continue;
+          }
+          next = nextWriteLocked(readiness.ready());
+          if (next == null) {
+            draining = false;
+            ready =
+                terminalFailure == null ? takeAdmissionWaitersLocked(readiness.ready()) : List.of();
+          } else {
+            ready = List.of();
+          }
+          break;
+        }
+      }
+      completeReady(ready);
+      if (next == null) {
+        return;
+      }
+      send = next;
+    }
+  }
+
+  private @Nullable PreparedSend nextWriteLocked(boolean transportReady) {
+    if (terminalFailure != null || pendingWrites.isEmpty() || budget <= 0 || !transportReady) {
+      return null;
+    }
+    PreparedSend send = pendingWrites.getFirst();
+    budget -= send.size();
+    return send;
+  }
+
+  private List> takeAdmissionWaitersLocked(
+      boolean transportReady) {
+    if (bufferMemory == 0) {
+      if (!canAdmitLocked(0, transportReady)) {
+        return List.of();
+      }
+      lock.notifyAll();
+      List> ready = new ArrayList<>(admissionWaiters.size());
+      for (AdmissionWaiter waiter : admissionWaiters) {
+        ready.add(waiter.future());
+      }
+      admissionWaiters.clear();
+      return ready;
+    }
+
+    long available = bufferMemory - bufferedBytes;
+    List> ready = new ArrayList<>();
+    for (Iterator it = admissionWaiters.iterator(); it.hasNext(); ) {
+      AdmissionWaiter waiter = it.next();
+      if (waiter.requiredBytes() <= available) {
+        ready.add(waiter.future());
+        it.remove();
+      }
+    }
+    return ready;
+  }
+
+  private boolean cannotBlockInline() {
+    return usageGuard.getHoldCount() > 1 || INLINE_CALLBACK.get() != null;
+  }
+
+  private static void runInlineCallbacks(Runnable action) {
+    boolean alreadyInline = INLINE_CALLBACK.get() != null;
+    if (!alreadyInline) {
+      INLINE_CALLBACK.set(true);
+    }
+    try {
+      action.run();
+    } finally {
+      if (!alreadyInline) {
+        INLINE_CALLBACK.remove();
+      }
+    }
+  }
+
+  private void completeReady(List> ready) {
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture<@Nullable Void> future : ready) {
+            @Nullable IntegrationClientException failure;
+            synchronized (lock) {
+              failure = terminalFailure;
+            }
+            if (failure == null) {
+              future.complete(null);
+            } else {
+              future.completeExceptionally(failure);
+            }
+          }
+        });
+  }
+
+  private void ensureOpenLocked() {
+    if (terminalFailure != null) {
+      throw new IllegalStateException("producer is closed", terminalFailure);
+    }
+  }
+
+  private ProducerBufferExhaustedException admissionTimeout() {
+    String condition =
+        bufferMemory == 0 ? "producer remained backpressured" : "producer buffer remained full";
+    return new ProducerBufferExhaustedException(condition + " for " + maxBlockTime);
+  }
+
+  private static long awaitFlush(CompletableFuture flush) {
+    try {
+      return flush.get();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "interrupted while flushing producer", e);
+    } catch (ExecutionException e) {
+      @Nullable Throwable cause = e.getCause();
+      if (cause instanceof RuntimeException runtimeException) {
+        throw runtimeException;
+      }
+      if (cause instanceof Error error) {
+        throw error;
+      }
+      if (cause == null) {
+        throw new IntegrationClientException(
+            IntegrationClientException.Kind.UNKNOWN, "producer flush failed");
+      }
+      throw new IntegrationClientException(
+          IntegrationClientException.Kind.UNKNOWN, "producer flush failed", cause);
+    }
+  }
+
+  private static long toNanosSaturated(Duration duration) {
+    try {
+      return duration.toNanos();
+    } catch (ArithmeticException ignored) {
+      return Long.MAX_VALUE;
+    }
+  }
+
+  private void onResponse(IngestionResponse resp) {
+    if (resp.hasError()) {
+      onErrorResponse(resp);
+      return;
+    }
+
+    Acknowledgements acknowledgements;
+    boolean drain = false;
+    synchronized (lock) {
+      if (terminalFailure != null) {
+        return;
+      }
+      acknowledgements = advanceCommittedLocked(resp);
+      if (resp.hasWindowUpdate()) {
+        // increment_bytes is a uint32; read it as unsigned.
+        long increment = Integer.toUnsignedLong(resp.getWindowUpdate().getIncrementBytes());
+        budget = addSaturated(budget, increment);
+        drain = true;
+      }
+    }
+    completeAcknowledgements(acknowledgements);
+    if (drain) {
+      drainFromCallback();
+    }
+  }
+
+  private void onErrorResponse(IngestionResponse resp) {
+    Acknowledgements acknowledgements = new Acknowledgements(-1, List.of());
+    @Nullable Termination termination = null;
+    synchronized (lock) {
+      if (terminalFailure == null) {
+        acknowledgements = advanceCommittedLocked(resp);
+        termination = beginTerminationLocked(mapError(resp.getError()));
+      }
+      scheduleOutboundActionLocked(OutboundAction.CANCEL);
+    }
+
+    // Shut down the transport before invoking user continuations: a continuation may block, but it
+    // must not prevent an Error response from cancelling the request stream. A response carrying
+    // both fields still completes the acknowledged records before failing the rest.
+    try {
+      flushPendingOutboundAction();
+    } finally {
+      completeAcknowledgements(acknowledgements);
+      if (termination != null) {
+        finishTermination(termination);
+      }
+    }
+  }
+
+  private Acknowledgements advanceCommittedLocked(IngestionResponse resp) {
+    if (!resp.hasLastCommitted()) {
+      return new Acknowledgements(lastCommitted, List.of());
+    }
+
+    long candidate = resp.getLastCommitted();
+    // last_committed is uint64 on the wire. Java exposes values above Long.MAX_VALUE as negative;
+    // every representable producer offset is necessarily covered by such a watermark.
+    if (candidate < 0) {
+      candidate = Long.MAX_VALUE;
+    }
+    if (candidate <= lastCommitted) {
+      return new Acknowledgements(lastCommitted, List.of());
+    }
+
+    lastCommitted = candidate;
+    List> completed = List.of();
+    if (!ackWaiters.isEmpty()) {
+      Map> head = ackWaiters.headMap(lastCommitted, true);
+      completed = new ArrayList<>(head.values());
+      head.clear();
+    }
+    return new Acknowledgements(lastCommitted, completed);
+  }
+
+  private void completeAcknowledgements(Acknowledgements acknowledgements) {
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture future : acknowledgements.futures()) {
+            future.complete(acknowledgements.watermark());
+          }
+        });
+  }
+
+  /** Mark the producer terminally closed and fail every pending future with {@code cause}. */
+  private void terminate(IntegrationClientException cause, OutboundAction outboundAction) {
+    @Nullable Termination termination;
+    synchronized (lock) {
+      termination = beginTerminationLocked(cause);
+      if (termination == null) {
+        return;
+      }
+      scheduleOutboundActionLocked(outboundAction);
+    }
+    // User continuations run inline when their futures complete. End the request stream first so a
+    // blocking continuation cannot delay close/cancellation, while still guaranteeing completion
+    // if the transport observer throws an Error.
+    try {
+      flushPendingOutboundAction();
+    } finally {
+      finishTermination(termination);
+    }
+  }
+
+  private @Nullable Termination beginTerminationLocked(IntegrationClientException cause) {
+    if (terminalFailure != null) {
+      return null;
+    }
+    terminalFailure = cause;
+
+    List> capacity = new ArrayList<>(admissionWaiters.size());
+    for (AdmissionWaiter waiter : admissionWaiters) {
+      capacity.add(waiter.future());
+    }
+    admissionWaiters.clear();
+    pendingWrites.clear();
+    bufferedBytes = 0;
+    lock.notifyAll();
+
+    List> acks = new ArrayList<>(ackWaiters.values());
+    ackWaiters.clear();
+    return new Termination(cause, capacity, acks);
+  }
+
+  private void finishTermination(Termination termination) {
+    runInlineCallbacks(
+        () -> {
+          for (CompletableFuture<@Nullable Void> future : termination.capacity()) {
+            future.completeExceptionally(termination.cause());
+          }
+          for (CompletableFuture future : termination.acknowledgements()) {
+            future.completeExceptionally(termination.cause());
+          }
+        });
+  }
+
+  private void scheduleOutboundActionLocked(OutboundAction action) {
+    if (!outboundTerminated && action.ordinal() > pendingOutboundAction.ordinal()) {
+      pendingOutboundAction = action;
+    }
+  }
+
+  /** Publish that a readiness query or write has released the outbound observer gate. */
+  private void publishOutboundAvailability() {
+    synchronized (lock) {
+      readinessEpoch++;
+      lock.notifyAll();
+    }
+  }
+
+  /** Publish a completed readiness observation and return the epoch its caller may consume. */
+  private long publishOutboundAvailability(
+      long observedEpoch, boolean noReadinessSignalBeforeRelease) {
+    synchronized (lock) {
+      boolean observationIsCurrent =
+          noReadinessSignalBeforeRelease && readinessEpoch == observedEpoch;
+      readinessEpoch++;
+      lock.notifyAll();
+      // If onReady ran at any point during the query, deliberately return a stale epoch so the
+      // caller samples again. Otherwise it may consume the post-release publication.
+      return observationIsCurrent ? readinessEpoch : observedEpoch;
+    }
+  }
+
+  private OutboundAction claimPendingOutboundActionLocked() {
+    if (outboundWriteActive || outboundTerminated || pendingOutboundAction == OutboundAction.NONE) {
+      return OutboundAction.NONE;
+    }
+
+    OutboundAction action = pendingOutboundAction;
+    pendingOutboundAction = OutboundAction.NONE;
+    outboundTerminated = true;
+    return action;
+  }
+
+  private void performOutboundAction(OutboundAction action) {
+    if (action == OutboundAction.NONE) {
+      return;
+    }
+    @Nullable ClientCallStreamObserver observer = callObserver;
+    if (observer == null) {
+      return;
+    }
+
+    if (action == OutboundAction.CANCEL) {
+      IntegrationClientException cause;
+      synchronized (lock) {
+        cause = Objects.requireNonNull(terminalFailure);
+      }
+      cancelTransport(observer, cause);
+    } else {
+      try {
+        observer.onCompleted();
+      } catch (RuntimeException ignored) {
+        // Already torn down transport-side; nothing to half-close.
+      }
+    }
+  }
+
+  private void flushPendingOutboundAction() {
+    outboundLock.lock();
+    try {
+      OutboundAction action;
+      synchronized (lock) {
+        action = claimPendingOutboundActionLocked();
+      }
+      performOutboundAction(action);
+    } finally {
+      outboundLock.unlock();
+    }
+  }
+
+  private void terminateFromTransport(IntegrationClientException cause) {
+    @Nullable Termination termination;
+    synchronized (lock) {
+      // The peer has already ended the RPC, so suppress any deferred local terminal action.
+      outboundTerminated = true;
+      pendingOutboundAction = OutboundAction.NONE;
+      termination = beginTerminationLocked(cause);
+    }
+    // If a writer reserved the outbound gate first, do not return the peer terminal callback until
+    // that write either observes the terminal state or finishes its already-started onNext.
+    outboundLock.lock();
+    outboundLock.unlock();
+    if (termination != null) {
+      finishTermination(termination);
+    }
+  }
+
+  private static long addSaturated(long value, long increment) {
+    return value > Long.MAX_VALUE - increment ? Long.MAX_VALUE : value + increment;
+  }
+
+  private static IntegrationClientException mapError(dev.restate.ingestion.v1.Error error) {
+    String detail =
+        error.hasInvocationOffset()
+            ? "[offset="
+                + Long.toUnsignedString(error.getInvocationOffset())
+                + "] "
+                + error.getMessage()
+            : error.getMessage();
+    return new IntegrationClientException(mapKind(error.getKind()), detail);
+  }
+
+  private static IntegrationClientException.Kind mapKind(ErrorKind kind) {
+    switch (kind) {
+      case ERROR_KIND_SHUTTING_DOWN:
+        return IntegrationClientException.Kind.SHUTTING_DOWN;
+      case ERROR_KIND_GO_AWAY:
+        return IntegrationClientException.Kind.GO_AWAY;
+      case ERROR_KIND_NOT_FOUND:
+        return IntegrationClientException.Kind.NOT_FOUND;
+      case ERROR_KIND_BAD_REQUEST:
+        return IntegrationClientException.Kind.BAD_REQUEST;
+      default:
+        return IntegrationClientException.Kind.UNKNOWN;
+    }
+  }
+
+  // ---- fail-fast guard ----
+
+  private void acquire() {
+    if (!usageGuard.tryLock()) {
+      throw new ConcurrentModificationException("Producer is not safe for multi-threaded access");
+    }
+  }
+
+  private void release() {
+    usageGuard.unlock();
+  }
+
+  private final class ResponseObserver
+      implements ClientResponseObserver {
+    @Override
+    public void beforeStart(ClientCallStreamObserver requestStream) {
+      callObserver = requestStream;
+      requestStream.setOnReadyHandler(ProducerImpl.this::onTransportReady);
+    }
+
+    @Override
+    public void onNext(IngestionResponse value) {
+      onResponse(value);
+    }
+
+    @Override
+    public void onError(Throwable t) {
+      terminateFromTransport(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN,
+              "ingestion stream failed: " + t.getMessage(),
+              t));
+    }
+
+    @Override
+    public void onCompleted() {
+      terminateFromTransport(
+          new IntegrationClientException(
+              IntegrationClientException.Kind.UNKNOWN, "ingestion stream closed by server"));
+    }
+  }
+
+  private record PreparedSend(long offset, IngestionRequest request, long size) {}
+
+  private record AcceptedSend(
+      CompletableFuture acknowledgement, @Nullable PreparedSend claimedWrite) {}
+
+  private record AdmissionWaiter(long requiredBytes, CompletableFuture<@Nullable Void> future) {}
+
+  private record Termination(
+      IntegrationClientException cause,
+      List> capacity,
+      List> acknowledgements) {}
+
+  private record Acknowledgements(long watermark, List> futures) {}
+
+  private record ReadinessObservation(boolean ready, long epoch) {}
+
+  private enum OutboundAction {
+    NONE,
+    HALF_CLOSE,
+    CANCEL
+  }
+
+  private record SendResultImpl(long offset) implements SendResult {}
+}
diff --git a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
index 6de12640..8f4e83c4 100644
--- a/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
+++ b/integration-client/src/test/java/dev/restate/integration/IntegrationClientTest.java
@@ -22,6 +22,7 @@
 import io.grpc.Channel;
 import io.grpc.ClientCall;
 import io.grpc.ForwardingClientCall;
+import io.grpc.ForwardingClientCallListener;
 import io.grpc.ManagedChannel;
 import io.grpc.Metadata;
 import io.grpc.MethodDescriptor;
@@ -34,6 +35,7 @@
 import io.grpc.inprocess.InProcessServerBuilder;
 import io.grpc.stub.StreamObserver;
 import java.io.IOException;
+import java.lang.reflect.Field;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
 import java.util.ConcurrentModificationException;
@@ -46,6 +48,7 @@
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -472,6 +475,30 @@ void closeDuringWriteDefersHalfClose(long bufferMemory) throws Exception {
     assertThat(duringWrite.halfClosedDuringWrite()).isFalse();
   }
 
+  @Test
+  void writeFailureAfterReentrantCloseCancelsWithTheFirstTerminalCause() throws Exception {
+    client.close();
+    DuringInvocationWriteChannel duringWrite = new DuringInvocationWriteChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(duringWrite).integration("test-integration", "1.0").build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(128).build());
+    duringWrite.runDuringInvocation(producer::close);
+    duringWrite.failAfterInvocationAction();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+
+    assertThatThrownBy(() -> producer.send(newBody("a")))
+        .isInstanceOf(IntegrationClientException.class)
+        .hasMessage("producer closed");
+    assertThatThrownBy(() -> get(producer.flushAsync()))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class)
+        .hasMessage("producer closed");
+    assertThat(duringWrite.halfCloseCount()).isZero();
+    assertThat(duringWrite.cancelCount()).isOne();
+  }
+
   @ParameterizedTest(name = "bufferMemory={0}")
   @ValueSource(longs = {0L, ProducerOptions.DEFAULT_BUFFER_MEMORY})
   void replayBelowKnownWatermarkIsAlreadyAcknowledged(long bufferMemory) throws Exception {
@@ -489,6 +516,22 @@ void replayBelowKnownWatermarkIsAlreadyAcknowledged(long bufferMemory) throws Ex
     assertThat(fake.take().getInvocation().getOffset()).isEqualTo(5L);
   }
 
+  @Test
+  void unsignedCommittedWatermarkCoversTheLongOffsetRange() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    CompletableFuture acknowledgement = producer.send(newBody("a"));
+    fake.take(); // Invocation
+
+    // uint64 2^63 is exposed by protobuf as the signed Java value Long.MIN_VALUE.
+    fake.ack(Long.MIN_VALUE);
+
+    assertThat(get(acknowledgement).offset()).isZero();
+    assertThat(producer.lastAcknowledgedOffset()).isEqualTo(Long.MAX_VALUE);
+    assertThat(get(producer.waitAcknowledged(Long.MAX_VALUE))).isEqualTo(Long.MAX_VALUE);
+  }
+
   @Test
   void reentrantTrySendPreservesTransportOrder() throws Exception {
     client.close();
@@ -590,6 +633,97 @@ void zeroBufferStreamErrorFailsReadinessWaiter() throws Exception {
         .isEqualTo(IntegrationClientException.Kind.GO_AWAY);
   }
 
+  @Test
+  void terminalCallbackDuringReadyCheckDoesNotAcceptOrWrite() throws Exception {
+    client.close();
+    DuringReadyCheckChannel duringReady = new DuringReadyCheckChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(duringReady).integration("test-integration", "1.0").build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    duringReady.runOnce(
+        () -> fake.errorWithoutCompleting(ErrorKind.ERROR_KIND_GO_AWAY, "terminal"));
+
+    assertThatThrownBy(() -> producer.trySend(newBody("a")))
+        .isInstanceOf(IllegalStateException.class)
+        .hasCauseInstanceOf(IntegrationClientException.class);
+    assertThat(producer.lastSentOffset()).isEqualTo(-1L);
+    fake.assertNoRequest();
+  }
+
+  @Test
+  void readySignalDuringReadinessCheckIsNotLost() throws Exception {
+    client.close();
+    ReadySignalDuringCheckChannel readyDuringCheck = new ReadySignalDuringCheckChannel(channel);
+    client =
+        GrpcIntegrationClient.builder(readyDuringCheck)
+            .integration("test-integration", "1.0")
+            .build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    fake.take(); // Start
+    fake.grantWindow(10_000);
+    readyDuringCheck.signalDuringNextCheck();
+
+    SendAttempt attempt = producer.trySend(newBody("a"));
+
+    assertThat(attempt).isInstanceOf(SendAttempt.Accepted.class);
+    assertThat(fake.take().getInvocation().getOffset()).isZero();
+  }
+
+  @Test
+  void zeroBufferBusyObservationCannotMissCompletedReadinessCheck() throws Exception {
+    client.close();
+    client =
+        GrpcIntegrationClient.builder(new AlwaysReadyWithoutSignalsChannel(channel))
+            .integration("test-integration", "1.0")
+            .build();
+    Producer producer = client.newProducer(ProducerOptions.builder().bufferMemory(0).build());
+    fake.take(); // Start
+
+    CoordinatedOutboundLock outboundLock = new CoordinatedOutboundLock();
+    Field field = ProducerImpl.class.getDeclaredField("outboundLock");
+    field.setAccessible(true);
+    field.set(producer, outboundLock);
+
+    CompletableFuture callback =
+        CompletableFuture.runAsync(
+            () -> {
+              outboundLock.designateCallbackThread();
+              fake.grantWindow(10_000);
+            });
+    CompletableFuture attempted = null;
+    try {
+      assertThat(outboundLock.awaitCallbackGate()).isTrue();
+      attempted =
+          CompletableFuture.supplyAsync(
+              () -> {
+                outboundLock.designateSenderThread();
+                return producer.trySend(newBody("a"));
+              });
+      assertThat(outboundLock.awaitSenderBusy()).isTrue();
+
+      // The callback has already reserved the outbound gate, but has not yet published that fact.
+      // Let it finish a successful readiness check while the sender is still returning BUSY. No
+      // later transport onReady signal will repair a waiter registered from that stale result.
+      outboundLock.releaseCallback();
+      get(callback);
+      outboundLock.releaseSender();
+
+      SendAttempt attempt = get(attempted);
+      if (attempt instanceof SendAttempt.Backpressured backpressured) {
+        get(backpressured.ready());
+        attempt = producer.trySend(newBody("a"));
+      }
+
+      assertThat(attempt).isInstanceOf(SendAttempt.Accepted.class);
+      assertThat(fake.take().getInvocation().getOffset()).isZero();
+    } finally {
+      outboundLock.releaseCallback();
+      outboundLock.releaseSender();
+    }
+  }
+
   @Test
   void exactlyOnceZeroBufferBackpressureDoesNotConsumeOffset() throws Exception {
     ExactlyOnceProducer producer =
@@ -815,6 +949,39 @@ void flushBlocksUntilEverythingSentIsCommitted() throws Exception {
     assertThat(get(flushed)).isEqualTo(0L);
   }
 
+  @Test
+  void interruptedFlushRestoresInterruptAndReleasesUsageGuard() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    producer.send(newBody("a"));
+
+    AtomicReference failure = new AtomicReference<>();
+    AtomicReference interrupted = new AtomicReference<>(false);
+    Thread flusher =
+        new Thread(
+            () -> {
+              try {
+                producer.flush();
+              } catch (Throwable t) {
+                failure.set(t);
+                interrupted.set(Thread.currentThread().isInterrupted());
+              }
+            });
+    flusher.setDaemon(true);
+    flusher.start();
+    awaitState(flusher, Thread.State.WAITING);
+
+    flusher.interrupt();
+    flusher.join(TimeUnit.SECONDS.toMillis(5));
+
+    assertThat(flusher.isAlive()).isFalse();
+    assertThat(failure.get())
+        .isInstanceOf(IntegrationClientException.class)
+        .hasCauseInstanceOf(InterruptedException.class);
+    assertThat(interrupted.get()).isTrue();
+    assertThat(producer.lastSentOffset()).isZero();
+  }
+
   @Test
   void concurrentUseFailsFastAndSequentialThreadHandoffWorks() throws Exception {
     Producer producer = client.newProducer();
@@ -900,6 +1067,22 @@ void streamErrorFailsPendingFuturesFast() throws Exception {
     assertThatThrownBy(() -> producer.send(newBody("b"))).isInstanceOf(IllegalStateException.class);
   }
 
+  @Test
+  void protocolErrorIncludesRecordOffsetAndCancelsOutboundStream() throws Exception {
+    Producer producer = client.newProducer();
+    fake.take(); // Start
+    CompletableFuture pending = producer.send(newBody("a"));
+
+    fake.errorAtWithoutCompleting(ErrorKind.ERROR_KIND_BAD_REQUEST, "nope", -1L);
+
+    assertThatThrownBy(() -> get(pending))
+        .isInstanceOf(ExecutionException.class)
+        .cause()
+        .isInstanceOf(IntegrationClientException.class)
+        .hasMessage("[offset=18446744073709551615] nope");
+    assertThat(fake.awaitRequestFailure()).isTrue();
+  }
+
   @Test
   void streamErrorFailsBufferedAcknowledgementsAndBackpressureWaiters() throws Exception {
     Producer producer =
@@ -1043,6 +1226,7 @@ private static void awaitState(Thread thread, Thread.State expected) throws Inte
   private static final class FakeIngestionService extends IngestionSvcGrpc.IngestionSvcImplBase {
 
     private final BlockingQueue received = new LinkedBlockingQueue<>();
+    private final CountDownLatch requestFailed = new CountDownLatch(1);
     private volatile StreamObserver responses;
 
     @Override
@@ -1056,7 +1240,9 @@ public void onNext(IngestionRequest value) {
         }
 
         @Override
-        public void onError(Throwable t) {}
+        public void onError(Throwable t) {
+          requestFailed.countDown();
+        }
 
         @Override
         public void onCompleted() {}
@@ -1104,6 +1290,29 @@ void error(ErrorKind kind, String message, long lastCommitted) {
               .build());
       responses.onCompleted();
     }
+
+    void errorWithoutCompleting(ErrorKind kind, String message) {
+      responses.onNext(
+          IngestionResponse.newBuilder()
+              .setError(
+                  dev.restate.ingestion.v1.Error.newBuilder().setKind(kind).setMessage(message))
+              .build());
+    }
+
+    void errorAtWithoutCompleting(ErrorKind kind, String message, long invocationOffset) {
+      responses.onNext(
+          IngestionResponse.newBuilder()
+              .setError(
+                  dev.restate.ingestion.v1.Error.newBuilder()
+                      .setKind(kind)
+                      .setMessage(message)
+                      .setInvocationOffset(invocationOffset))
+              .build());
+    }
+
+    boolean awaitRequestFailure() throws InterruptedException {
+      return requestFailed.await(5, TimeUnit.SECONDS);
+    }
   }
 
   private static final class FailingSecondWriteChannel extends Channel {
@@ -1175,6 +1384,187 @@ public String authority() {
     }
   }
 
+  private static final class DuringReadyCheckChannel extends Channel {
+
+    private final Channel delegate;
+    private final AtomicReference action = new AtomicReference<>();
+
+    private DuringReadyCheckChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    void runOnce(Runnable action) {
+      this.action.set(action);
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        @Override
+        public boolean isReady() {
+          boolean ready = super.isReady();
+          Runnable once = action.getAndSet(null);
+          if (once != null) {
+            once.run();
+          }
+          return ready;
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
+
+  private static final class ReadySignalDuringCheckChannel extends Channel {
+
+    private final Channel delegate;
+    private final AtomicInteger mode = new AtomicInteger();
+
+    private ReadySignalDuringCheckChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    void signalDuringNextCheck() {
+      mode.set(1);
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        private ClientCall.Listener listener;
+
+        @Override
+        public void start(ClientCall.Listener listener, Metadata headers) {
+          this.listener = listener;
+          super.start(listener, headers);
+        }
+
+        @Override
+        public boolean isReady() {
+          if (mode.compareAndSet(1, 2)) {
+            listener.onReady();
+            return false;
+          }
+          if (mode.get() == 2) {
+            return true;
+          }
+          return super.isReady();
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
+
+  /** Reports readiness when sampled, but deliberately suppresses asynchronous onReady signals. */
+  private static final class AlwaysReadyWithoutSignalsChannel extends Channel {
+
+    private final Channel delegate;
+
+    private AlwaysReadyWithoutSignalsChannel(Channel delegate) {
+      this.delegate = delegate;
+    }
+
+    @Override
+    public  ClientCall newCall(
+        MethodDescriptor methodDescriptor, CallOptions callOptions) {
+      return new ForwardingClientCall.SimpleForwardingClientCall<>(
+          delegate.newCall(methodDescriptor, callOptions)) {
+        @Override
+        public void start(ClientCall.Listener listener, Metadata headers) {
+          super.start(
+              new ForwardingClientCallListener.SimpleForwardingClientCallListener<>(listener) {
+                @Override
+                public void onReady() {
+                  // This fixture only exposes readiness through isReady().
+                }
+              },
+              headers);
+        }
+
+        @Override
+        public boolean isReady() {
+          return true;
+        }
+      };
+    }
+
+    @Override
+    public String authority() {
+      return delegate.authority();
+    }
+  }
+
+  /** Pauses the two tryLock calls around the stale BUSY observation under test. */
+  private static final class CoordinatedOutboundLock extends ReentrantLock {
+
+    private final CountDownLatch callbackHasGate = new CountDownLatch(1);
+    private final CountDownLatch allowCallback = new CountDownLatch(1);
+    private final CountDownLatch senderSawBusy = new CountDownLatch(1);
+    private final CountDownLatch allowSender = new CountDownLatch(1);
+    private volatile Thread callbackThread;
+    private volatile Thread senderThread;
+
+    void designateCallbackThread() {
+      callbackThread = Thread.currentThread();
+    }
+
+    void designateSenderThread() {
+      senderThread = Thread.currentThread();
+    }
+
+    boolean awaitCallbackGate() throws InterruptedException {
+      return callbackHasGate.await(5, TimeUnit.SECONDS);
+    }
+
+    boolean awaitSenderBusy() throws InterruptedException {
+      return senderSawBusy.await(5, TimeUnit.SECONDS);
+    }
+
+    void releaseCallback() {
+      allowCallback.countDown();
+    }
+
+    void releaseSender() {
+      allowSender.countDown();
+    }
+
+    @Override
+    public boolean tryLock() {
+      boolean acquired = super.tryLock();
+      Thread current = Thread.currentThread();
+      if (current == callbackThread && acquired) {
+        callbackHasGate.countDown();
+        await(allowCallback, "callback readiness check");
+      } else if (current == senderThread && !acquired) {
+        senderSawBusy.countDown();
+        await(allowSender, "sender BUSY observation");
+      }
+      return acquired;
+    }
+
+    private static void await(CountDownLatch latch, String operation) {
+      try {
+        if (!latch.await(5, TimeUnit.SECONDS)) {
+          throw new AssertionError("timed out coordinating " + operation);
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        throw new AssertionError("interrupted while coordinating " + operation, e);
+      }
+    }
+  }
+
   private static final class GatedWriteChannel extends Channel {
 
     private final Channel delegate;
@@ -1242,6 +1632,8 @@ private static final class DuringInvocationWriteChannel extends Channel {
     private volatile boolean insideInvocationWrite;
     private volatile boolean halfClosedDuringWrite;
     private volatile int halfCloseCount;
+    private volatile int cancelCount;
+    private volatile boolean failAfterInvocationAction;
 
     private DuringInvocationWriteChannel(Channel delegate) {
       this.delegate = delegate;
@@ -1251,6 +1643,10 @@ void runDuringInvocation(Runnable action) {
       this.duringInvocation = action;
     }
 
+    void failAfterInvocationAction() {
+      failAfterInvocationAction = true;
+    }
+
     boolean halfClosedDuringWrite() {
       return halfClosedDuringWrite;
     }
@@ -1259,6 +1655,10 @@ int halfCloseCount() {
       return halfCloseCount;
     }
 
+    int cancelCount() {
+      return cancelCount;
+    }
+
     @Override
     public  ClientCall newCall(
         MethodDescriptor methodDescriptor, CallOptions callOptions) {
@@ -1275,6 +1675,9 @@ public void sendMessage(RequestT message) {
           insideInvocationWrite = true;
           try {
             duringInvocation.run();
+            if (failAfterInvocationAction) {
+              throw new IllegalStateException("simulated failure after invocation action");
+            }
             super.sendMessage(message);
           } finally {
             insideInvocationWrite = false;
@@ -1287,6 +1690,12 @@ public void halfClose() {
           halfClosedDuringWrite |= insideInvocationWrite;
           super.halfClose();
         }
+
+        @Override
+        public void cancel(String message, Throwable cause) {
+          cancelCount++;
+          super.cancel(message, cause);
+        }
       };
     }