Skip to content

Commit

Permalink
Simplify upload, reduce multipart upload part size [#89] (#187)
Browse files Browse the repository at this point in the history
* Simplify upload, reduce multipart upload part size [#89]
  • Loading branch information
bdon authored Sep 15, 2024
1 parent 1ea0ffd commit 5882fe8
Show file tree
Hide file tree
Showing 2 changed files with 26 additions and 32 deletions.
47 changes: 15 additions & 32 deletions pmtiles/upload.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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,
}

Expand All @@ -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
Expand Down
11 changes: 11 additions & 0 deletions pmtiles/upload_test.go
Original file line number Diff line number Diff line change
@@ -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))
}

0 comments on commit 5882fe8

Please sign in to comment.