diff --git a/trpc-proto/trpc-proto-http/src/main/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutor.java b/trpc-proto/trpc-proto-http/src/main/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutor.java index 73eeb11f9..f4a39dab8 100644 --- a/trpc-proto/trpc-proto-http/src/main/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutor.java +++ b/trpc-proto/trpc-proto-http/src/main/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutor.java @@ -38,11 +38,15 @@ import com.tencent.trpc.proto.http.common.RpcServerContextWithHttp; import com.tencent.trpc.proto.http.common.TrpcServletRequestWrapper; import com.tencent.trpc.proto.http.common.TrpcServletResponseWrapper; +import java.io.IOException; import java.lang.reflect.Type; import java.net.InetSocketAddress; import java.nio.charset.StandardCharsets; import java.util.Enumeration; import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.CompletionStage; import java.util.concurrent.CountDownLatch; import java.util.concurrent.TimeUnit; @@ -65,15 +69,15 @@ public abstract class AbstractHttpExecutor { protected void execute(HttpServletRequest request, HttpServletResponse response, RpcMethodInfoAndInvoker methodInfoAndInvoker) { - + AtomicBoolean responded = new AtomicBoolean(false); try { DefRequest rpcRequest = buildDefRequest(request, response, methodInfoAndInvoker); - CountDownLatch countDownLatch = new CountDownLatch(1); + CompletableFuture completionFuture = new CompletableFuture<>(); // use a thread pool for asynchronous processing - invokeRpcRequest(methodInfoAndInvoker.getInvoker(), rpcRequest, countDownLatch); + invokeRpcRequest(methodInfoAndInvoker.getInvoker(), rpcRequest, completionFuture, responded); // If the request carries a timeout, use this timeout to wait for the request to be processed. // If not carried, use the default timeout. @@ -81,18 +85,25 @@ protected void execute(HttpServletRequest request, HttpServletResponse response, if (requestTimeout <= 0) { requestTimeout = methodInfoAndInvoker.getInvoker().getConfig().getRequestTimeout(); } - if (requestTimeout > 0 && !countDownLatch.await(requestTimeout, TimeUnit.MILLISECONDS)) { - throw TRpcException.newFrameException(ErrorCode.TRPC_SERVER_TIMEOUT_ERR, - "wait http request execute timeout"); + if (requestTimeout > 0) { + try { + completionFuture.get(requestTimeout, TimeUnit.MILLISECONDS); + } catch (TimeoutException ex) { + if (responded.compareAndSet(false, true)) { + doErrorReply(request, response, + TRpcException.newFrameException(ErrorCode.TRPC_SERVER_TIMEOUT_ERR, + "wait http request execute timeout")); + } + } } else { - countDownLatch.await(); + completionFuture.get(); } - } catch (Exception ex) { logger.error("dispatch request [{}] error", request, ex); - doErrorReply(request, response, ex); + if (responded.compareAndSet(false, true)) { + doErrorReply(request, response, ex); + } } - } /** @@ -107,55 +118,83 @@ protected void execute(HttpServletRequest request, HttpServletResponse response, /** * Request processing * - * @param countDownLatch latch used to wait for the request processing + * @param invoker the invoker + * @param rpcRequest the rpc request + * @param completionFuture the completion future + * @param responded the responded flag */ - private void invokeRpcRequest(ProviderInvoker invoker, DefRequest rpcRequest, CountDownLatch countDownLatch) { + private void invokeRpcRequest(ProviderInvoker invoker, DefRequest rpcRequest, + CompletableFuture completionFuture, + AtomicBoolean responded) { WorkerPool workerPool = invoker.getConfig().getWorkerPoolObj(); if (null == workerPool) { logger.error("dispatch rpcRequest [{}] error, workerPool is empty", rpcRequest); - throw TRpcException.newFrameException(ErrorCode.TRPC_SERVER_NOSERVICE_ERR, - "not found service, workerPool is empty"); + completionFuture.completeExceptionally(TRpcException.newFrameException(ErrorCode.TRPC_SERVER_NOSERVICE_ERR, + "not found service, workerPool is empty")); + return; } workerPool.execute(() -> { - - // Get the original http response - HttpServletResponse response = getOriginalResponse(rpcRequest); - - // Invoke the routing implementation method to handle the request. - CompletionStage future = invoker.invoke(rpcRequest); - future.whenComplete((result, t) -> { - try { - // Throw the call exception, which will be handled uniformly by the exception handling program. - if (t != null) { - throw t; - } - - // Throw a business logic exception, which will be handled uniformly - // by the exception handling program. - Throwable ex = result.getException(); - if (ex != null) { - throw ex; + try { + // Get the original http response + HttpServletResponse response = getOriginalResponse(rpcRequest); + // Invoke the routing implementation method to handle the request. + CompletionStage rpcFuture = invoker.invoke(rpcRequest); + + rpcFuture.whenComplete((result, throwable) -> { + try { + if (responded.get()) { + return; + } + + // Throw the call exception, which will be handled uniformly by the exception handling program. + if (throwable != null) { + throw throwable; + } + + // Throw a business logic exception, which will be handled uniformly + // by the exception handling program. + if (result.getException() != null) { + throw result.getException(); + } + + // normal response + if (responded.compareAndSet(false, true)) { + response.setStatus(HttpStatus.SC_OK); + httpCodec.writeHttpResponse(response, result); + response.flushBuffer(); + } + + completionFuture.complete(null); + } catch (Throwable t) { + handleError(t, rpcRequest, response, responded, completionFuture); } + }); - // normal response - response.setStatus(HttpStatus.SC_OK); - httpCodec.writeHttpResponse(response, result); - response.flushBuffer(); - } catch (Throwable e) { - HttpServletRequest request = getOriginalRequest(rpcRequest); - logger.warn("reply message error, channel: [{}], msg:[{}]", request.getRemoteAddr(), request, e); - httpErrorReply(request, response, - ErrorResponse.create(request, HttpStatus.SC_SERVICE_UNAVAILABLE, e)); - } finally { - countDownLatch.countDown(); - } - }); + } catch (Exception e) { + handleError(e, rpcRequest, getOriginalResponse(rpcRequest), responded, completionFuture); + } }); } + /** + * Handle error + */ + private void handleError(Throwable t, DefRequest rpcRequest, HttpServletResponse response, + AtomicBoolean responded, CompletableFuture completionFuture) { + try { + if (responded.compareAndSet(false, true)) { + HttpServletRequest request = getOriginalRequest(rpcRequest); + logger.warn("reply message error, channel: [{}], msg:[{}]", request.getRemoteAddr(), request, t); + httpErrorReply(request, response, ErrorResponse.create(request, HttpStatus.SC_SERVICE_UNAVAILABLE, t)); + } + } finally { + completionFuture.completeExceptionally(t); + } + } + /** * Build the context request. * diff --git a/trpc-proto/trpc-proto-http/src/test/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutorTest.java b/trpc-proto/trpc-proto-http/src/test/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutorTest.java index 2abbe5643..2a4e1926a 100644 --- a/trpc-proto/trpc-proto-http/src/test/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutorTest.java +++ b/trpc-proto/trpc-proto-http/src/test/java/com/tencent/trpc/proto/http/server/AbstractHttpExecutorTest.java @@ -13,16 +13,32 @@ import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; +import static org.mockito.Matchers.any; +import static org.mockito.Matchers.eq; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; import static org.powermock.api.mockito.PowerMockito.doReturn; import static org.powermock.api.mockito.PowerMockito.mock; import static org.powermock.api.mockito.PowerMockito.when; import com.tencent.trpc.core.rpc.RpcInvocation; import com.tencent.trpc.core.rpc.common.RpcMethodInfo; +import com.tencent.trpc.core.rpc.def.DefRequest; +import com.tencent.trpc.proto.http.common.ErrorResponse; import com.tencent.trpc.proto.http.common.HttpConstants; +import java.io.IOException; +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.atomic.AtomicBoolean; import javax.servlet.http.HttpServletRequest; +import javax.servlet.http.HttpServletResponse; import org.junit.Test; import org.junit.runner.RunWith; +import org.powermock.api.mockito.PowerMockito; import org.powermock.core.classloader.annotations.PowerMockIgnore; import org.powermock.core.classloader.annotations.PrepareForTest; import org.powermock.modules.junit4.PowerMockRunner; @@ -48,4 +64,118 @@ public void buildRpcInvocation_shouldSuccess() throws Exception { methodInfo); assertEquals(rpcInvocation.getFunc(), "/trpc.demo.server/hello"); } + + @Test + public void handleError_shouldCallHttpErrorReply_whenRespondedIsFalse() throws Exception { + // Arrange + HttpServletRequest request = mock(HttpServletRequest.class); + HttpServletResponse response = mock(HttpServletResponse.class); + DefRequest rpcRequest = mock(DefRequest.class); + AtomicBoolean responded = new AtomicBoolean(false); + CompletableFuture completionFuture = new CompletableFuture<>(); + Throwable throwable = new RuntimeException("test error"); + + AbstractHttpExecutor abstractHttpExecutor = mock(AbstractHttpExecutor.class); + + // Mock getOriginalRequest method + Map attachments = new HashMap<>(); + attachments.put(HttpConstants.TRPC_ATTACH_SERVLET_REQUEST, request); + when(rpcRequest.getAttachments()).thenReturn(attachments); + when(request.getRemoteAddr()).thenReturn("127.0.0.1"); + + // Mock httpErrorReply method + PowerMockito.doNothing().when(abstractHttpExecutor, "httpErrorReply", any(HttpServletRequest.class), + any(HttpServletResponse.class), any(ErrorResponse.class)); + + // Call real method + when(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, responded, completionFuture) + .thenCallRealMethod(); + + // Act + Whitebox.invokeMethod(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, + responded, completionFuture); + + // Assert + assertTrue("responded should be set to true", responded.get()); + assertTrue("completionFuture should be completed exceptionally", completionFuture.isCompletedExceptionally()); + + // Verify httpErrorReply was called + PowerMockito.verifyPrivate(abstractHttpExecutor, times(1)) + .invoke("httpErrorReply", eq(request), eq(response), any(ErrorResponse.class)); + } + + @Test + public void handleError_shouldNotCallHttpErrorReply_whenRespondedIsTrue() throws Exception { + // Arrange + HttpServletRequest request = mock(HttpServletRequest.class); + HttpServletResponse response = mock(HttpServletResponse.class); + DefRequest rpcRequest = mock(DefRequest.class); + AtomicBoolean responded = new AtomicBoolean(true); // Already responded + CompletableFuture completionFuture = new CompletableFuture<>(); + Throwable throwable = new RuntimeException("test error"); + + AbstractHttpExecutor abstractHttpExecutor = mock(AbstractHttpExecutor.class); + + // Mock httpErrorReply method + PowerMockito.doNothing().when(abstractHttpExecutor, "httpErrorReply", any(HttpServletRequest.class), + any(HttpServletResponse.class), any(ErrorResponse.class)); + + // Call real method + when(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, responded, completionFuture) + .thenCallRealMethod(); + + // Act + Whitebox.invokeMethod(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, + responded, completionFuture); + + // Assert + assertTrue("responded should remain true", responded.get()); + assertTrue("completionFuture should be completed exceptionally", completionFuture.isCompletedExceptionally()); + + // Verify httpErrorReply was NOT called + PowerMockito.verifyPrivate(abstractHttpExecutor, never()) + .invoke("httpErrorReply", any(HttpServletRequest.class), any(HttpServletResponse.class), + any(ErrorResponse.class)); + } + + @Test + public void handleError_shouldCompleteExceptionally_evenWhenHttpErrorReplyThrows() throws Exception { + // Arrange + HttpServletRequest request = mock(HttpServletRequest.class); + HttpServletResponse response = mock(HttpServletResponse.class); + DefRequest rpcRequest = mock(DefRequest.class); + AtomicBoolean responded = new AtomicBoolean(false); + CompletableFuture completionFuture = new CompletableFuture<>(); + Throwable throwable = new RuntimeException("test error"); + + AbstractHttpExecutor abstractHttpExecutor = mock(AbstractHttpExecutor.class); + + // Mock getOriginalRequest method + Map attachments = new HashMap<>(); + attachments.put(HttpConstants.TRPC_ATTACH_SERVLET_REQUEST, request); + when(rpcRequest.getAttachments()).thenReturn(attachments); + when(request.getRemoteAddr()).thenReturn("127.0.0.1"); + + // Mock httpErrorReply method to throw exception + PowerMockito.doThrow(new IOException("response error")).when(abstractHttpExecutor, "httpErrorReply", + any(HttpServletRequest.class), any(HttpServletResponse.class), any(ErrorResponse.class)); + + // Call real method + when(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, responded, completionFuture) + .thenCallRealMethod(); + + Whitebox.invokeMethod(abstractHttpExecutor, "handleError", throwable, rpcRequest, response, + responded, completionFuture); + + assertTrue("responded should be set to true", responded.get()); + assertTrue("completionFuture should be completed exceptionally", completionFuture.isCompletedExceptionally()); + + try { + completionFuture.get(); + fail("Should have thrown exception"); + } catch (ExecutionException e) { + assertEquals("Should complete with original throwable", throwable, e.getCause()); + } + } + } \ No newline at end of file