From 39bcb45569769d50db56eba96af5c89fc22e68f2 Mon Sep 17 00:00:00 2001 From: Mounir Date: Tue, 6 Oct 2026 17:20:22 +0200 Subject: [PATCH] Fix large extract pmtiles from large source file --- pmtiles/bucket.go | 27 ++++++++++++--------------- pmtiles/bucket_test.go | 31 +++++++++++++++++++++++++++++++ pmtiles/extract.go | 4 ++-- pmtiles/show.go | 3 ++- pmtiles/verify.go | 2 +- 5 files changed, 48 insertions(+), 19 deletions(-) diff --git a/pmtiles/bucket.go b/pmtiles/bucket.go index b5b60ab..c786826 100644 --- a/pmtiles/bucket.go +++ b/pmtiles/bucket.go @@ -121,34 +121,31 @@ func generateEtagFromInts(ns ...int64) string { func (b FileBucket) NewRangeReaderEtag(_ context.Context, key string, offset, length int64, etag string) (io.ReadCloser, string, int, error) { name := filepath.Join(b.path, key) file, err := os.Open(name) - defer file.Close() if err != nil { return nil, "", 404, err } info, err := file.Stat() if err != nil { + file.Close() return nil, "", 404, err } newEtag := generateEtagFromInts(info.ModTime().UnixNano(), info.Size()) if len(etag) > 0 && etag != newEtag { + file.Close() return nil, "", 412, &RefreshRequiredError{} } - result := make([]byte, length) - read, err := file.ReadAt(result, offset) - if err == io.EOF { - part := result[0:read] - return io.NopCloser(bytes.NewReader(part)), newEtag, 206, nil - } - - if err != nil { - return nil, "", 500, err - } - if read != int(length) { - return nil, "", 416, fmt.Errorf("Expected to read %d bytes but only read %d", length, read) - } + // stream the range instead of buffering it: ranges can be hundreds of GB. + // reads past the end of the file are truncated (io.EOF). + return §ionReadCloser{ + Reader: io.NewSectionReader(file, offset, length), + Closer: file, + }, newEtag, 206, nil +} - return io.NopCloser(bytes.NewReader(result)), newEtag, 206, nil +type sectionReadCloser struct { + io.Reader + io.Closer } func (b FileBucket) Close() error { diff --git a/pmtiles/bucket_test.go b/pmtiles/bucket_test.go index c8c16be..38f2e1a 100644 --- a/pmtiles/bucket_test.go +++ b/pmtiles/bucket_test.go @@ -152,6 +152,7 @@ func TestFileBucketReplace(t *testing.T) { assert.Equal(t, 206, status) assert.Nil(t, err) data, err := io.ReadAll(reader) + reader.Close() assert.Nil(t, err) assert.Equal(t, []byte{2}, data) @@ -161,6 +162,7 @@ func TestFileBucketReplace(t *testing.T) { assert.Equal(t, 206, status) assert.Nil(t, err) data, err = io.ReadAll(reader) + reader.Close() assert.Nil(t, err) assert.NotEqual(t, etag1, etag2) assert.Equal(t, []byte{5}, data) @@ -188,6 +190,7 @@ func TestFileBucketRename(t *testing.T) { assert.Equal(t, 206, status) assert.Nil(t, err) data, err := io.ReadAll(reader) + reader.Close() assert.Nil(t, err) assert.Equal(t, []byte{2}, data) @@ -198,6 +201,7 @@ func TestFileBucketRename(t *testing.T) { assert.Equal(t, 206, status) assert.Nil(t, err) data, err = io.ReadAll(reader) + reader.Close() assert.Nil(t, err) assert.NotEqual(t, etag1, etag2) assert.Equal(t, []byte{5}, data) @@ -219,6 +223,7 @@ func TestFileShorterThan16K(t *testing.T) { assert.Equal(t, 206, status) assert.Nil(t, err) data, err := io.ReadAll(reader) + reader.Close() assert.Nil(t, err) assert.Equal(t, 3, len(data)) } @@ -276,3 +281,29 @@ func TestGenerationEtag(t *testing.T) { assert.Equal(t, int64(123), etagToGeneration("123")) assert.Equal(t, "123", generationToEtag(int64(123))) } + +func TestFileBucketRangeReader(t *testing.T) { + dir := t.TempDir() + assert.Nil(t, os.WriteFile(filepath.Join(dir, "a.bin"), []byte("0123456789"), 0666)) + bucket := NewFileBucket(dir) + + read := func(offset, length int64) string { + r, err := bucket.NewRangeReader(context.Background(), "a.bin", offset, length) + assert.Nil(t, err) + b, err := io.ReadAll(r) + assert.Nil(t, err) + assert.Nil(t, r.Close()) + return string(b) + } + + assert.Equal(t, "2345", read(2, 4)) + assert.Equal(t, "89", read(8, 100)) + assert.Equal(t, "", read(50, 4)) + + _, err := bucket.NewRangeReader(context.Background(), "missing.bin", 0, 4) + assert.NotNil(t, err) + + _, _, status, err := bucket.NewRangeReaderEtag(context.Background(), "a.bin", 0, 4, "\"stale\"") + assert.Equal(t, 412, status) + assert.NotNil(t, err) +} diff --git a/pmtiles/extract.go b/pmtiles/extract.go index 2d119a3..96daa8b 100644 --- a/pmtiles/extract.go +++ b/pmtiles/extract.go @@ -277,10 +277,10 @@ func Extract(ctx context.Context, logger *log.Logger, bucketURL string, key stri return fmt.Errorf("Failed to create range reader for %s, %w", key, err) } b, err := io.ReadAll(r) + r.Close() if err != nil { return err } - r.Close() header, err := DeserializeHeader(b[0:HeaderV3LenBytes]) @@ -539,6 +539,7 @@ func Extract(ctx context.Context, logger *log.Logger, bucketURL string, key stri if err != nil { return err } + defer tileReader.Close() offsetWriter := io.NewOffsetWriter(outfile, int64(header.TileDataOffset)+int64(or.Rng.DstOffset)) for _, cd := range or.CopyDiscards { @@ -553,7 +554,6 @@ func Extract(ctx context.Context, logger *log.Logger, bucketURL string, key stri return err } } - tileReader.Close() return nil } diff --git a/pmtiles/show.go b/pmtiles/show.go index 3207e25..3378f7d 100644 --- a/pmtiles/show.go +++ b/pmtiles/show.go @@ -37,10 +37,10 @@ func Show(_ *log.Logger, output io.Writer, bucketURL string, key string, showHea return fmt.Errorf("Failed to create range reader for %s, %w", key, err) } b, err := io.ReadAll(r) + r.Close() if err != nil { return fmt.Errorf("Failed to read %s, %w", key, err) } - r.Close() header, err := DeserializeHeader(b[0:HeaderV3LenBytes]) if err != nil { @@ -60,6 +60,7 @@ func Show(_ *log.Logger, output io.Writer, bucketURL string, key string, showHea } metadataBytes, err := DeserializeMetadataBytes(metadataReader, header.InternalCompression) + metadataReader.Close() if err != nil { return fmt.Errorf("Failed to read %s, %w", key, err) } diff --git a/pmtiles/verify.go b/pmtiles/verify.go index c10c293..7def92c 100644 --- a/pmtiles/verify.go +++ b/pmtiles/verify.go @@ -36,10 +36,10 @@ func Verify(logger *log.Logger, file string) error { return fmt.Errorf("failed to create range reader for %s, %w", key, err) } b, err := io.ReadAll(r) + r.Close() if err != nil { return fmt.Errorf("failed to read %s, %w", key, err) } - r.Close() header, err := DeserializeHeader(b[0:HeaderV3LenBytes])