Use sync.Pools for allocations (#1568)

This should help a lot with memory usage when building lots of things in
the same process.

Signed-off-by: Jon Johnson <[email protected]>

Author: jonjohnsonjr

pkg/apk/expandapk/expandapk.go

@@ -28,6 +28,40 @@ import (
 	"go.opentelemetry.io/otel"
 )
 
+var slicePool = sync.Pool{
+	New: func() interface{} {
+		return make([]byte, 1<<20)
+	},
+}
+
+func pooledSlice() []byte {
+	return slicePool.Get().([]byte)
+}
+
+var readerPool = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewReaderSize(nil, 1<<20)
+	},
+}
+
+func pooledBufioReader(r io.Reader) *bufio.Reader {
+	br := readerPool.Get().(*bufio.Reader)
+	br.Reset(r)
+	return br
+}
+
+var writerPool = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewWriterSize(nil, 1<<20)
+	},
+}
+
+func pooledBufioWriter(w io.Writer) *bufio.Writer {
+	bw := writerPool.Get().(*bufio.Writer)
+	bw.Reset(w)
+	return bw
+}
+
 // APKExpanded contains information about and reference to an expanded APK package.
 // Close() deletes all temporary files and directories created during the expansion process.
 type APKExpanded struct {
@@ -66,8 +100,6 @@ type APKExpanded struct {
 	controlData []byte
 }
 
-const meg = 1 << 20
-
 func (a *APKExpanded) ControlData() ([]byte, error) {
 	a.Lock()
 	defer a.Unlock()
@@ -100,20 +132,16 @@ func (a *APKExpanded) PackageData() (*os.File, error) {
 		return nil, fmt.Errorf("opening package data file: %w", err)
 	}
 
-	// Use min(1MB, a.Size) bufio to avoid GC pressure for small packages.
-	bufSize := meg
-	if total := int(a.Size); total != 0 && total < bufSize {
-		bufSize = total
-	}
-
 	// Handle old caches without the uncompressed file.
 	f, err := os.Open(a.PackageFile)
 	if err != nil {
 		return nil, fmt.Errorf("opening %q: %w", a.PackageFile, err)
 	}
 	defer f.Close()
 
-	br := bufio.NewReaderSize(f, bufSize)
+	br := pooledBufioReader(f)
+	defer readerPool.Put(br)
+
 	zr, err := gzip.NewReader(br)
 	if err != nil {
 		return nil, fmt.Errorf("parsing %q: %w", a.PackageFile, err)
@@ -124,7 +152,9 @@ func (a *APKExpanded) PackageData() (*os.File, error) {
 		return nil, fmt.Errorf("opening tar file %q: %w", a.TarFile, err)
 	}
 
-	buf := make([]byte, bufSize)
+	buf := pooledSlice()
+	defer slicePool.Put(buf)
+
 	if _, err := io.CopyBuffer(uf, zr, buf); err != nil {
 		return nil, fmt.Errorf("decompressing %q: %w", a.PackageFile, err)
 	}
@@ -390,7 +420,9 @@ func ExpandApk(ctx context.Context, source io.Reader, cacheDir string) (*APKExpa
 			if err != nil {
 				return nil, fmt.Errorf("opening tar file: %w", err)
 			}
-			bw := bufio.NewWriterSize(tarfile, 1<<20)
+			bw := pooledBufioWriter(tarfile)
+			defer writerPool.Put(bw)
+
 			tr := io.TeeReader(gzi, bw)
 
 			if err := checkSums(ctx, tr); err != nil {

pkg/apk/internal/tarfs/tarfs.go

@@ -24,9 +24,22 @@ import (
 	"io/fs"
 	"path"
 	"slices"
+	"sync"
 	"time"
 )
 
+var readerPool = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewReaderSize(nil, 1<<20)
+	},
+}
+
+func pooledBufioReader(r io.Reader) *bufio.Reader {
+	br := readerPool.Get().(*bufio.Reader)
+	br.Reset(r)
+	return br
+}
+
 type Entry struct {
 	Header tar.Header
 	Offset int64
@@ -204,7 +217,11 @@ func New(ra io.ReaderAt, size int64) (*FS, error) {
 
 	// TODO: Consider caching this across builds.
 	r := io.NewSectionReader(ra, 0, size)
-	cr := &countReader{bufio.NewReaderSize(r, 1<<20), 0}
+
+	br := pooledBufioReader(r)
+	defer readerPool.Put(br)
+
+	cr := &countReader{br, 0}
 	tr := tar.NewReader(cr)
 	for {
 		hdr, err := tr.Next()

pkg/build/build_implementation.go

@@ -25,6 +25,7 @@ import (
 	"os"
 	"path/filepath"
 	"runtime"
+	"sync"
 
 	gzip "github.com/klauspost/pgzip"
 	"github.com/sigstore/cosign/v2/pkg/oci"
@@ -48,12 +49,33 @@ import (
 // concurrent builds on giant machines, and uses only 1 core on tiny machines.
 var pgzipThreads = min(runtime.GOMAXPROCS(0), 8)
 
-func min(l, r int) int {
-	if l < r {
-		return l
-	}
+var pgzipPool = sync.Pool{
+	New: func() interface{} {
+		zw := gzip.NewWriter(nil)
+		if err := zw.SetConcurrency(1<<20, pgzipThreads); err != nil {
+			// This should never happen.
+			panic(fmt.Errorf("tried to set pgzip concurrency to %d: %w", pgzipThreads, err))
+		}
+		return zw
+	},
+}
+
+func pooledGzipWriter(w io.Writer) *gzip.Writer {
+	zw := pgzipPool.Get().(*gzip.Writer)
+	zw.Reset(w)
+	return zw
+}
 
-	return r
+var bufioPool = sync.Pool{
+	New: func() interface{} {
+		return bufio.NewWriterSize(nil, 1<<22)
+	},
+}
+
+func pooledBufioWriter(w io.Writer) *bufio.Writer {
+	bw := bufioPool.Get().(*bufio.Writer)
+	bw.Reset(w)
+	return bw
 }
 
 // BuildTarball takes the fully populated working directory and saves it to
@@ -85,11 +107,11 @@ func (bc *Context) BuildTarball(ctx context.Context) (string, hash.Hash, hash.Ha
 
 	digest := sha256.New()
 
-	buf := bufio.NewWriterSize(outfile, 1<<22)
-	gzw := gzip.NewWriter(io.MultiWriter(digest, buf))
-	if err := gzw.SetConcurrency(1<<20, pgzipThreads); err != nil {
-		return "", nil, nil, 0, fmt.Errorf("tried to set pgzip concurrency to %d: %w", pgzipThreads, err)
-	}
+	buf := pooledBufioWriter(outfile)
+	defer bufioPool.Put(buf)
+
+	gzw := pooledGzipWriter(io.MultiWriter(digest, buf))
+	defer pgzipPool.Put(gzw)
 
 	diffid := sha256.New()