Fix large extract pmtiles from large source file

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 &sectionReadCloser{
+		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 {

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)
+}

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
 		}
 

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)
 		}

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])