diff --git a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannel.java b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannel.java index 0135818fb3b..292c9b2f2ff 100644 --- a/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannel.java +++ b/proxy/src/main/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannel.java @@ -51,6 +51,7 @@ import org.apache.rocketmq.proxy.service.relay.ProxyRelayResult; import org.apache.rocketmq.proxy.service.relay.ProxyRelayService; import org.apache.rocketmq.proxy.service.transaction.TransactionData; +import org.apache.rocketmq.remoting.protocol.ResponseCode; import org.apache.rocketmq.remoting.protocol.RemotingCommand; import org.apache.rocketmq.remoting.protocol.body.ConsumeMessageDirectlyResult; import org.apache.rocketmq.remoting.protocol.body.ConsumerRunningInfo; @@ -235,11 +236,15 @@ protected CompletableFuture processGetConsumerRunningInfo(RemotingCommand if (Objects.isNull(header) || !header.isJstackEnable()) { return CompletableFuture.completedFuture(null); } - this.writeTelemetryCommand(TelemetryCommand.newBuilder() + String nonce = this.grpcChannelManager.addResponseFuture(responseFuture); + boolean written = this.tryWriteTelemetryCommand(TelemetryCommand.newBuilder() .setPrintThreadStackTraceCommand(PrintThreadStackTraceCommand.newBuilder() - .setNonce(this.grpcChannelManager.addResponseFuture(responseFuture)) + .setNonce(nonce) .build()) .build()); + if (!written) { + this.completeResponseFutureOnWriteFailure(nonce, responseFuture); + } return CompletableFuture.completedFuture(null); } @@ -247,12 +252,16 @@ protected CompletableFuture processGetConsumerRunningInfo(RemotingCommand protected CompletableFuture processConsumeMessageDirectly(RemotingCommand command, ConsumeMessageDirectlyResultRequestHeader header, MessageExt messageExt, CompletableFuture> responseFuture) { - this.writeTelemetryCommand(TelemetryCommand.newBuilder() + String nonce = this.grpcChannelManager.addResponseFuture(responseFuture); + boolean written = this.tryWriteTelemetryCommand(TelemetryCommand.newBuilder() .setVerifyMessageCommand(VerifyMessageCommand.newBuilder() - .setNonce(this.grpcChannelManager.addResponseFuture(responseFuture)) + .setNonce(nonce) .setMessage(GrpcConverter.getInstance().buildMessage(messageExt)) .build()) .build()); + if (!written) { + this.completeResponseFutureOnWriteFailure(nonce, responseFuture); + } return CompletableFuture.completedFuture(null); } @@ -261,26 +270,41 @@ public String getClientId() { } public void writeTelemetryCommand(TelemetryCommand command) { + this.tryWriteTelemetryCommand(command); + } + + private boolean tryWriteTelemetryCommand(TelemetryCommand command) { StreamObserver observer = this.telemetryCommandRef.get(); if (observer == null) { log.warn("telemetry command observer is null when try to write data. command:{}, channel:{}", TextFormat.shortDebugString(command), this); - return; + return false; } synchronized (this.telemetryWriteLock) { observer = this.telemetryCommandRef.get(); if (observer == null) { log.warn("telemetry command observer is null when try to write data. command:{}, channel:{}", TextFormat.shortDebugString(command), this); - return; + return false; } try { observer.onNext(command); + return true; } catch (StatusRuntimeException | IllegalStateException exception) { log.warn("write telemetry failed. command:{}", command, exception); this.clearClientObserver(observer); + return false; } } } + private void completeResponseFutureOnWriteFailure(String nonce, + CompletableFuture> responseFuture) { + CompletableFuture> registeredFuture = + this.grpcChannelManager.getAndRemoveResponseFuture(nonce); + CompletableFuture> future = + registeredFuture == null ? responseFuture : registeredFuture; + future.complete(new ProxyRelayResult<>(ResponseCode.SYSTEM_BUSY, "write telemetry command failed", null)); + } + @Override public String toString() { return MoreObjects.toStringHelper(this) diff --git a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannelTest.java b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannelTest.java index 1bdbdd9befe..a88e385b557 100644 --- a/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannelTest.java +++ b/proxy/src/test/java/org/apache/rocketmq/proxy/grpc/v2/channel/GrpcClientChannelTest.java @@ -20,14 +20,26 @@ import apache.rocketmq.v2.Publishing; import apache.rocketmq.v2.Resource; import apache.rocketmq.v2.Settings; +import apache.rocketmq.v2.TelemetryCommand; +import io.grpc.stub.StreamObserver; +import java.nio.charset.StandardCharsets; +import java.util.concurrent.CompletableFuture; import org.apache.commons.lang3.RandomStringUtils; +import org.apache.rocketmq.common.message.MessageExt; import org.apache.rocketmq.proxy.common.ProxyContext; import org.apache.rocketmq.proxy.config.InitConfigTest; import org.apache.rocketmq.proxy.grpc.v2.common.GrpcClientSettingsManager; import org.apache.rocketmq.proxy.processor.channel.ChannelProtocolType; import org.apache.rocketmq.proxy.processor.channel.RemoteChannel; import org.apache.rocketmq.proxy.remoting.channel.RemotingChannel; +import org.apache.rocketmq.proxy.service.relay.ProxyRelayResult; import org.apache.rocketmq.proxy.service.relay.ProxyRelayService; +import org.apache.rocketmq.remoting.protocol.RemotingCommand; +import org.apache.rocketmq.remoting.protocol.ResponseCode; +import org.apache.rocketmq.remoting.protocol.body.ConsumeMessageDirectlyResult; +import org.apache.rocketmq.remoting.protocol.body.ConsumerRunningInfo; +import org.apache.rocketmq.remoting.protocol.header.ConsumeMessageDirectlyResultRequestHeader; +import org.apache.rocketmq.remoting.protocol.header.GetConsumerRunningInfoRequestHeader; import org.junit.Before; import org.junit.Test; import org.junit.runner.RunWith; @@ -35,9 +47,14 @@ import org.mockito.junit.MockitoJUnitRunner; import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @RunWith(MockitoJUnitRunner.class) @@ -79,4 +96,58 @@ public void testChannelExtendAttributeParse() { assertEquals(clientSettings, GrpcClientChannel.parseChannelExtendAttribute(this.grpcClientChannel)); assertNull(GrpcClientChannel.parseChannelExtendAttribute(mock(RemotingChannel.class))); } -} \ No newline at end of file + + @Test + public void testGetConsumerRunningInfoShouldFailFastWhenObserverIsMissing() throws Exception { + CompletableFuture> responseFuture = new CompletableFuture<>(); + when(grpcChannelManager.addResponseFuture(eq(responseFuture))).thenReturn("nonce-1"); + when(grpcChannelManager.getAndRemoveResponseFuture(eq("nonce-1"))).thenReturn(responseFuture); + + GetConsumerRunningInfoRequestHeader header = new GetConsumerRunningInfoRequestHeader(); + header.setJstackEnable(true); + + grpcClientChannel.processGetConsumerRunningInfo(mock(RemotingCommand.class), header, responseFuture).get(); + + assertTrue(responseFuture.isDone()); + ProxyRelayResult result = responseFuture.get(); + assertEquals(ResponseCode.SYSTEM_BUSY, result.getCode()); + assertEquals("write telemetry command failed", result.getRemark()); + verify(grpcChannelManager).getAndRemoveResponseFuture("nonce-1"); + } + + @Test + public void testConsumeMessageDirectlyShouldFailFastWhenObserverWriteFails() throws Exception { + StreamObserver observer = mock(StreamObserver.class); + doThrow(new IllegalStateException("stream closed")).when(observer).onNext(any(TelemetryCommand.class)); + grpcClientChannel.setClientObserver(observer); + + CompletableFuture> responseFuture = new CompletableFuture<>(); + when(grpcChannelManager.addResponseFuture(eq(responseFuture))).thenReturn("nonce-2"); + when(grpcChannelManager.getAndRemoveResponseFuture(eq("nonce-2"))).thenReturn(responseFuture); + + grpcClientChannel.processConsumeMessageDirectly( + mock(RemotingCommand.class), + new ConsumeMessageDirectlyResultRequestHeader(), + buildMessageExt(), + responseFuture + ).get(); + + assertTrue(responseFuture.isDone()); + ProxyRelayResult result = responseFuture.get(); + assertEquals(ResponseCode.SYSTEM_BUSY, result.getCode()); + assertEquals("write telemetry command failed", result.getRemark()); + verify(grpcChannelManager).getAndRemoveResponseFuture("nonce-2"); + assertFalse(grpcClientChannel.isOpen()); + } + + private MessageExt buildMessageExt() { + MessageExt messageExt = new MessageExt(); + messageExt.setTopic("test-topic"); + messageExt.setBody("hello".getBytes(StandardCharsets.UTF_8)); + messageExt.setMsgId("msg-id"); + messageExt.setBornTimestamp(System.currentTimeMillis()); + messageExt.setStoreTimestamp(System.currentTimeMillis()); + messageExt.putUserProperty("test", "true"); + return messageExt; + } +}