Run firecracker workspace conversion in action cgroup

enterprise/server/remote_execution/cgroup/cgroup.go

@@ -113,6 +113,16 @@ func ParentEnabledControllers(path string) (map[string]bool, error) {
 	return enabled, nil
 }
 
+// WrapCommand takes a command and returns a new command which executes the
+// original command inside the given cgroup.
+func WrapCommand(cgroup, command string, arguments ...string) (wrappedCommand string, wrappedArgs []string) {
+	script := fmt.Sprintf(
+		`echo $$ > %q && exec "$0" "$@"`,
+		filepath.Join(RootPath, cgroup, "cgroup.procs"),
+	)
+	return "sh", append([]string{"-ec", script, command}, arguments...)
+}
+
 func settingsMap(s *scpb.CgroupSettings, blockDevice *block_io.Device) (map[string]string, error) {
 	m := map[string]string{}
 	if s == nil {

enterprise/server/remote_execution/containers/firecracker/firecracker.go

@@ -1211,7 +1211,7 @@ func (c *FirecrackerContainer) createWorkspaceImage(ctx context.Context, workspa
 		}
 		workspaceDiskSizeBytes := ext4.MinDiskImageSizeBytes + workspaceSizeBytes + *workspaceDiskSlackSpaceMB*1e6
 		workspaceDiskSizeBytes = alignToMultiple(workspaceDiskSizeBytes, int64(os.Getpagesize()))
-		if err := ext4.DirectoryToImage(ctx, workspaceDir, ext4ImagePath, workspaceDiskSizeBytes); err != nil {
+		if err := ext4.DirectoryToImage(ctx, workspaceDir, ext4ImagePath, workspaceDiskSizeBytes, c.cgroup()); err != nil {
 			return status.WrapError(err, "failed to convert workspace dir to ext4 image")
 		}
 	}
@@ -1505,8 +1505,15 @@ func (c *FirecrackerContainer) getConfig(ctx context.Context, rootFS, containerF
 	return cfg, nil
 }
 
+// cgroup returns the cgroup path relative to the cgroup root.
+func (c *FirecrackerContainer) cgroup() string {
+	return filepath.Join(c.cgroupParent, c.id)
+}
+
+// cgroupPath returns the absolute cgroup path, including the cgroupfs root
+// path.
 func (c *FirecrackerContainer) cgroupPath() string {
-	return filepath.Join("/sys/fs/cgroup", c.cgroupParent, c.id)
+	return filepath.Join("/sys/fs/cgroup", c.cgroup())
 }
 
 func (c *FirecrackerContainer) getJailerConfig(ctx context.Context, kernelImagePath string) (*fcclient.JailerConfig, error) {

enterprise/server/util/ext4/BUILD

@@ -9,6 +9,7 @@ go_library(
     target_compatible_with = ["@platforms//os:linux"],
     deps = select({
         "@io_bazel_rules_go//go/platform:linux": [
+            "//enterprise/server/remote_execution/cgroup",
             "//server/util/log",
             "//server/util/status",
             "//server/util/tracing",

enterprise/server/util/ext4/ext4.go

@@ -12,6 +12,7 @@ import (
 	"path/filepath"
 	"syscall"
 
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/cgroup"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
 	"github.com/buildbuddy-io/buildbuddy/server/util/tracing"
@@ -35,16 +36,19 @@ const (
 
 // DirectoryToImage creates an ext4 image of the specified size from inputDir
 // and writes it to outputFile.
-func DirectoryToImage(ctx context.Context, inputDir, outputFile string, sizeBytes int64) error {
+//
+// If cgroupName is nonempty then the conversion command is executed in the
+// given cgroup, which should be a path relative to the cgroupfs root. The
+// caller must ensure that the cgroup exists.
+func DirectoryToImage(ctx context.Context, inputDir, outputFile string, sizeBytes int64, cgroupName string) error {
 	if err := checkImageOutputPath(outputFile); err != nil {
 		return err
 	}
 
 	ctx, span := tracing.StartSpan(ctx)
 	defer span.End()
 
-	args := []string{
-		"/sbin/mke2fs",
+	command, args := "/sbin/mke2fs", []string{
 		"-L", "''",
 		"-N", "0",
 		"-O", "^64bit",
@@ -56,10 +60,12 @@ func DirectoryToImage(ctx context.Context, inputDir, outputFile string, sizeByte
 		outputFile,
 		fmt.Sprintf("%dK", sizeBytes/iecKilobyte),
 	}
-	cmd := exec.CommandContext(ctx, args[0], args[1:]...)
-	if out, err := cmd.CombinedOutput(); err != nil {
-		log.Errorf("Error running %q: %s %s", cmd.String(), err, out)
-		return status.InternalErrorf("%s: %s", err, out)
+	if cgroupName != "" {
+		command, args = cgroup.WrapCommand(cgroupName, command, args...)
+	}
+	cmd := exec.CommandContext(ctx, command, args...)
+	if b, err := cmd.CombinedOutput(); err != nil {
+		return status.InternalErrorf("mke2fs failed (%s): %q", err, string(b))
 	}
 	return nil
 }
@@ -140,14 +146,14 @@ func DiskSizeBytes(ctx context.Context, inputDir string) (int64, error) {
 
 // DirectoryToImageAutoSize is like DirectoryToImage, but it will attempt to
 // automatically pick a file size that is "big enough".
-func DirectoryToImageAutoSize(ctx context.Context, inputDir, outputFile string) error {
+func DirectoryToImageAutoSize(ctx context.Context, inputDir, outputFile, cgroup string) error {
 	dirSizeBytes, err := DiskSizeBytes(ctx, inputDir)
 	if err != nil {
 		return status.WrapError(err, "estimate disk usage")
 	}
 
 	imageSizeBytes := int64(float64(dirSizeBytes)*1.2) + MinDiskImageSizeBytes
-	return DirectoryToImage(ctx, inputDir, outputFile, imageSizeBytes)
+	return DirectoryToImage(ctx, inputDir, outputFile, imageSizeBytes, cgroup)
 }
 
 // isDirEmpty returns a bool indicating if a directory contains no files, or

enterprise/server/util/ociconv/ociconv.go

@@ -272,7 +272,10 @@ func convertContainerToExt4FS(ctx context.Context, workspaceDir, containerImage
 	}
 	defer f.Close()
 	imageFile := f.Name()
-	if err := ext4.DirectoryToImageAutoSize(ctx, rootFSDir, imageFile); err != nil {
+	// Keep the image conversion work in the current process cgroup (i.e.
+	// executor cgroup) for now.
+	cgroup := ""
+	if err := ext4.DirectoryToImageAutoSize(ctx, rootFSDir, imageFile, cgroup); err != nil {
 		return "", err
 	}
 	log.Debugf("Wrote container %q to image file: %q", containerImage, imageFile)