Add support for custom wcow and lcow sandbox scratch sizes

Signed-off-by: Justin Terry (VM) <juterry@microsoft.com>

snapshots/lcow/lcow.go

@@ -57,6 +57,10 @@ func init() {
 	})
 }
 
+const (
+	rootfsSizeLabel = "containerd.io/snapshot/io.microsoft.container.storage.rootfs.size-gb"
+)
+
 type snapshotter struct {
 	root string
 	ms   *storage.MetaStore
@@ -321,7 +325,21 @@ func (s *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, k
 			return nil, err
 		}
 
-		scratchSource, err := s.openOrCreateScratch(ctx)
+		var snapshotInfo snapshots.Info
+		for _, o := range opts {
+			o(&snapshotInfo)
+		}
+
+		var sizeGB int
+		if sizeGBstr, ok := snapshotInfo.Labels[rootfsSizeLabel]; ok {
+			i32, err := strconv.ParseInt(sizeGBstr, 10, 32)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to parse annotation %q=%q", rootfsSizeLabel, sizeGBstr)
+			}
+			sizeGB = int(i32)
+		}
+
+		scratchSource, err := s.openOrCreateScratch(ctx, sizeGB)
 		if err != nil {
 			return nil, err
 		}
@@ -352,19 +370,24 @@ func (s *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, k
 	return s.mounts(newSnapshot), nil
 }
 
-func (s *snapshotter) openOrCreateScratch(ctx context.Context) (_ *os.File, err error) {
+func (s *snapshotter) openOrCreateScratch(ctx context.Context, sizeGB int) (_ *os.File, err error) {
 	// Create the scratch.vhdx cache file if it doesn't already exit.
 	s.scratchLock.Lock()
 	defer s.scratchLock.Unlock()
 
-	scratchFinalPath := filepath.Join(s.root, "scratch.vhdx")
+	vhdFileName := "scratch.vhdx"
+	if sizeGB > 0 {
+		vhdFileName = fmt.Sprintf("scratch_%d.vhdx", sizeGB)
+	}
+
+	scratchFinalPath := filepath.Join(s.root, vhdFileName)
 	scratchSource, err := os.OpenFile(scratchFinalPath, os.O_RDONLY, 0700)
 	if err != nil {
 		if !os.IsNotExist(err) {
-			return nil, errors.Wrap(err, "failed to open scratch.vhdx for read")
+			return nil, errors.Wrapf(err, "failed to open vhd %s for read", vhdFileName)
 		}
 
-		log.G(ctx).Debug("scratch.vhdx not found, creating a new one")
+		log.G(ctx).Debugf("vhd %s not found, creating a new one", vhdFileName)
 
 		// Golang logic for ioutil.TempFile without the file creation
 		r := uint32(time.Now().UnixNano() + int64(os.Getpid()))
@@ -380,7 +403,12 @@ func (s *snapshotter) openOrCreateScratch(ctx context.Context) (_ *os.File, err
 			LogFormat: runhcs.JSON,
 			Owner:     "containerd",
 		}
-		if err := rhcs.CreateScratch(ctx, scratchTempPath); err != nil {
+
+		opt := runhcs.CreateScratchOpts{
+			SizeGB: sizeGB,
+		}
+
+		if err := rhcs.CreateScratchWithOpts(ctx, scratchTempPath, &opt); err != nil {
 			_ = os.Remove(scratchTempPath)
 			return nil, errors.Wrapf(err, "failed to create '%s' temp file", scratchTempName)
 		}

snapshots/windows/windows.go

@@ -23,6 +23,7 @@ import (
 	"encoding/json"
 	"os"
 	"path/filepath"
+	"strconv"
 	"strings"
 
 	winfs "github.com/Microsoft/go-winio/pkg/fs"
@@ -48,6 +49,10 @@ func init() {
 	})
 }
 
+const (
+	rootfsSizeLabel = "containerd.io/snapshot/io.microsoft.container.storage.rootfs.size-gb"
+)
+
 type snapshotter struct {
 	root string
 	info hcsshim.DriverInfo
@@ -332,7 +337,26 @@ func (s *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, k
 			return nil, errors.Wrap(err, "failed to create sandbox layer")
 		}
 
-		// TODO(darrenstahlmsft): Allow changing sandbox size
+		var snapshotInfo snapshots.Info
+		for _, o := range opts {
+			o(&snapshotInfo)
+		}
+
+		var sizeGB int
+		if sizeGBstr, ok := snapshotInfo.Labels[rootfsSizeLabel]; ok {
+			i32, err := strconv.ParseInt(sizeGBstr, 10, 32)
+			if err != nil {
+				return nil, errors.Wrapf(err, "failed to parse annotation %q=%q", rootfsSizeLabel, sizeGBstr)
+			}
+			sizeGB = int(i32)
+		}
+
+		if sizeGB > 0 {
+			const gbToByte = 1024 * 1024 * 1024
+			if err := hcsshim.ExpandSandboxSize(s.info, newSnapshot.ID, uint64(gbToByte*sizeGB)); err != nil {
+				return nil, errors.Wrapf(err, "failed to expand scratch size to %d GB", sizeGB)
+			}
+		}
 	}
 
 	if err := t.Commit(); err != nil {