Fix upgrade size calculation inside k8s (#537)

Co-authored-by: Dimitris Karakasilis <dimitris@karakasilis.me>
(cherry picked from commit 7c6c195bd5)

pkg/config/spec.go

@@ -808,11 +808,33 @@ func GetSourceSize(config *Config, source *v1.ImageSource) (int64, error) {
 		size = int64(float64(size) * 2.5)
 	case source.IsDir():
 		filesVisited = make(map[string]bool, 30000) // An Ubuntu system has around 27k files. This improves performance by not having to resize the map for every file visited
+		// In kubernetes we use the suc script to upgrade (https://github.com/kairos-io/packages/blob/main/packages/system/suc-upgrade/suc-upgrade.sh)
+		// , which mounts the host root into $HOST_DIR
+		// we should skip that dir when calculating the size as we would be doubling the calculated size
+		// Plus we will hit the usual things when checking a running system. Processes that go away, tmpfiles, etc...
 
+		// This is always set for pods running under kubernetes
+		underKubernetes := os.Getenv("KUBERNETES_SERVICE_HOST")
+		// Try to get the HOST_DIR in case we are not using the default one
+		hostDir := os.Getenv("HOST_DIR")
+		// If we are under kubernetes but the HOST_DIR var is empty, default to /host as system-upgrade-controller mounts
+		// the host in that dir by default
+		if underKubernetes != "" && hostDir == "" {
+			hostDir = "/host"
+		}
 		err = fsutils.WalkDirFs(config.Fs, source.Value(), func(path string, d fs.DirEntry, err error) error {
-			v := getSize(&size, filesVisited, path, d, err)
 
-			return v
+			// If its empty we are just not setting it, so probably out of the k8s upgrade path
+			if hostDir != "" && strings.HasPrefix(path, hostDir) {
+				config.Logger.Logger.Debug().Str("path", path).Str("hostDir", hostDir).Msg("Skipping file as it is a host directory")
+			} else if strings.HasPrefix(path, "/proc") || strings.HasPrefix(path, "/dev") || strings.HasPrefix(path, "/run") {
+				config.Logger.Logger.Debug().Str("path", path).Str("hostDir", hostDir).Msg("Skipping dir as it is a system directory (/proc, /dev or /run)")
+			} else {
+				v := getSize(&size, filesVisited, path, d, err)
+				return v
+			}
+
+			return nil
 		})
 
 	case source.IsFile():

pkg/config/spec_test.go

@@ -17,6 +17,7 @@ limitations under the License.
 package config_test
 
 import (
+	"bytes"
 	"fmt"
 	"os"
 	"path/filepath"
@@ -653,19 +654,23 @@ func createFileOfSizeInMB(filename string, sizeInMB int) error {
 	return nil
 }
 
-var _ = Describe("GetSourceSize", func() {
+var _ = Describe("GetSourceSize", Label("GetSourceSize"), func() {
 	var tempDir string
 	var tempFilePath string
 	var err error
 	var logger sdkTypes.KairosLogger
 	var conf *config.Config
 	var imageSource *v1.ImageSource
+	var memLog bytes.Buffer
 
 	BeforeEach(func() {
 		tempDir, err = os.MkdirTemp("/tmp", "kairos-test")
 		Expect(err).To(BeNil())
 
-		logger = sdkTypes.NewNullLogger()
+		//logger = sdkTypes.NewNullLogger()
+		memLog = bytes.Buffer{}
+		logger = sdkTypes.NewBufferLogger(&memLog)
+		logger.SetLevel("debug")
 		conf = config.NewConfig(
 			config.WithLogger(logger),
 		)
@@ -678,6 +683,7 @@ var _ = Describe("GetSourceSize", func() {
 	})
 
 	AfterEach(func() {
+		fmt.Println(memLog.String())
 		defer os.RemoveAll(tempDir)
 	})
 
@@ -693,4 +699,35 @@ var _ = Describe("GetSourceSize", func() {
 		Expect(err).ToNot(HaveOccurred())
 		Expect(sizeAfter).To(Equal(sizeBefore))
 	})
+	It("Skips the kubernetes host dir when calculating the sizes if set", func() {
+		sizeBefore, err := config.GetSourceSize(conf, imageSource)
+		Expect(err).To(BeNil())
+		Expect(sizeBefore).ToNot(BeZero())
+
+		Expect(os.Mkdir(filepath.Join(tempDir, "host"), os.ModePerm)).ToNot(HaveOccurred())
+		Expect(createFileOfSizeInMB(filepath.Join(tempDir, "host", "what.txt"), 200)).ToNot(HaveOccurred())
+		// Set env var like the suc upgrade does
+		Expect(os.Setenv("HOST_DIR", filepath.Join(tempDir, "host"))).ToNot(HaveOccurred())
+
+		sizeAfter, err := config.GetSourceSize(conf, imageSource)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(sizeAfter).To(Equal(sizeBefore))
+	})
+	It("Counts the kubernetes host dir when calculating the sizes if not set", func() {
+		sizeBefore, err := config.GetSourceSize(conf, imageSource)
+		Expect(err).To(BeNil())
+		Expect(sizeBefore).ToNot(BeZero())
+
+		Expect(os.Mkdir(filepath.Join(tempDir, "host"), os.ModePerm)).ToNot(HaveOccurred())
+		Expect(createFileOfSizeInMB(filepath.Join(tempDir, "host", "what.txt"), 200)).ToNot(HaveOccurred())
+
+		sizeAfter, err := config.GetSourceSize(conf, imageSource)
+		Expect(err).ToNot(HaveOccurred())
+		Expect(sizeAfter).ToNot(Equal(sizeBefore))
+		Expect(sizeAfter).ToNot(BeZero())
+		// Size is 2 files of 200 + 100Mb on top, normalized from bytes to MB
+		// So take those 200Mb, converts to bytes by multiplying them (400*1024*1024), then back to MB by dividing
+		// what we get (/1000/1000) then we finish by adding and extra 100MB on top, like the GetSourceSize does internally
+		Expect(sizeAfter).To(Equal(int64((400 * 1024 * 1024 / 1000 / 1000) + 100)))
+	})
 })