This is an automated email from the ASF dual-hosted git repository. dmsolr pushed a commit to branch fix/grpc-bounded-aggregated-response in repository https://gitbox.apache.org/repos/asf/shenyu.git
commit dfd309ebecadc44f101a45c371c0b65e4c3579a8 Author: Luke Haochao Zhuang <[email protected]> AuthorDate: Wed Oct 21 18:11:29 2026 +0800 fix: bound gRPC aggregated streaming responses with configurable budgets The gRPC plugin aggregates SERVER_STREAMING/BIDI_STREAMING responses into an unbounded ArrayList (ShenyuGrpcResponse.results) that only completes when the RPC terminates, so retained results can grow without limit and the HTTP client consumption rate cannot regulate upstream receipt. Add opt-in aggregate budgets to the gRPC plugin config: - GrpcRegisterConfig gains maxAggregatedResultCount and maxAggregatedResultBytes (null/non-positive = unlimited, preserving existing behavior). - MessageWriter tracks retained item count and UTF-8 payload bytes; on budget exhaustion it releases retained results, stops appending and cancels the underlying ClientCall, so the call future completes exceptionally and a documented error is returned instead of a success-shaped partial aggregate. - ShenyuGrpcClient resolves the budgets from config and binds the ClientCall to the writer (invoke now creates the call once). The byte budget bounds serialized payload bytes, not exact JVM heap. Each call owns its writer/response, so concurrent RPCs stay isolated. Fixes #7928 --- .../dto/convert/plugin/GrpcRegisterConfig.java | 61 +++++++- .../dto/convert/plugin/GrpcRegisterConfigTest.java | 21 +++ .../plugin/grpc/client/ShenyuGrpcClient.java | 42 +++++- .../shenyu/plugin/grpc/proto/MessageWriter.java | 106 +++++++++++++- .../plugin/grpc/proto/MessageWriterTest.java | 157 +++++++++++++++++++-- 5 files changed, 365 insertions(+), 22 deletions(-) diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfig.java b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfig.java index 7ca6a41275..08fb4daa95 100644 --- a/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfig.java +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfig.java @@ -29,6 +29,21 @@ public class GrpcRegisterConfig implements Serializable { private String threadpool; + /** + * Maximum number of aggregated response items retained by a single gRPC call + * (including SERVER_STREAMING / BIDI_STREAMING). A value that is {@code null} + * or not positive means unlimited, which preserves the historical behavior. + */ + private Long maxAggregatedResultCount; + + /** + * Maximum accumulated payload bytes (UTF-8) of the aggregated response items retained + * by a single gRPC call. This bounds the serialized payload bytes, not the exact JVM heap + * footprint. A value that is {@code null} or not positive means unlimited, which preserves + * the historical behavior. + */ + private Long maxAggregatedResultBytes; + /** * get threadpool. * @@ -47,6 +62,42 @@ public class GrpcRegisterConfig implements Serializable { this.threadpool = threadpool; } + /** + * get maxAggregatedResultCount. + * + * @return maxAggregatedResultCount + */ + public Long getMaxAggregatedResultCount() { + return maxAggregatedResultCount; + } + + /** + * set maxAggregatedResultCount. + * + * @param maxAggregatedResultCount maxAggregatedResultCount + */ + public void setMaxAggregatedResultCount(final Long maxAggregatedResultCount) { + this.maxAggregatedResultCount = maxAggregatedResultCount; + } + + /** + * get maxAggregatedResultBytes. + * + * @return maxAggregatedResultBytes + */ + public Long getMaxAggregatedResultBytes() { + return maxAggregatedResultBytes; + } + + /** + * set maxAggregatedResultBytes. + * + * @param maxAggregatedResultBytes maxAggregatedResultBytes + */ + public void setMaxAggregatedResultBytes(final Long maxAggregatedResultBytes) { + this.maxAggregatedResultBytes = maxAggregatedResultBytes; + } + @Override public boolean equals(final Object o) { if (this == o) { @@ -56,12 +107,14 @@ public class GrpcRegisterConfig implements Serializable { return false; } GrpcRegisterConfig that = (GrpcRegisterConfig) o; - return Objects.equals(threadpool, that.threadpool); + return Objects.equals(threadpool, that.threadpool) + && Objects.equals(maxAggregatedResultCount, that.maxAggregatedResultCount) + && Objects.equals(maxAggregatedResultBytes, that.maxAggregatedResultBytes); } @Override public int hashCode() { - return Objects.hash(threadpool); + return Objects.hash(threadpool, maxAggregatedResultCount, maxAggregatedResultBytes); } @Override @@ -70,6 +123,10 @@ public class GrpcRegisterConfig implements Serializable { + "threadpool='" + threadpool + '\'' + + ", maxAggregatedResultCount=" + + maxAggregatedResultCount + + ", maxAggregatedResultBytes=" + + maxAggregatedResultBytes + '}'; } } diff --git a/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfigTest.java b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfigTest.java index 8b4f4dcecd..07db99c55f 100644 --- a/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfigTest.java +++ b/shenyu-common/src/test/java/org/apache/shenyu/common/dto/convert/plugin/GrpcRegisterConfigTest.java @@ -33,8 +33,12 @@ public class GrpcRegisterConfigTest { public void testGetterSetter() { GrpcRegisterConfig config = new GrpcRegisterConfig(); config.setThreadpool("threadPool"); + config.setMaxAggregatedResultCount(100L); + config.setMaxAggregatedResultBytes(1024L); assertThat(config.getThreadpool(), is("threadPool")); + assertThat(config.getMaxAggregatedResultCount(), is(100L)); + assertThat(config.getMaxAggregatedResultBytes(), is(1024L)); } @Test @@ -44,8 +48,25 @@ public class GrpcRegisterConfigTest { config1.setThreadpool("threadPool"); config2.setThreadpool("threadPool"); + config1.setMaxAggregatedResultCount(100L); + config2.setMaxAggregatedResultCount(100L); + config1.setMaxAggregatedResultBytes(1024L); + config2.setMaxAggregatedResultBytes(1024L); assertThat(ImmutableSet.of(config1, config2), hasSize(1)); } + @Test + public void testNotEqualsWhenBudgetDiffers() { + GrpcRegisterConfig config1 = new GrpcRegisterConfig(); + GrpcRegisterConfig config2 = new GrpcRegisterConfig(); + + config1.setThreadpool("threadPool"); + config2.setThreadpool("threadPool"); + config1.setMaxAggregatedResultCount(100L); + config2.setMaxAggregatedResultCount(200L); + + assertThat(ImmutableSet.of(config1, config2), hasSize(2)); + } + } diff --git a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/client/ShenyuGrpcClient.java b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/client/ShenyuGrpcClient.java index 52fb41a8c5..ac1701bd60 100644 --- a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/client/ShenyuGrpcClient.java +++ b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/client/ShenyuGrpcClient.java @@ -30,7 +30,9 @@ import io.grpc.stub.StreamObserver; import org.apache.commons.collections4.CollectionUtils; import org.apache.commons.lang3.StringUtils; import org.apache.shenyu.common.dto.MetaData; +import org.apache.shenyu.common.dto.convert.plugin.GrpcRegisterConfig; import org.apache.shenyu.common.utils.GsonUtils; +import org.apache.shenyu.common.utils.Singleton; import org.apache.shenyu.plugin.grpc.exception.ShenyuGrpcException; import org.apache.shenyu.plugin.grpc.proto.CompleteObserver; import org.apache.shenyu.plugin.grpc.proto.MessageWriter; @@ -44,7 +46,10 @@ import org.slf4j.LoggerFactory; import java.io.Closeable; import java.util.Collections; import java.util.List; +import java.util.Objects; +import java.util.Optional; import java.util.concurrent.CompletableFuture; +import java.util.function.Function; import static io.grpc.stub.ClientCalls.asyncServerStreamingCall; import static io.grpc.stub.ClientCalls.asyncUnaryCall; @@ -92,7 +97,8 @@ public class ShenyuGrpcClient implements Closeable { jsonResponse); ShenyuGrpcResponse shenyuGrpcResponse = new ShenyuGrpcResponse(); - StreamObserver<DynamicMessage> streamObserver = MessageWriter.newInstance(shenyuGrpcResponse); + StreamObserver<DynamicMessage> streamObserver = MessageWriter.newInstance(shenyuGrpcResponse, + resolveMaxAggregatedResultCount(), resolveMaxAggregatedResultBytes()); ShenyuGrpcCallRequest callParams = new ShenyuGrpcCallRequest(); callParams.setMethodDescriptor(jsonMarshallerMethodDescriptor); @@ -136,22 +142,30 @@ public class ShenyuGrpcClient implements Closeable { StreamObserver<DynamicMessage> responseObserver = callParams.getResponseObserver(); CompleteObserver<DynamicMessage> doneObserver = new CompleteObserver<>(); StreamObserver<DynamicMessage> compositeObserver = CompositeStreamObserver.of(responseObserver, doneObserver); - + + ClientCall<DynamicMessage, DynamicMessage> call = createCall(callParams); + if (responseObserver instanceof MessageWriter) { + // allow the aggregator to cancel the RPC when a response budget is exhausted. + @SuppressWarnings("unchecked") + MessageWriter<DynamicMessage> writer = (MessageWriter<DynamicMessage>) responseObserver; + writer.bindCall(call); + } + StreamObserver<DynamicMessage> requestObserver; switch (methodType) { case UNARY: - asyncUnaryCall(createCall(callParams), requestList.get(0), compositeObserver); + asyncUnaryCall(call, requestList.get(0), compositeObserver); return doneObserver.getCompletionFuture(); case SERVER_STREAMING: - asyncServerStreamingCall(createCall(callParams), requestList.get(0), compositeObserver); + asyncServerStreamingCall(call, requestList.get(0), compositeObserver); return doneObserver.getCompletionFuture(); case CLIENT_STREAMING: - requestObserver = asyncClientStreamingCall(createCall(callParams), compositeObserver); + requestObserver = asyncClientStreamingCall(call, compositeObserver); requestList.forEach(requestObserver::onNext); requestObserver.onCompleted(); return doneObserver.getCompletionFuture(); case BIDI_STREAMING: - requestObserver = asyncBidiStreamingCall(createCall(callParams), compositeObserver); + requestObserver = asyncBidiStreamingCall(call, compositeObserver); requestList.forEach(requestObserver::onNext); requestObserver.onCompleted(); return doneObserver.getCompletionFuture(); @@ -160,6 +174,22 @@ public class ShenyuGrpcClient implements Closeable { return null; } } + + private long resolveMaxAggregatedResultCount() { + return resolveBudget(config -> config.getMaxAggregatedResultCount()); + } + + private long resolveMaxAggregatedResultBytes() { + return resolveBudget(config -> config.getMaxAggregatedResultBytes()); + } + + private long resolveBudget(final Function<GrpcRegisterConfig, Long> getter) { + GrpcRegisterConfig config = Singleton.INST.get(GrpcRegisterConfig.class); + if (Objects.isNull(config)) { + return 0L; + } + return Optional.ofNullable(getter.apply(config)).orElse(0L); + } @Override public void close() { diff --git a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/proto/MessageWriter.java b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/proto/MessageWriter.java index 1a1bf7a871..73d7f49c19 100644 --- a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/proto/MessageWriter.java +++ b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/main/java/org/apache/shenyu/plugin/grpc/proto/MessageWriter.java @@ -17,9 +17,13 @@ package org.apache.shenyu.plugin.grpc.proto; +import java.nio.charset.StandardCharsets; import java.util.HashMap; +import java.util.Objects; +import java.util.concurrent.atomic.AtomicReference; import org.apache.commons.lang3.StringUtils; +import org.apache.shenyu.plugin.grpc.exception.ShenyuGrpcException; import org.apache.shenyu.protocol.grpc.message.JsonMessage; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -29,10 +33,25 @@ import com.google.gson.reflect.TypeToken; import com.google.protobuf.DynamicMessage; import com.google.protobuf.Message; +import io.grpc.ClientCall; import io.grpc.stub.StreamObserver; /** * MessageWriter. + * + * <p>Aggregates the response items of a single gRPC call into {@link ShenyuGrpcResponse}. + * Optionally bounds the retained aggregate by an item-count budget and a payload-byte budget. + * When a budget is exhausted the underlying RPC is cancelled, the retained results are released + * and the call future completes exceptionally, so a partial aggregate is never returned as a + * success-shaped result.</p> + * + * <p>The budgets are opt-in: a non-positive limit means unlimited and preserves the historical + * behavior. The byte budget measures the UTF-8 payload bytes of the retained items, which is a + * payload-size bound rather than an exact JVM heap bound.</p> + * + * <p>A {@code StreamObserver} receives serialized callbacks for one call, and each call owns its + * own {@code MessageWriter} and {@link ShenyuGrpcResponse}; therefore the mutable counters here + * need no locking and concurrent RPCs are naturally isolated.</p> */ public final class MessageWriter<T extends Message> implements StreamObserver<T> { @@ -42,24 +61,69 @@ public final class MessageWriter<T extends Message> implements StreamObserver<T> private final ShenyuGrpcResponse grpcResponse; - private MessageWriter(final ShenyuGrpcResponse grpcResponse) { + private final long maxResultCount; + + private final long maxResultBytes; + + private final AtomicReference<ClientCall<?, ?>> callRef = new AtomicReference<>(); + + private long currentCount; + + private long currentBytes; + + private boolean overflowed; + + private MessageWriter(final ShenyuGrpcResponse grpcResponse, final long maxResultCount, final long maxResultBytes) { this.grpcResponse = grpcResponse; + this.maxResultCount = maxResultCount; + this.maxResultBytes = maxResultBytes; } /** - * New instance. + * New instance without any aggregate budget (unlimited), preserving the historical behavior. * * @param results results * @param <T> t * @return message message */ public static <T extends Message> MessageWriter<T> newInstance(final ShenyuGrpcResponse results) { - return new MessageWriter<>(results); + return new MessageWriter<>(results, 0L, 0L); + } + + /** + * New instance with an aggregate item-count and payload-byte budget. + * + * <p>A non-positive limit means unlimited for that dimension.</p> + * + * @param results results + * @param maxResultCount maximum retained item count, non-positive means unlimited + * @param maxResultBytes maximum retained payload bytes, non-positive means unlimited + * @param <T> t + * @return message message + */ + public static <T extends Message> MessageWriter<T> newInstance(final ShenyuGrpcResponse results, + final long maxResultCount, + final long maxResultBytes) { + return new MessageWriter<>(results, maxResultCount, maxResultBytes); + } + + /** + * Bind the underlying call so that it can be cancelled when a budget is exhausted. + * + * @param call the client call + */ + public void bindCall(final ClientCall<?, ?> call) { + this.callRef.set(call); } @Override public void onNext(final T value) { + if (overflowed) { + return; + } String respData = JsonMessage.getDataFromDynamicMessage((DynamicMessage) value); + long payloadBytes; + boolean retained; if (StringUtils.isNotBlank(respData)) { respData = respData.trim(); if (StringUtils.startsWith(respData, "{") && StringUtils.endsWith(respData, "}")) { @@ -67,9 +131,25 @@ public final class MessageWriter<T extends Message> implements StreamObserver<T> grpcResponse.getResults().add(GSON.fromJson(respData, new TypeToken<HashMap<String, Object>>() { }.getType())); + payloadBytes = respData.getBytes(StandardCharsets.UTF_8).length; + retained = true; + } else { + // non-json, non-blank payloads are not retained by the aggregate contract. + payloadBytes = 0L; + retained = false; } } else { grpcResponse.getResults().add(respData); + payloadBytes = 0L; + retained = true; + } + if (!retained) { + return; + } + currentCount++; + currentBytes += payloadBytes; + if (exceedsBudget()) { + triggerOverflow(); } } @@ -82,4 +162,24 @@ public final class MessageWriter<T extends Message> implements StreamObserver<T> public void onCompleted() { LOG.info("Messages write complete"); } + + private boolean exceedsBudget() { + return (maxResultCount > 0 && currentCount > maxResultCount) + || (maxResultBytes > 0 && currentBytes > maxResultBytes); + } + + private void triggerOverflow() { + overflowed = true; + String reason = String.format( + "gRPC aggregated response exceeded budget: retainedCount=%d (max=%d), retainedBytes=%d (max=%d)", + currentCount, maxResultCount, currentBytes, maxResultBytes); + // release retained results so the memory can be reclaimed before the RPC terminates. + grpcResponse.getResults().clear(); + currentBytes = 0L; + LOG.warn(reason); + ClientCall<?, ?> call = callRef.get(); + if (Objects.nonNull(call)) { + call.cancel(reason, new ShenyuGrpcException(reason)); + } + } } diff --git a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/test/java/org/apache/shenyu/plugin/grpc/proto/MessageWriterTest.java b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/test/java/org/apache/shenyu/plugin/grpc/proto/MessageWriterTest.java index 1c6eb99d0e..25f152e3d9 100644 --- a/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/test/java/org/apache/shenyu/plugin/grpc/proto/MessageWriterTest.java +++ b/shenyu-plugin/shenyu-plugin-proxy/shenyu-plugin-rpc/shenyu-plugin-grpc/src/test/java/org/apache/shenyu/plugin/grpc/proto/MessageWriterTest.java @@ -17,37 +17,172 @@ package org.apache.shenyu.plugin.grpc.proto; -import com.google.protobuf.Message; -import org.junit.jupiter.api.BeforeEach; +import com.google.protobuf.DynamicMessage; +import io.grpc.ClientCall; +import org.apache.shenyu.protocol.grpc.message.JsonMessage; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.junit.jupiter.MockitoExtension; +import java.nio.charset.StandardCharsets; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; /** * The Test Case For {@link MessageWriter}. */ @ExtendWith(MockitoExtension.class) public class MessageWriterTest { - - private MessageWriter<Message> messageWriter; - @BeforeEach - public void setUp() { - ShenyuGrpcResponse shenyuGrpcResponse = mock(ShenyuGrpcResponse.class); - messageWriter = MessageWriter.newInstance(shenyuGrpcResponse); + private static DynamicMessage jsonMessage(final String payload) { + return JsonMessage.buildJsonMessage(payload); + } + + private static int bytes(final String payload) { + return payload.getBytes(StandardCharsets.UTF_8).length; } @Test public void onError() { - Throwable throwable = new Throwable(); - messageWriter.onError(throwable); - + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response); + messageWriter.onError(new Throwable()); } @Test public void onCompleted() { + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response); messageWriter.onCompleted(); } + + @Test + @SuppressWarnings("unchecked") + public void unlimitedByDefaultRetainsAllResults() { + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response); + ClientCall<?, ?> call = mock(ClientCall.class); + messageWriter.bindCall(call); + + for (int i = 0; i < 100; i++) { + messageWriter.onNext(jsonMessage("{\"i\":\"" + i + "\"}")); + } + + assertEquals(100, response.getResults().size()); + verify(call, never()).cancel(anyString(), any()); + } + + @Test + @SuppressWarnings("unchecked") + public void withinCountBudgetDoesNotCancel() { + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response, 3L, 0L); + ClientCall<?, ?> call = mock(ClientCall.class); + messageWriter.bindCall(call); + + messageWriter.onNext(jsonMessage("{\"a\":\"1\"}")); + messageWriter.onNext(jsonMessage("{\"a\":\"2\"}")); + messageWriter.onNext(jsonMessage("{\"a\":\"3\"}")); + + assertEquals(3, response.getResults().size()); + verify(call, never()).cancel(anyString(), any()); + } + + @Test + @SuppressWarnings("unchecked") + public void exceedingCountBudgetCancelsAndReleasesResults() { + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response, 2L, 0L); + ClientCall<?, ?> call = mock(ClientCall.class); + messageWriter.bindCall(call); + + messageWriter.onNext(jsonMessage("{\"a\":\"1\"}")); + messageWriter.onNext(jsonMessage("{\"a\":\"2\"}")); + // this one crosses the count budget. + messageWriter.onNext(jsonMessage("{\"a\":\"3\"}")); + + assertTrue(response.getResults().isEmpty()); + verify(call, times(1)).cancel(anyString(), any()); + + // subsequent messages are ignored and do not cancel again. + messageWriter.onNext(jsonMessage("{\"a\":\"4\"}")); + assertTrue(response.getResults().isEmpty()); + verify(call, times(1)).cancel(anyString(), any()); + } + + @Test + @SuppressWarnings("unchecked") + public void atByteBudgetDoesNotCancel() { + String payload = "{\"a\":\"bb\"}"; + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response, 0L, bytes(payload)); + ClientCall<?, ?> call = mock(ClientCall.class); + messageWriter.bindCall(call); + + messageWriter.onNext(jsonMessage(payload)); + + assertEquals(1, response.getResults().size()); + verify(call, never()).cancel(anyString(), any()); + } + + @Test + @SuppressWarnings("unchecked") + public void exceedingByteBudgetCancelsAndReleasesResults() { + String payload = "{\"a\":\"bb\"}"; + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response, 0L, bytes(payload)); + ClientCall<?, ?> call = mock(ClientCall.class); + messageWriter.bindCall(call); + + messageWriter.onNext(jsonMessage(payload)); + // accumulated bytes now exceed the budget. + messageWriter.onNext(jsonMessage(payload)); + + assertTrue(response.getResults().isEmpty()); + verify(call, times(1)).cancel(anyString(), any()); + } + + @Test + @SuppressWarnings("unchecked") + public void overflowDoesNotCancelWhenCallNotBound() { + ShenyuGrpcResponse response = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> messageWriter = MessageWriter.newInstance(response, 1L, 0L); + + messageWriter.onNext(jsonMessage("{\"a\":\"1\"}")); + messageWriter.onNext(jsonMessage("{\"a\":\"2\"}")); + + // results are still released even without a bound call; no exception is thrown. + assertTrue(response.getResults().isEmpty()); + } + + @Test + @SuppressWarnings("unchecked") + public void concurrentWritersAreIsolated() { + ShenyuGrpcResponse responseA = new ShenyuGrpcResponse(); + ShenyuGrpcResponse responseB = new ShenyuGrpcResponse(); + MessageWriter<DynamicMessage> writerA = MessageWriter.newInstance(responseA, 1L, 0L); + MessageWriter<DynamicMessage> writerB = MessageWriter.newInstance(responseB, 0L, 0L); + ClientCall<?, ?> callA = mock(ClientCall.class); + ClientCall<?, ?> callB = mock(ClientCall.class); + writerA.bindCall(callA); + writerB.bindCall(callB); + + // writer A overflows, writer B keeps accumulating. + writerA.onNext(jsonMessage("{\"a\":\"1\"}")); + writerA.onNext(jsonMessage("{\"a\":\"2\"}")); + writerB.onNext(jsonMessage("{\"b\":\"1\"}")); + writerB.onNext(jsonMessage("{\"b\":\"2\"}")); + + assertTrue(responseA.getResults().isEmpty()); + verify(callA, times(1)).cancel(anyString(), any()); + assertEquals(2, responseB.getResults().size()); + verify(callB, never()).cancel(anyString(), any()); + } }
