Merge pull request #300 from databendlabs/feat/enhance-stream-upload

feat: enhance stream upload

Author: hantmac

databend-jdbc/src/main/java/com/databend/jdbc/cloud/DatabendPresignClientV1.java

@@ -5,6 +5,7 @@
 import okhttp3.MediaType;
 import okhttp3.MultipartBody;
 import okhttp3.OkHttpClient;
+import okhttp3.Protocol;
 import okhttp3.Request;
 import okhttp3.RequestBody;
 import okhttp3.Response;
@@ -15,12 +16,14 @@
 import org.checkerframework.checker.nullness.qual.NonNull;
 
 import javax.annotation.Nullable;
+
 import java.io.File;
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.SocketTimeoutException;
 import java.nio.file.Files;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.concurrent.TimeUnit;
 import java.util.logging.Level;
 import java.util.logging.Logger;
@@ -31,7 +34,9 @@
 import static java.util.Objects.requireNonNull;
 import static java.util.concurrent.TimeUnit.MILLISECONDS;
 
-public class DatabendPresignClientV1 implements DatabendPresignClient {
+public class DatabendPresignClientV1
+        implements DatabendPresignClient
+{
 
     private static final int MaxRetryAttempts = 20;
 
@@ -40,15 +45,52 @@ public class DatabendPresignClientV1 implements DatabendPresignClient {
     private final String uri;
     private static final Logger logger = Logger.getLogger(DatabendPresignClientV1.class.getPackage().getName());
 
-    public DatabendPresignClientV1(OkHttpClient client, String uri) {
+    public DatabendPresignClientV1(OkHttpClient client, String uri)
+    {
         Logger.getLogger(OkHttpClient.class.getName()).setLevel(Level.FINEST);
         OkHttpClient.Builder builder = client.newBuilder();
         this.client = builder.
-                connectTimeout(120, TimeUnit.SECONDS).build();
+                connectTimeout(600, TimeUnit.SECONDS)
+                .writeTimeout(900, TimeUnit.SECONDS)
+                .readTimeout(600, TimeUnit.SECONDS)
+                .retryOnConnectionFailure(true)
+                .protocols(Arrays.asList(Protocol.HTTP_1_1))
+                .addInterceptor(chain -> {
+                    Request request = chain.request();
+                    int retryCount = 0;
+                    Response response = null;
+                    while (retryCount < 3) {
+                        try {
+                            response = chain.proceed(request);
+                            if (response.isSuccessful()) {
+                                return response;
+                            }
+                            response.close();
+                        }
+                        catch (IOException e) {
+                            if (retryCount == 2) {
+                                throw e;
+                            }
+                        }
+                        retryCount++;
+
+                        long waitTimeMs = (long) (Math.pow(2, retryCount) * 1000);
+                        try {
+                            TimeUnit.MILLISECONDS.sleep(waitTimeMs);
+                        }
+                        catch (InterruptedException e) {
+                            Thread.currentThread().interrupt();
+                            throw new IOException("Upload interrupted", e);
+                        }
+                    }
+                    return response;
+                }).build();
         this.uri = uri;
     }
 
-    private void uploadFromStream(InputStream inputStream, String stageName, String relativePath, String name, long fileSize) throws IOException {
+    private void uploadFromStream(InputStream inputStream, String stageName, String relativePath, String name, long fileSize)
+            throws IOException
+    {
         // multipart upload input stream into /v1/upload_to_stage
         RequestBody requestBody = new MultipartBody.Builder()
                 .setType(MultipartBody.FORM)
@@ -73,24 +115,31 @@ private void uploadFromStream(InputStream inputStream, String stageName, String
                 .build();
         try {
             executeInternal(request, true);
-
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             throw new IOException("uploadFromStreamAPI failed", e);
         }
-
     }
 
-    private void uploadFromStream(InputStream inputStream, Headers headers, String presignedUrl, long fileSize) throws IOException {
-        requireNonNull(inputStream, "inputStream is null");
-        Request r = putRequest(headers, presignedUrl, inputStream, fileSize);
+    private void uploadFromStream(InputStream inputStream, Headers headers, String presignedUrl, long fileSize)
+            throws IOException
+    {
+        logger.info("Starting upload: size=" + fileSize + " bytes, url=" + presignedUrl);
+        long startTime = System.currentTimeMillis();
         try {
+            Request r = putRequest(headers, presignedUrl, inputStream, fileSize);
             executeInternal(r, true);
-        } catch (IOException e) {
-            throw new IOException(format(" uploadFromStream failed, file size is %s kb, error is %s", fileSize / 1024.0, presignedUrl, e.toString()));
+            logger.info("Upload completed in " + (System.currentTimeMillis() - startTime) + "ms");
+        }
+        catch (IOException e) {
+            logger.severe("Upload failed after " + (System.currentTimeMillis() - startTime) + "ms: " + e.getMessage());
+            throw e;
         }
     }
 
-    private ResponseBody executeInternal(Request request, boolean shouldClose) throws IOException {
+    private ResponseBody executeInternal(Request request, boolean shouldClose)
+            throws IOException
+    {
         requireNonNull(request, "request is null");
         long start = System.nanoTime();
         long attempts = 0;
@@ -110,9 +159,11 @@ private ResponseBody executeInternal(Request request, boolean shouldClose) throw
 
                 try {
                     MILLISECONDS.sleep(attempts * 100);
-                } catch (InterruptedException e) {
+                }
+                catch (InterruptedException e) {
                     try {
-                    } finally {
+                    }
+                    finally {
                         Thread.currentThread().interrupt();
                     }
                     throw new RuntimeException("StatementClient thread was interrupted");
@@ -124,122 +175,169 @@ private ResponseBody executeInternal(Request request, boolean shouldClose) throw
                 response = client.newCall(request).execute();
                 if (response.isSuccessful()) {
                     return response.body();
-                } else if (response.code() == 401) {
+                }
+                else if (response.code() == 401) {
                     throw new RuntimeException("Error exeucte presign, Unauthorized user: " + response.code() + " " + response.message());
-                } else if (response.code() >= 503) {
+                }
+                else if (response.code() >= 503) {
                     cause = new RuntimeException("Error execute presign, service unavailable: " + response.code() + " " + response.message());
-                } else if (response.code() >= 400) {
+                }
+                else if (response.code() >= 400) {
                     cause = new RuntimeException("Error execute presign, configuration error: " + response.code() + " " + response.message());
                 }
-            } catch (SocketTimeoutException e) {
+            }
+            catch (SocketTimeoutException e) {
                 logger.warning("Error execute presign, socket timeout: " + e.getMessage());
                 cause = new RuntimeException("Error execute presign, request is " + request.toString() + "socket timeout: " + e.getMessage());
-            } catch (RuntimeException e) {
+            }
+            catch (RuntimeException e) {
                 cause = e;
-            } finally {
+            }
+            finally {
                 if (shouldClose) {
                     try {
                         if (response != null) {
                             response.close();
                         }
-                    } catch (Exception e) {
+                    }
+                    catch (Exception e) {
                         // ignore
                     }
                 }
             }
-
         }
     }
 
     @Override
     public void presignUpload(File srcFile, InputStream inputStream, Headers headers,
-                              String presignedUrl, long fileSize, boolean uploadFromStream) throws IOException {
+            String presignedUrl, long fileSize, boolean uploadFromStream)
+            throws IOException
+    {
 
         InputStream it = null;
         if (!uploadFromStream) {
             it = Files.newInputStream(srcFile.toPath());
-        } else {
+        }
+        else {
             it = inputStream;
         }
         uploadFromStream(it, headers, presignedUrl, fileSize);
     }
 
     @Override
-    public void presignUpload(File srcFile, InputStream inputStream, String stageName, String relativePath, String name, long fileSize, boolean uploadFromStream) throws IOException {
+    public void presignUpload(File srcFile, InputStream inputStream, String stageName, String relativePath, String name, long fileSize, boolean uploadFromStream)
+            throws IOException
+    {
         if (!uploadFromStream) {
             try (InputStream it = Files.newInputStream(srcFile.toPath())) {
                 uploadFromStream(it, stageName, relativePath, name, fileSize);
             }
-        } else {
+        }
+        else {
             uploadFromStream(inputStream, stageName, relativePath, name, fileSize);
         }
     }
 
     @Override
-    public void presignDownload(String destFileName, Headers headers, String presignedUrl) {
+    public void presignDownload(String destFileName, Headers headers, String presignedUrl)
+    {
         Request r = getRequest(headers, presignedUrl);
         try (ResponseBody body = executeInternal(r, false)) {
             BufferedSink sink = Okio.buffer(Okio.sink(new File(destFileName)));
             sink.writeAll(body.source());
             sink.close();
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             throw new RuntimeException("presignDownload failed", e);
         }
     }
 
     @Override
-    public InputStream presignDownloadStream(Headers headers, String presignedUrl) {
+    public InputStream presignDownloadStream(Headers headers, String presignedUrl)
+    {
         Request r = getRequest(headers, presignedUrl);
         try {
             ResponseBody responseBody = executeInternal(r, false);
             return responseBody.byteStream();
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             throw new RuntimeException("presignDownloadStream failed", e);
         }
     }
 
-    private Request getRequest(Headers headers, String url) {
+    private Request getRequest(Headers headers, String url)
+    {
         return new Request.Builder().headers(headers).url(url).get().build();
     }
 
-    private Request putRequest(Headers headers, String url, InputStream inputStream, long fileSize)
-            throws IOException {
-        RequestBody input = new InputStreamRequestBody(null, inputStream, fileSize);
-        return new Request.Builder().headers(headers).url(url).put(input).build();
-    }
+    private Request putRequest(Headers headers, String presignedUrl, InputStream inputStream, long fileSize) {
+        RequestBody requestBody = new RequestBody() {
+            @Override
+            public MediaType contentType() {
+                return MediaType.parse("application/octet-stream");
+            }
+
+            @Override
+            public long contentLength() {
+                return fileSize;
+            }
 
+            @Override
+            public void writeTo(BufferedSink sink) throws IOException {
+                try (Source source = Okio.source(inputStream)) {
+                    sink.writeAll(source);
+                }
+            }
+        };
+
+        return new Request.Builder()
+                .url(presignedUrl)
+                .put(requestBody)
+                .headers(headers)
+                .build();
+    }
 }
 
-class InputStreamRequestBody extends RequestBody {
+class InputStreamRequestBody
+        extends RequestBody
+{
     private final InputStream inputStream;
     private final MediaType contentType;
     private final long fileSize;
     private static final Logger logger = Logger.getLogger(InputStreamRequestBody.class.getPackage().getName());
 
-    public InputStreamRequestBody(MediaType contentType, InputStream inputStream, long fileSize) {
-        if (inputStream == null) throw new NullPointerException("inputStream == null");
+    public InputStreamRequestBody(MediaType contentType, InputStream inputStream, long fileSize)
+    {
+        if (inputStream == null) {
+            throw new NullPointerException("inputStream == null");
+        }
         this.contentType = contentType;
         this.inputStream = inputStream;
         this.fileSize = fileSize;
     }
 
     @Nullable
     @Override
-    public MediaType contentType() {
+    public MediaType contentType()
+    {
         return contentType;
     }
 
     @Override
-    public long contentLength() {
+    public long contentLength()
+    {
         return fileSize; // return the actual file size
 //        return inputStream.available() == 0 ? -1 : inputStream.available();
     }
 
     @Override
-    public void writeTo(@NonNull BufferedSink sink) throws IOException {
+    public void writeTo(@NonNull BufferedSink sink)
+            throws IOException
+    {
         try (Source source = Okio.source(inputStream)) {
             sink.writeAll(source);
-        } catch (IOException e) {
+        }
+        catch (IOException e) {
             logger.warning(format("writeTo failed, error is %s, cause is %s", e.getMessage(), e.getCause()));
         }
     }