Auto calculate image size on actions (#122)

Author: Itxaka

go.mod

@@ -12,7 +12,7 @@ require (
 	github.com/hashicorp/go-multierror v1.1.1
 	github.com/jaypipes/ghw v0.12.0
 	github.com/joho/godotenv v1.5.1
-	github.com/kairos-io/kairos-sdk v0.0.11
+	github.com/kairos-io/kairos-sdk v0.0.12
 	github.com/labstack/echo/v4 v4.10.2
 	github.com/mitchellh/mapstructure v1.5.0
 	github.com/mudler/go-nodepair v0.0.0-20221223092639-ba399a66fdfb

go.sum

@@ -371,6 +371,8 @@ github.com/kairos-io/kairos-sdk v0.0.10 h1:TUgrGSGP6Z1CPfA4gjmbb+cCaJg1OR18c+LD+
 github.com/kairos-io/kairos-sdk v0.0.10/go.mod h1:AK9poWisuhnzri04diXnTG8L7EkOSUArSeeDn2LYFU0=
 github.com/kairos-io/kairos-sdk v0.0.11 h1:+EWuO4gWMzBa81s8gbF1L7wOvEjbph1z8UkfUrMyvxc=
 github.com/kairos-io/kairos-sdk v0.0.11/go.mod h1:AK9poWisuhnzri04diXnTG8L7EkOSUArSeeDn2LYFU0=
+github.com/kairos-io/kairos-sdk v0.0.12 h1:+uibTjjsAh8klupXHWpWse2AKww6p8ROu8Ii28t3k48=
+github.com/kairos-io/kairos-sdk v0.0.12/go.mod h1:AK9poWisuhnzri04diXnTG8L7EkOSUArSeeDn2LYFU0=
 github.com/kbinani/screenshot v0.0.0-20210720154843-7d3a670d8329 h1:qq2nCpSrXrmvDGRxW0ruW9BVEV1CN2a9YDOExdt+U0o=
 github.com/kbinani/screenshot v0.0.0-20210720154843-7d3a670d8329/go.mod h1:2VPVQDR4wO7KXHwP+DAypEy67rXf+okUx2zjgpCxZw4=
 github.com/kendru/darwin/go/depgraph v0.0.0-20221105232959-877d6a81060c h1:eKb4PqwAMhlqwXw0W3atpKaYaPGlXE/Fwh+xpCEYaPk=

pkg/config/spec.go

@@ -18,15 +18,16 @@ package config
 
 import (
 	"fmt"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
-	"github.com/kairos-io/kairos-agent/v2/pkg/utils/partitions"
+	"io/fs"
 	"path/filepath"
 	"reflect"
 	"strings"
 
 	"github.com/kairos-io/kairos-agent/v2/internal/common"
 	"github.com/kairos-io/kairos-agent/v2/pkg/constants"
 	v1 "github.com/kairos-io/kairos-agent/v2/pkg/types/v1"
+	"github.com/kairos-io/kairos-agent/v2/pkg/utils/fs"
+	"github.com/kairos-io/kairos-agent/v2/pkg/utils/partitions"
 	"github.com/mitchellh/mapstructure"
 	"github.com/sanity-io/litter"
 	"github.com/sirupsen/logrus"
@@ -95,7 +96,19 @@ func NewInstallSpec(cfg *Config) *v1.InstallSpec {
 		Recovery:  recoveryImg,
 		Passive:   passiveImg,
 	}
-	// Calculate the partitions sizes automatically based on the images set size
+
+	// Get the actual source size to calculate the image size and partitions size
+	size, err := GetSourceSize(cfg, spec.Active.Source)
+	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	} else {
+		cfg.Logger.Infof("Setting image size to %dMb", size)
+		spec.Active.Size = uint(size)
+		spec.Passive.Size = uint(size)
+		spec.Recovery.Size = uint(size)
+	}
+
+	// Calculate the partitions afterwards so they use the image sizes for the final partition sizes
 	spec.Partitions = NewInstallElementalPartitions(spec)
 
 	return spec
@@ -237,13 +250,26 @@ func NewUpgradeSpec(cfg *Config) (*v1.UpgradeSpec, error) {
 		}
 	}
 
-	return &v1.UpgradeSpec{
+	spec := &v1.UpgradeSpec{
 		Active:     active,
 		Recovery:   recovery,
 		Passive:    passive,
 		Partitions: ep,
 		State:      installState,
-	}, nil
+	}
+
+	// Get the actual source size to calculate the image size and partitions size
+	size, err := GetSourceSize(cfg, spec.Active.Source)
+	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	} else {
+		cfg.Logger.Infof("Setting image size to %dMb", size)
+		// On upgrade only the active or recovery will be upgraded, so we dont need to override passive
+		spec.Active.Size = uint(size)
+		spec.Recovery.Size = uint(size)
+	}
+
+	return spec, nil
 }
 
 // NewResetSpec returns a ResetSpec struct all based on defaults and current host state
@@ -343,7 +369,7 @@ func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 	}
 
 	activeFile := filepath.Join(ep.State.MountPoint, "cOS", constants.ActiveImgFile)
-	return &v1.ResetSpec{
+	spec := &v1.ResetSpec{
 		Target:           target,
 		Partitions:       ep,
 		Efi:              efiExists,
@@ -367,7 +393,19 @@ func NewResetSpec(cfg *Config) (*v1.ResetSpec, error) {
 			FS:     constants.LinuxImgFs,
 		},
 		State: installState,
-	}, nil
+	}
+
+	// Get the actual source size to calculate the image size and partitions size
+	size, err := GetSourceSize(cfg, spec.Active.Source)
+	if err != nil {
+		cfg.Logger.Warnf("Failed to infer size for images: %s", err.Error())
+	} else {
+		cfg.Logger.Infof("Setting image size to %dMb", size)
+		spec.Active.Size = uint(size)
+		spec.Passive.Size = uint(size)
+	}
+
+	return spec, nil
 }
 
 // ReadResetSpecFromConfig will return a proper v1.ResetSpec based on an agent Config
@@ -400,6 +438,44 @@ func ReadInstallSpecFromConfig(c *Config) (*v1.InstallSpec, error) {
 	return installSpec, nil
 }
 
+// GetSourceSize will try to gather the actual size of the source
+// Useful to create the exact size of images and by side effect the partition size
+// This helps adjust the size to be juuuuust right.
+// It can still be manually override from the cloud config by setting all values manually
+// But by default it should adjust the sizes properly
+func GetSourceSize(config *Config, source *v1.ImageSource) (int64, error) {
+	var size int64
+	var err error
+
+	switch {
+	case source.IsDocker():
+		size, err = config.ImageExtractor.GetOCIImageSize(source.Value(), config.Platform.String())
+	case source.IsDir():
+		err = fsutils.WalkDirFs(config.Fs, source.Value(), func(path string, d fs.DirEntry, err error) error {
+			if err != nil {
+				return err
+			}
+			info, err := d.Info()
+			if err != nil {
+				return err
+			}
+			size += info.Size()
+			return nil
+		})
+
+	case source.IsFile():
+		file, err := config.Fs.Stat(source.Value())
+		if err == nil {
+			size = file.Size()
+		}
+	}
+	// Normalize size to Mb before returning and add 100Mb to round the size from bytes to mb+extra files like grub stuff
+	if size != 0 {
+		size = (size / 1000 / 1000) + 100
+	}
+	return size, err
+}
+
 // ReadUpgradeSpecFromConfig will return a proper v1.UpgradeSpec based on an agent Config
 func ReadUpgradeSpecFromConfig(c *Config) (*v1.UpgradeSpec, error) {
 	sp, err := ReadSpecFromCloudConfig(c, "upgrade")

pkg/elemental/elemental.go

@@ -252,7 +252,7 @@ func (e Elemental) UnmountImage(img *v1.Image) error {
 
 // CreateFileSystemImage creates the image file for config.target
 func (e Elemental) CreateFileSystemImage(img *v1.Image) error {
-	e.config.Logger.Infof("Creating file system image %s", img.File)
+	e.config.Logger.Infof("Creating file system image %s with size %dMb", img.File, img.Size)
 	err := fsutils.MkdirAll(e.config.Fs, filepath.Dir(img.File), cnst.DirPerm)
 	if err != nil {
 		return err

pkg/elemental/elemental_test.go

@@ -387,13 +387,13 @@ var _ = Describe("Elemental", Label("elemental"), func() {
 						"mkpart", "oem", "ext4", "133120", "264191",
 					}, {"mkfs.ext4", "-L", "COS_OEM", "/some/device2"}, {
 						"parted", "--script", "--machine", "--", "/some/device", "unit", "s",
-						"mkpart", "recovery", "ext4", "264192", "6760447",
+						"mkpart", "recovery", "ext4", "264192", "468991",
 					}, {"mkfs.ext4", "-L", "COS_RECOVERY", "/some/device3"}, {
 						"parted", "--script", "--machine", "--", "/some/device", "unit", "s",
-						"mkpart", "state", "ext4", "6760448", "19548159",
+						"mkpart", "state", "ext4", "468992", "673791",
 					}, {"mkfs.ext4", "-L", "COS_STATE", "/some/device4"}, {
 						"parted", "--script", "--machine", "--", "/some/device", "unit", "s",
-						"mkpart", "persistent", "ext4", "19548160", "100%",
+						"mkpart", "persistent", "ext4", "673792", "100%",
 					}, {"mkfs.ext4", "-L", "COS_PERSISTENT", "/some/device5"},
 				}
 

pkg/types/v1/image.go

@@ -22,6 +22,7 @@ import (
 
 type ImageExtractor interface {
 	ExtractImage(imageRef, destination, platformRef string) error
+	GetOCIImageSize(imageRef, platformRef string) (int64, error)
 }
 
 type OCIImageExtractor struct{}
@@ -31,3 +32,7 @@ var _ ImageExtractor = OCIImageExtractor{}
 func (e OCIImageExtractor) ExtractImage(imageRef, destination, platformRef string) error {
 	return utils.ExtractOCIImage(imageRef, destination, platformRef)
 }
+
+func (e OCIImageExtractor) GetOCIImageSize(imageRef, platformRef string) (int64, error) {
+	return utils.GetOCIImageSize(imageRef, platformRef)
+}

pkg/utils/fs/fs.go

@@ -20,6 +20,7 @@ limitations under the License.
 package fsutils
 
 import (
+	"errors"
 	"io/fs"
 	"os"
 	"path/filepath"
@@ -198,7 +199,7 @@ func WalkDirFs(fs v1.FS, root string, fn fs.WalkDirFunc) error {
 	} else {
 		err = walkDir(fs, root, &statDirEntry{info}, fn)
 	}
-	if err == filepath.SkipDir {
+	if errors.Is(err, filepath.SkipDir) {
 		return nil
 	}
 	return err
@@ -225,7 +226,7 @@ func walkDir(fs v1.FS, path string, d fs.DirEntry, walkDirFn fs.WalkDirFunc) err
 	for _, d1 := range dirs {
 		path1 := filepath.Join(path, d1.Name())
 		if err := walkDir(fs, path1, d1, walkDirFn); err != nil {
-			if err == filepath.SkipDir {
+			if errors.Is(err, filepath.SkipDir) {
 				break
 			}
 			return err

tests/mocks/extractor_mock.go

@@ -23,6 +23,10 @@ type FakeImageExtractor struct {
 	SideEffect func(imageRef, destination, platformRef string) error
 }
 
+func (f FakeImageExtractor) GetOCIImageSize(imageRef, platformRef string) (int64, error) {
+	return 0, nil
+}
+
 var _ v1.ImageExtractor = FakeImageExtractor{}
 
 func NewFakeImageExtractor(logger v1.Logger) *FakeImageExtractor {