Fix warmer memory leak.

pkg/cache/warm.go

@@ -17,7 +17,6 @@ limitations under the License.
 package cache
 
 import (
-	"bytes"
 	"fmt"
 	"io"
 	"net/http"
@@ -58,35 +57,11 @@ func WarmCache(opts *config.WarmerOptions) error {
 
 	errs := 0
 	for _, img := range images {
-		tarBuf := new(bytes.Buffer)
-		manifestBuf := new(bytes.Buffer)
-
-		cw := &Warmer{
-			Remote:         remote.RetrieveRemoteImage,
-			Local:          LocalSource,
-			TarWriter:      tarBuf,
-			ManifestWriter: manifestBuf,
-		}
-
-		digest, err := cw.Warm(img, opts)
+		err := warmToFile(cacheDir, img, opts)
 		if err != nil {
-			if !IsAlreadyCached(err) {
-				logrus.Warnf("Error while trying to warm image: %v %v", img, err)
-				errs++
-			}
-
-			continue
-		}
-
-		cachePath := path.Join(cacheDir, digest.String())
-
-		if err := writeBufsToFile(cachePath, tarBuf, manifestBuf); err != nil {
-			logrus.Warnf("Error while writing %v to cache: %v", img, err)
+			logrus.Warnf("Error while trying to warm image: %v %v", img, err)
 			errs++
-			continue
 		}
-
-		logrus.Debugf("Wrote %s to cache", img)
 	}
 
 	if len(images) == errs {
@@ -96,22 +71,53 @@ func WarmCache(opts *config.WarmerOptions) error {
 	return nil
 }
 
-func writeBufsToFile(cachePath string, tarBuf, manifestBuf *bytes.Buffer) error {
-	f, err := os.Create(cachePath)
+// Download image in temporary files then move files to final destination
+func warmToFile(cacheDir, img string, opts *config.WarmerOptions) error {
+	f, err := os.CreateTemp(cacheDir, "warmingImage.*")
 	if err != nil {
 		return err
 	}
+	// defer called in reverse order
+	defer os.Remove(f.Name())
 	defer f.Close()
 
-	if _, err := f.Write(tarBuf.Bytes()); err != nil {
-		return errors.Wrap(err, "Failed to save tar to file")
+	mtfsFile, err := os.CreateTemp(cacheDir, "warmingManifest.*")
+	if err != nil {
+		return err
+	}
+	defer os.Remove(mtfsFile.Name())
+	defer mtfsFile.Close()
+
+	cw := &Warmer{
+		Remote:         remote.RetrieveRemoteImage,
+		Local:          LocalSource,
+		TarWriter:      f,
+		ManifestWriter: mtfsFile,
 	}
 
-	mfstPath := cachePath + ".json"
-	if err := os.WriteFile(mfstPath, manifestBuf.Bytes(), 0666); err != nil {
-		return errors.Wrap(err, "Failed to save manifest to file")
+	digest, err := cw.Warm(img, opts)
+	if err != nil {
+		if !IsAlreadyCached(err) {
+			logrus.Warnf("Error while trying to warm image: %v %v", img, err)
+		}
+
+		return err
+	}
+
+	finalCachePath := path.Join(cacheDir, digest.String())
+	finalMfstPath := finalCachePath + ".json"
+
+	err = os.Rename(f.Name(), finalCachePath)
+	if err != nil {
+		return err
+	}
+
+	err = os.Rename(mtfsFile.Name(), finalMfstPath)
+	if err != nil {
+		return errors.Wrap(err, "Failed to rename manifest file")
 	}
 
+	logrus.Debugf("Wrote %s to cache", img)
 	return nil
 }
 

scripts/boxed_warm_in_docker.sh

@@ -0,0 +1,31 @@
+#!/bin/bash
+
+# Copyright 2018 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Test the warmer in boxed memory conditions.
+# Attempt to run the warmer inside a container limited to 16MB of RAM. Use gcr.io/kaniko-project/warmer:latest image."
+# Example: ./boxed_warm_in_docker.sh --image debian:trixie-slim
+# 
+set -e
+
+rc=0
+docker run \
+	--memory=16m --memory-swappiness=0 \
+        gcr.io/kaniko-project/warmer:latest \
+	"$@" || rc=$?
+
+>&2 echo "RC=$rc"
+exit $rc
+