chunked: use temporary file for tar-split data

Replace the in-memory buffer with a O_TMPFILE file.  This reduces the
memory requirements for a partial pull since the tar-split data can be
written to disk.

Signed-off-by: Giuseppe Scrivano <[email protected]>

Author: giuseppe

drivers/driver.go

@@ -216,7 +216,7 @@ type DriverWithDifferOutput struct {
 	CompressedDigest   digest.Digest
 	Metadata           string
 	BigData            map[string][]byte
-	TarSplit           []byte // nil if not available
+	TarSplit           *os.File // nil if not available
 	TOCDigest          digest.Digest
 	// RootDirMode is the mode of the root directory of the layer, if specified.
 	RootDirMode *os.FileMode

layers.go

@@ -2550,10 +2550,14 @@ func (r *layerStore) applyDiffFromStagingDirectory(id string, diffOutput *driver
 		if err != nil {
 			compressor = pgzip.NewWriter(&tsdata)
 		}
+		if _, err := diffOutput.TarSplit.Seek(0, 0); err != nil {
+			return err
+		}
+
 		if err := compressor.SetConcurrency(1024*1024, 1); err != nil { // 1024*1024 is the hard-coded default; we're not changing that
 			logrus.Infof("setting compression concurrency threads to 1: %v; ignoring", err)
 		}
-		if _, err := compressor.Write(diffOutput.TarSplit); err != nil {
+		if _, err := diffOutput.TarSplit.WriteTo(compressor); err != nil {
 			compressor.Close()
 			return err
 		}

pkg/chunked/compression_linux.go

@@ -6,7 +6,11 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"io/fs"
 	"maps"
+	"math/rand"
+	"os"
+	"path/filepath"
 	"slices"
 	"strconv"
 	"time"
@@ -18,6 +22,7 @@ import (
 	"github.com/vbatts/tar-split/archive/tar"
 	"github.com/vbatts/tar-split/tar/asm"
 	"github.com/vbatts/tar-split/tar/storage"
+	"golang.org/x/sys/unix"
 )
 
 const (
@@ -157,10 +162,41 @@ func readEstargzChunkedManifest(blobStream ImageSourceSeekable, blobSize int64,
 	return manifestUncompressed, tocOffset, nil
 }
 
+func openTmpFile(tmpDir string) (fd int, err error) {
+	fd, err = unix.Open(tmpDir, unix.O_TMPFILE|unix.O_RDWR|unix.O_CLOEXEC|unix.O_EXCL, 0o600)
+	if err == nil {
+		return fd, nil
+	}
+	return openTmpFileNoTmpFile(tmpDir)
+}
+
+// openTmpFileNoTmpFile is a fallback used by openTmpFile when the underlying file system does not
+// support O_TMPFILE.
+func openTmpFileNoTmpFile(tmpDir string) (fd int, err error) {
+	for i := 0; i < 100; i++ {
+		name := fmt.Sprintf(".tmpfile-%d", rand.Int63())
+		path := filepath.Join(tmpDir, name)
+
+		fd, err := unix.Open(path, unix.O_RDWR|unix.O_CREAT|unix.O_EXCL|unix.O_CLOEXEC, 0o600)
+		if err == nil {
+			// Unlink the file immediately so that only the open fd refers to it.
+			_ = os.Remove(path)
+			return fd, nil
+		}
+		if !errors.Is(err, os.ErrExist) {
+			return -1, &fs.PathError{Op: "open", Path: tmpDir, Err: err}
+		}
+	}
+	// report the original error if the fallback failed
+	return -1, &fs.PathError{Op: "open O_TMPFILE", Path: tmpDir, Err: err}
+}
+
 // readZstdChunkedManifest reads the zstd:chunked manifest from the seekable stream blobStream.
+// tmpDir is a directory where the tar-split temporary file is written to.  The file is opened with
+// O_TMPFILE so that it is automatically removed when it is closed.
 // Returns (manifest blob, parsed manifest, tar-split blob or nil, manifest offset).
 // It may return an error matching ErrFallbackToOrdinaryLayerDownload / errFallbackCanConvert.
-func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Digest, annotations map[string]string) (_ []byte, _ *minimal.TOC, _ []byte, _ int64, retErr error) {
+func readZstdChunkedManifest(tmpDir string, blobStream ImageSourceSeekable, tocDigest digest.Digest, annotations map[string]string) (_ []byte, _ *minimal.TOC, _ *os.File, _ int64, retErr error) {
 	offsetMetadata := annotations[minimal.ManifestInfoKey]
 	if offsetMetadata == "" {
 		return nil, nil, nil, 0, fmt.Errorf("%q annotation missing", minimal.ManifestInfoKey)
@@ -245,7 +281,7 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Di
 		return nil, nil, nil, 0, fmt.Errorf("unmarshaling TOC: %w", err)
 	}
 
-	var decodedTarSplit []byte = nil
+	var decodedTarSplit *os.File
 	if toc.TarSplitDigest != "" {
 		if tarSplitChunk.Offset <= 0 {
 			return nil, nil, nil, 0, fmt.Errorf("TOC requires a tar-split, but the %s annotation does not describe a position", minimal.TarSplitInfoKey)
@@ -254,14 +290,20 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Di
 		if err != nil {
 			return nil, nil, nil, 0, err
 		}
-		decodedTarSplit, err = decodeAndValidateBlob(tarSplit, tarSplitLengthUncompressed, toc.TarSplitDigest.String())
+		fd, err := openTmpFile(tmpDir)
 		if err != nil {
+			return nil, nil, nil, 0, err
+		}
+		decodedTarSplit = os.NewFile(uintptr(fd), "decoded-tar-split")
+		if err := decodeAndValidateBlobToStream(tarSplit, decodedTarSplit, toc.TarSplitDigest.String()); err != nil {
+			decodedTarSplit.Close()
 			return nil, nil, nil, 0, fmt.Errorf("validating and decompressing tar-split: %w", err)
 		}
 		// We use the TOC for creating on-disk files, but the tar-split for creating metadata
 		// when exporting the layer contents. Ensure the two match, otherwise local inspection of a container
 		// might be misleading about the exported contents.
 		if err := ensureTOCMatchesTarSplit(toc, decodedTarSplit); err != nil {
+			decodedTarSplit.Close()
 			return nil, nil, nil, 0, fmt.Errorf("tar-split and TOC data is inconsistent: %w", err)
 		}
 	} else if tarSplitChunk.Offset > 0 {
@@ -278,7 +320,7 @@ func readZstdChunkedManifest(blobStream ImageSourceSeekable, tocDigest digest.Di
 }
 
 // ensureTOCMatchesTarSplit validates that toc and tarSplit contain _exactly_ the same entries.
-func ensureTOCMatchesTarSplit(toc *minimal.TOC, tarSplit []byte) error {
+func ensureTOCMatchesTarSplit(toc *minimal.TOC, tarSplit *os.File) error {
 	pendingFiles := map[string]*minimal.FileMetadata{} // Name -> an entry in toc.Entries
 	for i := range toc.Entries {
 		e := &toc.Entries[i]
@@ -290,7 +332,11 @@ func ensureTOCMatchesTarSplit(toc *minimal.TOC, tarSplit []byte) error {
 		}
 	}
 
-	unpacker := storage.NewJSONUnpacker(bytes.NewReader(tarSplit))
+	if _, err := tarSplit.Seek(0, 0); err != nil {
+		return err
+	}
+
+	unpacker := storage.NewJSONUnpacker(tarSplit)
 	if err := asm.IterateHeaders(unpacker, func(hdr *tar.Header) error {
 		e, ok := pendingFiles[hdr.Name]
 		if !ok {
@@ -320,10 +366,10 @@ func ensureTOCMatchesTarSplit(toc *minimal.TOC, tarSplit []byte) error {
 }
 
 // tarSizeFromTarSplit computes the total tarball size, using only the tarSplit metadata
-func tarSizeFromTarSplit(tarSplit []byte) (int64, error) {
+func tarSizeFromTarSplit(tarSplit io.Reader) (int64, error) {
 	var res int64 = 0
 
-	unpacker := storage.NewJSONUnpacker(bytes.NewReader(tarSplit))
+	unpacker := storage.NewJSONUnpacker(tarSplit)
 	for {
 		entry, err := unpacker.Next()
 		if err != nil {
@@ -464,3 +510,18 @@ func decodeAndValidateBlob(blob []byte, lengthUncompressed uint64, expectedCompr
 	b := make([]byte, 0, lengthUncompressed)
 	return decoder.DecodeAll(blob, b)
 }
+
+func decodeAndValidateBlobToStream(blob []byte, w *os.File, expectedCompressedChecksum string) error {
+	if err := validateBlob(blob, expectedCompressedChecksum); err != nil {
+		return err
+	}
+
+	decoder, err := zstd.NewReader(bytes.NewReader(blob)) //nolint:contextcheck
+	if err != nil {
+		return err
+	}
+	defer decoder.Close()
+
+	_, err = decoder.WriteTo(w)
+	return err
+}

pkg/chunked/compression_linux_test.go

@@ -2,14 +2,17 @@ package chunked
 
 import (
 	"bytes"
+	"fmt"
 	"io"
+	"os"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"github.com/vbatts/tar-split/archive/tar"
 	"github.com/vbatts/tar-split/tar/asm"
 	"github.com/vbatts/tar-split/tar/storage"
+	"golang.org/x/sys/unix"
 )
 
 func TestTarSizeFromTarSplit(t *testing.T) {
@@ -39,7 +42,27 @@ func TestTarSizeFromTarSplit(t *testing.T) {
 	_, err = io.Copy(io.Discard, tsReader)
 	require.NoError(t, err)
 
-	res, err := tarSizeFromTarSplit(tarSplit.Bytes())
+	res, err := tarSizeFromTarSplit(&tarSplit)
 	require.NoError(t, err)
 	assert.Equal(t, expectedTarSize, res)
 }
+
+func TestOpenTmpFile(t *testing.T) {
+	for i := 0; i < 1000; i++ {
+		// scope for cleanup
+		f := func(fn func(tmpDir string) (int, error)) {
+			fd, err := fn(t.TempDir())
+			assert.NoError(t, err)
+			defer unix.Close(fd)
+
+			path, err := os.Readlink(fmt.Sprintf("/proc/self/fd/%d", fd))
+			assert.NoError(t, err)
+
+			// the path under /proc/self/fd/$FD has the prefix "(deleted)" when the file
+			// is unlinked
+			assert.Contains(t, path, "(deleted)")
+		}
+		f(openTmpFile)
+		f(openTmpFileNoTmpFile)
+	}
+}

pkg/chunked/storage_linux.go

@@ -2,7 +2,6 @@ package chunked
 
 import (
 	archivetar "archive/tar"
-	"bytes"
 	"context"
 	"encoding/base64"
 	"errors"
@@ -89,7 +88,7 @@ type chunkedDiffer struct {
 	tocOffset           int64
 	manifest            []byte
 	toc                 *minimal.TOC // The parsed contents of manifest, or nil if not yet available
-	tarSplit            []byte
+	tarSplit            *os.File
 	uncompressedTarSize int64 // -1 if unknown
 	// skipValidation is set to true if the individual files in
 	// the layer are trusted and should not be validated.
@@ -194,6 +193,11 @@ func (c *chunkedDiffer) convertTarToZstdChunked(destDirectory string, payload *o
 }
 
 func (c *chunkedDiffer) Close() error {
+	if c.tarSplit != nil {
+		err := c.tarSplit.Close()
+		c.tarSplit = nil
+		return err
+	}
 	return nil
 }
 
@@ -337,13 +341,16 @@ func makeConvertFromRawDiffer(store storage.Store, blobDigest digest.Digest, blo
 // makeZstdChunkedDiffer sets up a chunkedDiffer for a zstd:chunked layer.
 // It may return an error matching ErrFallbackToOrdinaryLayerDownload / errFallbackCanConvert.
 func makeZstdChunkedDiffer(store storage.Store, blobSize int64, tocDigest digest.Digest, annotations map[string]string, iss ImageSourceSeekable, pullOptions pullOptions) (*chunkedDiffer, error) {
-	manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(iss, tocDigest, annotations)
+	manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(store.RunRoot(), iss, tocDigest, annotations)
 	if err != nil { // May be ErrFallbackToOrdinaryLayerDownload / errFallbackCanConvert
 		return nil, fmt.Errorf("read zstd:chunked manifest: %w", err)
 	}
 
 	var uncompressedTarSize int64 = -1
 	if tarSplit != nil {
+		if _, err := tarSplit.Seek(0, 0); err != nil {
+			return nil, err
+		}
 		uncompressedTarSize, err = tarSizeFromTarSplit(tarSplit)
 		if err != nil {
 			return nil, fmt.Errorf("computing size from tar-split: %w", err)
@@ -1439,7 +1446,7 @@ func (c *chunkedDiffer) ApplyDiff(dest string, options *archive.TarOptions, diff
 		if tocDigest == nil {
 			return graphdriver.DriverWithDifferOutput{}, fmt.Errorf("internal error: just-created zstd:chunked missing TOC digest")
 		}
-		manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(fileSource, *tocDigest, annotations)
+		manifest, toc, tarSplit, tocOffset, err := readZstdChunkedManifest(dest, fileSource, *tocDigest, annotations)
 		if err != nil {
 			return graphdriver.DriverWithDifferOutput{}, fmt.Errorf("read zstd:chunked manifest: %w", err)
 		}
@@ -1846,7 +1853,10 @@ func (c *chunkedDiffer) ApplyDiff(dest string, options *archive.TarOptions, diff
 		case c.pullOptions.insecureAllowUnpredictableImageContents:
 			// Oh well.  Skip the costly digest computation.
 		case output.TarSplit != nil:
-			metadata := tsStorage.NewJSONUnpacker(bytes.NewReader(output.TarSplit))
+			if _, err := output.TarSplit.Seek(0, 0); err != nil {
+				return output, err
+			}
+			metadata := tsStorage.NewJSONUnpacker(output.TarSplit)
 			fg := newStagedFileGetter(dirFile, flatPathNameMap)
 			digester := digest.Canonical.Digester()
 			if err := asm.WriteOutputTarStream(fg, metadata, digester.Hash()); err != nil {

pkg/chunked/zstdchunked_test.go

@@ -179,7 +179,7 @@ func TestGenerateAndParseManifest(t *testing.T) {
 	tocDigest, err := toc.GetTOCDigest(annotations)
 	require.NoError(t, err)
 	require.NotNil(t, tocDigest)
-	manifest, decodedTOC, _, _, err := readZstdChunkedManifest(s, *tocDigest, annotations)
+	manifest, decodedTOC, _, _, err := readZstdChunkedManifest(t.TempDir(), s, *tocDigest, annotations)
 	require.NoError(t, err)
 
 	var toc minimal.TOC