From 5882fe8b7ee3fa8a3618d8c7d8dc3ea7fc411ec0 Mon Sep 17 00:00:00 2001
From: Brandon Liu <bdon@users.noreply.github.com>
Date: Sun, 15 Sep 2024 16:00:52 +0800
Subject: [PATCH] Simplify upload, reduce multipart upload part size [#89]
 (#187)

* Simplify upload, reduce multipart upload part size [#89]
---
 pmtiles/upload.go      | 47 ++++++++++++++----------------------------
 pmtiles/upload_test.go | 11 ++++++++++
 2 files changed, 26 insertions(+), 32 deletions(-)
 create mode 100644 pmtiles/upload_test.go

diff --git a/pmtiles/upload.go b/pmtiles/upload.go
index 53703b3..53c52eb 100644
--- a/pmtiles/upload.go
+++ b/pmtiles/upload.go
@@ -10,9 +10,18 @@ import (
 	"os"
 )
 
+// Determine the multipart block size based on the total file size.
+func partSizeBytes(totalSize int64) int {
+	if totalSize/(5*1024*1024) >= 10_000 {
+		return int(totalSize/10_000 + 1)
+	}
+	return 5 * 1024 * 1024
+}
+
 // Upload a pmtiles archive to a bucket.
 func Upload(logger *log.Logger, input string, bucket string, key string, maxConcurrency int) error {
 	ctx := context.Background()
+
 	b, err := blob.OpenBucket(ctx, bucket)
 	if err != nil {
 		return fmt.Errorf("Failed to setup bucket: %w", err)
@@ -24,17 +33,14 @@ func Upload(logger *log.Logger, input string, bucket string, key string, maxConc
 		return fmt.Errorf("Failed to open file: %w", err)
 	}
 	defer f.Close()
+
 	filestat, err := f.Stat()
 	if err != nil {
-		return fmt.Errorf("Failed to open file: %w", err)
+		return fmt.Errorf("Failed to stat file: %w", err)
 	}
-	bar := progressbar.Default(filestat.Size())
-
-	nChunks := int64(0)
-	buffer := make([]byte, 8*1024)
 
 	opts := &blob.WriterOptions{
-		BufferSize:     256 * 1024 * 1024,
+		BufferSize:     partSizeBytes(filestat.Size()),
 		MaxConcurrency: maxConcurrency,
 	}
 
@@ -43,34 +49,11 @@ func Upload(logger *log.Logger, input string, bucket string, key string, maxConc
 		return fmt.Errorf("Failed to obtain writer: %w", err)
 	}
 
-	for {
-		n, err := f.Read(buffer)
-
-		if n == 0 {
-			if err == nil {
-				continue
-			}
-			if err == io.EOF {
-				break
-			}
-			logger.Fatal(err)
-		}
-
-		nChunks++
-
-		_, err = w.Write(buffer[:n])
-		if err != nil {
-			return fmt.Errorf("Failed to write to bucket: %w", err)
-		}
-		bar.Add(n)
-
-		if err != nil && err != io.EOF {
-			return fmt.Errorf("Failed to write data, %w", err)
-		}
-	}
+	bar := progressbar.Default(filestat.Size())
+	io.Copy(io.MultiWriter(w, bar), f)
 
 	if err := w.Close(); err != nil {
-		return fmt.Errorf("Failed to close: %w", err)
+		return fmt.Errorf("Failed to complete upload: %w", err)
 	}
 
 	return nil
diff --git a/pmtiles/upload_test.go b/pmtiles/upload_test.go
new file mode 100644
index 0000000..4e0da4a
--- /dev/null
+++ b/pmtiles/upload_test.go
@@ -0,0 +1,11 @@
+package pmtiles
+
+import (
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestPartSizeBytes(t *testing.T) {
+	assert.Equal(t, 5*1024*1024, partSizeBytes(100))
+	assert.Equal(t, 6442451, partSizeBytes(60*1024*1024*1024))
+}