SCB-2102 Fix the problem of unable to return due to abnormal file upl…

…oad (#2017)
apache · Oct 26, 2020 · 5616701 · 5616701
1 parent 2ea1340
commit 5616701
Show file tree

Hide file tree

Showing 7 changed files with 28 additions and 10 deletions.
diff --git a/...t/src/main/java/org/apache/servicecomb/common/rest/codec/param/RestClientRequestImpl.java b/...t/src/main/java/org/apache/servicecomb/common/rest/codec/param/RestClientRequestImpl.java
@@ -44,6 +44,7 @@
 import com.google.common.collect.Multimap;
 
 import io.vertx.core.Context;
+import io.vertx.core.Handler;
 import io.vertx.core.MultiMap;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.HttpClientRequest;
@@ -66,10 +67,18 @@ public class RestClientRequestImpl implements RestClientRequest {
 
   protected Buffer bodyBuffer;
 
+  private Handler<Throwable> throwableHandler;
+
   public RestClientRequestImpl(HttpClientRequest request, Context context, AsyncResponse asyncResp) {
+    this(request, context, asyncResp, null);
+  }
+
+  public RestClientRequestImpl(HttpClientRequest request, Context context, AsyncResponse asyncResp,
+                               Handler<Throwable> throwableHandler) {
     this.context = context;
     this.asyncResp = asyncResp;
     this.request = request;
+    this.throwableHandler = throwableHandler;
   }
 
   @Override
@@ -186,7 +195,7 @@ private void attachFile(String boundary, Iterator<Entry<String, Part>> uploadsIt
     Buffer fileHeader = fileBoundaryInfo(boundary, name, part);
     request.write(fileHeader);
 
-    new PumpFromPart(context, part).toWriteStream(request).whenComplete((v, e) -> {
+    new PumpFromPart(context, part).toWriteStream(request, throwableHandler).whenComplete((v, e) -> {
       if (e != null) {
         LOGGER.debug("Failed to sending file [{}:{}].", name, filename, e);
         asyncResp.consumerFail(e);

diff --git a/...tion-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/ReadStreamPart.java b/...tion-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/http/ReadStreamPart.java
@@ -91,7 +91,7 @@ public ReadStream<Buffer> getReadStream() {
    * so the return future only means finished read from readStream.
    */
   public CompletableFuture<Void> saveToWriteStream(WriteStream<Buffer> writeStream) {
-    return new PumpCommon().pump(context, readStream, writeStream);
+    return new PumpCommon().pump(context, readStream, writeStream, null);
   }
 
   public CompletableFuture<byte[]> saveAsBytes() {

diff --git a/...rg/apache/servicecomb/foundation/vertx/http/VertxServerResponseToHttpServletResponse.java b/...rg/apache/servicecomb/foundation/vertx/http/VertxServerResponseToHttpServletResponse.java
@@ -126,7 +126,7 @@ public void internalFlushBuffer() {
   public CompletableFuture<Void> sendPart(Part part) {
     DownloadUtils.prepareDownloadHeader(this, part);
 
-    return new PumpFromPart(context, part).toWriteStream(serverResponse);
+    return new PumpFromPart(context, part).toWriteStream(serverResponse, null);
   }
 
   @Override

diff --git a/...dation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/stream/PumpCommon.java b/...dation-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/stream/PumpCommon.java
@@ -21,6 +21,7 @@
 import org.apache.servicecomb.foundation.common.io.AsyncCloseable;
 
 import io.vertx.core.Context;
+import io.vertx.core.Handler;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.HttpClientResponse;
 import io.vertx.core.streams.Pump;
@@ -39,7 +40,8 @@ public class PumpCommon {
    * <p>  if writeStream is not AsyncCloseable, future only means read complete
    */
   @SuppressWarnings("unchecked")
-  public CompletableFuture<Void> pump(Context context, ReadStream<Buffer> readStream, WriteStream<Buffer> writeStream) {
+  public CompletableFuture<Void> pump(Context context, ReadStream<Buffer> readStream, WriteStream<Buffer> writeStream,
+      Handler<Throwable> throwableHandler) {
     CompletableFuture<Void> readFuture = new CompletableFuture<>();
 
     writeStream.exceptionHandler(e -> {
@@ -54,6 +56,9 @@ public CompletableFuture<Void> pump(Context context, ReadStream<Buffer> readStre
         // so can only close the connection.
         ((HttpClientResponse) readStream).request().connection().close();
       }
+      if (throwableHandler != null) {
+        throwableHandler.handle(e);
+      }
       readFuture.completeExceptionally(e);
     });
     readStream.exceptionHandler(readFuture::completeExceptionally);

diff --git a/...tion-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/stream/PumpFromPart.java b/...tion-vertx/src/main/java/org/apache/servicecomb/foundation/vertx/stream/PumpFromPart.java
@@ -28,6 +28,7 @@
 import org.apache.servicecomb.foundation.vertx.http.ReadStreamPart;
 
 import io.vertx.core.Context;
+import io.vertx.core.Handler;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.streams.ReadStream;
 import io.vertx.core.streams.WriteStream;
@@ -62,9 +63,9 @@ private CompletableFuture<ReadStream<Buffer>> prepareReadStream() {
     return future;
   }
 
-  public CompletableFuture<Void> toWriteStream(WriteStream<Buffer> writeStream) {
+  public CompletableFuture<Void> toWriteStream(WriteStream<Buffer> writeStream, Handler<Throwable> throwableHandler) {
     return prepareReadStream()
-        .thenCompose(readStream -> new PumpCommon().pump(context, readStream, writeStream))
+        .thenCompose(readStream -> new PumpCommon().pump(context, readStream, writeStream, throwableHandler))
         .whenComplete((v, e) -> DownloadUtils.clearPartResource(part));
   }
 
@@ -79,7 +80,7 @@ public CompletableFuture<Void> toOutputStream(OutputStream outputStream, boolean
   private CompletableFuture<Void> toOutputStreamAsync(OutputStream outputStream, boolean autoCloseOutputStream) {
     OutputStreamToWriteStream outputStreamToWriteStream = new OutputStreamToWriteStream(context, outputStream,
         autoCloseOutputStream);
-    return toWriteStream(outputStreamToWriteStream);
+    return toWriteStream(outputStreamToWriteStream, null);
   }
 
   // DO NOT use a mocked sync context to unify the pump logic

diff --git a/...pache/servicecomb/foundation/vertx/http/TestVertxServerResponseToHttpServletResponse.java b/...pache/servicecomb/foundation/vertx/http/TestVertxServerResponseToHttpServletResponse.java
@@ -372,7 +372,7 @@ public void sendPart_ReadStreamPart(@Mocked ReadStreamPart part) {
     CompletableFuture<Void> future = new CompletableFuture<>();
     new MockUp<PumpFromPart>() {
       @Mock
-      CompletableFuture<Void> toWriteStream(WriteStream<Buffer> writeStream) {
+      CompletableFuture<Void> toWriteStream(WriteStream<Buffer> writeStream, Handler<Throwable> throwableHandler) {
         return future;
       }
     };

diff --git a/...src/main/java/org/apache/servicecomb/transport/rest/client/http/RestClientInvocation.java b/...src/main/java/org/apache/servicecomb/transport/rest/client/http/RestClientInvocation.java
@@ -47,6 +47,7 @@
 import org.slf4j.LoggerFactory;
 import org.springframework.util.StringUtils;
 
+import io.vertx.core.Handler;
 import io.vertx.core.buffer.Buffer;
 import io.vertx.core.http.HttpClientRequest;
 import io.vertx.core.http.HttpClientResponse;
@@ -76,6 +77,8 @@ public class RestClientInvocation {
 
   private HttpClientResponse clientResponse;
 
+  private Handler<Throwable> throwableHandler = e -> fail((ConnectionBase) clientRequest.connection(), e);
+
   public RestClientInvocation(HttpClientWithContext httpClientWithContext, List<HttpClientFilter> httpClientFilters) {
     this.httpClientWithContext = httpClientWithContext;
     this.httpClientFilters = httpClientFilters;
@@ -94,7 +97,7 @@ public void invoke(Invocation invocation, AsyncResponse asyncResp) throws Except
     createRequest(ipPort, path);
     clientRequest.putHeader(org.apache.servicecomb.core.Const.TARGET_MICROSERVICE, invocation.getMicroserviceName());
     RestClientRequestImpl restClientRequest =
-        new RestClientRequestImpl(clientRequest, httpClientWithContext.context(), asyncResp);
+        new RestClientRequestImpl(clientRequest, httpClientWithContext.context(), asyncResp, throwableHandler);
     invocation.getHandlerContext().put(RestConst.INVOCATION_HANDLER_REQUESTCLIENT, restClientRequest);
 
     Buffer requestBodyBuffer = restClientRequest.getBodyBuffer();
@@ -109,7 +112,7 @@ public void invoke(Invocation invocation, AsyncResponse asyncResp) throws Except
     clientRequest.exceptionHandler(e -> {
       LOGGER.error(invocation.getMarker(), "Failed to send request, local:{}, remote:{}.",
           getLocalAddress(), ipPort.getSocketAddress(), e);
-      fail((ConnectionBase) clientRequest.connection(), e);
+      throwableHandler.handle(e);
     });
     clientRequest.connectionHandler(connection -> {
       LOGGER.debug("http connection connected, local:{}, remote:{}.",