Firecracker: Use docker pull for images (#1565)

enterprise/server/remote_execution/containers/docker/docker.go

@@ -290,6 +290,10 @@ func (r *dockerCommandContainer) IsImageCached(ctx context.Context) (bool, error
 }
 
 func (r *dockerCommandContainer) PullImage(ctx context.Context, creds container.PullCredentials) error {
+	return PullImage(ctx, r.client, r.image, creds)
+}
+
+func PullImage(ctx context.Context, client *dockerclient.Client, image string, creds container.PullCredentials) error {
 	if !creds.IsEmpty() {
 		authCfg := dockertypes.AuthConfig{
 			Username: creds.Username,
@@ -299,7 +303,7 @@ func (r *dockerCommandContainer) PullImage(ctx context.Context, creds container.
 		if err != nil {
 			return err
 		}
-		rc, err := r.client.ImagePull(ctx, r.image, dockertypes.ImagePullOptions{
+		rc, err := client.ImagePull(ctx, image, dockertypes.ImagePullOptions{
 			RegistryAuth: auth,
 		})
 		if err != nil {
@@ -315,18 +319,36 @@ func (r *dockerCommandContainer) PullImage(ctx context.Context, creds container.
 	// TODO: find a way to implement this without calling the Docker CLI.
 	// Currently it's a bit involved to replicate the exact protocols that the
 	// CLI uses to pull images.
-	cmd := exec.CommandContext(ctx, "docker", "pull", r.image)
+	cmd := exec.CommandContext(ctx, "docker", "pull", image)
 	stderr := &bytes.Buffer{}
 	cmd.Stderr = stderr
 	if err := cmd.Run(); err != nil {
 		return wrapDockerErr(
 			err,
-			fmt.Sprintf("docker pull %q: %s -- stderr:\n%s", r.image, err, string(stderr.Bytes())),
+			fmt.Sprintf("docker pull %q: %s -- stderr:\n%s", image, err, string(stderr.Bytes())),
 		)
 	}
 	return nil
 }
 
+// SaveImage saves an image from the Docker cache to a tarball file at the given
+// path. The image must have already been pulled, otherwise the operation will
+// fail.
+func SaveImage(ctx context.Context, client *dockerclient.Client, imageRef, path string) error {
+	r, err := client.ImageSave(ctx, []string{imageRef})
+	if err != nil {
+		return err
+	}
+	defer r.Close()
+	f, err := os.Create(path)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+	_, err = io.Copy(f, r)
+	return err
+}
+
 func generateContainerName() (string, error) {
 	suffix, err := random.RandomString(20)
 	if err != nil {

enterprise/server/remote_execution/containers/firecracker/BUILD

@@ -13,7 +13,9 @@ go_library(
     ],
     importpath = "github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/containers/firecracker",
     visibility = ["//visibility:public"],
-    deps = select({
+    deps = [
+        "@com_github_docker_docker//client:go_default_library",
+    ] + select({
         "@io_bazel_rules_go//go/platform:darwin": [
             "//enterprise/server/remote_execution/container",
             "//proto:remote_execution_go_proto",

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

@@ -1,9 +1,20 @@
 package firecracker
 
+import (
+	dockerclient "github.com/docker/docker/client"
+)
+
 type ContainerOpts struct {
 	// The OCI container image. ex "alpine:latest".
 	ContainerImage string
 
+	// DockerClient can optionally be specified to pull container images via
+	// Docker. This is useful for de-duping in-flight image pull operations and
+	// making use of the local Docker cache for images. If not specified, images
+	// will be pulled directly by skopeo and no image pull de-duping will be
+	// performed.
+	DockerClient *dockerclient.Client
+
 	// The action directory with inputs / outputs.
 	ActionWorkingDirectory string
 

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

@@ -45,6 +45,7 @@ import (
 	repb "github.com/buildbuddy-io/buildbuddy/proto/remote_execution"
 	vmxpb "github.com/buildbuddy-io/buildbuddy/proto/vmexec"
 	vmfspb "github.com/buildbuddy-io/buildbuddy/proto/vmvfs"
+	dockerclient "github.com/docker/docker/client"
 	fcclient "github.com/firecracker-microvm/firecracker-go-sdk"
 	fcmodels "github.com/firecracker-microvm/firecracker-go-sdk/client/models"
 )
@@ -276,6 +277,10 @@ type FirecrackerContainer struct {
 	workspaceFSPath  string // the path to the workspace ext4 image
 	containerFSPath  string // the path to the container ext4 image
 
+	// dockerClient is used to optimize image pulls by reusing image layers from
+	// the Docker cache as well as deduping multiple requests for the same image.
+	dockerClient *dockerclient.Client
+
 	// when VFS is enabled, this contains the layout for the next execution
 	fsLayout  *container.FileSystemLayout
 	vfsServer *vfs_server.Server
@@ -354,6 +359,7 @@ func NewContainer(env environment.Env, imageCacheAuth *container.ImageCacheAuthe
 			DebugMode:        opts.DebugMode,
 		},
 		jailerRoot:         opts.JailerRoot,
+		dockerClient:       opts.DockerClient,
 		containerImage:     opts.ContainerImage,
 		actionWorkingDir:   opts.ActionWorkingDirectory,
 		env:                env,
@@ -1230,7 +1236,7 @@ func (c *FirecrackerContainer) PullImage(ctx context.Context, creds container.Pu
 	if c.containerFSPath != "" {
 		return nil
 	}
-	containerFSPath, err := containerutil.CreateDiskImage(ctx, c.jailerRoot, c.containerImage, creds)
+	containerFSPath, err := containerutil.CreateDiskImage(ctx, c.dockerClient, c.jailerRoot, c.containerImage, creds)
 	if err != nil {
 		return err
 	}

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

@@ -34,7 +34,7 @@ import (
 )
 
 const (
-	imageWithDockerInstalled = "gcr.io/flame-public/bb-devtools@sha256:16ef96fe7efc61b8b703ee45c4dc2ef53d1df9011e72c362e61b40f3d03780f4"
+	imageWithDockerInstalled = "gcr.io/flame-public/executor-docker-default:enterprise-v1.6.0"
 )
 
 func getTestEnv(ctx context.Context, t *testing.T) *testenv.TestEnv {
@@ -525,11 +525,12 @@ func TestFirecrackerExecWithDockerFromSnapshot(t *testing.T) {
 		ContainerImage:         imageWithDockerInstalled,
 		ActionWorkingDirectory: workDir,
 		NumCPUs:                1,
-		MemSizeMB:              200, // small to make snapshotting faster.
+		MemSizeMB:              2500,
 		InitDockerd:            true,
 		EnableNetworking:       true,
 		DiskSlackSpaceMB:       100,
 		JailerRoot:             tempJailerRoot(t),
+		DebugMode:              true,
 	}
 	c, err := firecracker.NewContainer(env, cacheAuth, opts)
 	if err != nil {

enterprise/server/remote_execution/runner/runner.go

@@ -795,6 +795,7 @@ func (p *Pool) newContainer(ctx context.Context, props *platform.Properties, tas
 		sizeEstimate := tasksize.Estimate(task)
 		opts := firecracker.ContainerOpts{
 			ContainerImage:         props.ContainerImage,
+			DockerClient:           p.dockerClient,
 			ActionWorkingDirectory: p.hostBuildRoot(),
 			NumCPUs:                int64(math.Max(1.0, float64(sizeEstimate.GetEstimatedMilliCpu())/1000)),
 			MemSizeMB:              int64(math.Max(1.0, float64(sizeEstimate.GetEstimatedMemoryBytes())/1e6)),

enterprise/server/util/container/BUILD

@@ -7,9 +7,11 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//enterprise/server/remote_execution/container",
+        "//enterprise/server/remote_execution/containers/docker",
         "//enterprise/server/util/ext4",
         "//server/util/disk",
         "//server/util/log",
         "//server/util/status",
+        "@com_github_docker_docker//client:go_default_library",
     ],
 )

enterprise/server/util/container/container.go

@@ -11,10 +11,13 @@ import (
 	"sort"
 
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/container"
+	"github.com/buildbuddy-io/buildbuddy/enterprise/server/remote_execution/containers/docker"
 	"github.com/buildbuddy-io/buildbuddy/enterprise/server/util/ext4"
 	"github.com/buildbuddy-io/buildbuddy/server/util/disk"
 	"github.com/buildbuddy-io/buildbuddy/server/util/log"
 	"github.com/buildbuddy-io/buildbuddy/server/util/status"
+
+	dockerclient "github.com/docker/docker/client"
 )
 
 const (
@@ -77,13 +80,19 @@ func CachedDiskImagePath(ctx context.Context, workspaceDir, containerImage strin
 	return diskImagePath, nil
 }
 
-// CreateDiskImage pulls the image from the container registry and creates an
-// ext4 disk image from the container image in the user's local cache directory.
-// If the image is already available locally, the image is not re-downloaded
-// from the registry, but the credentials are still authenticated with the
-// remote registry to ensure that the image can be accessed. The path to the
-// disk image is returned.
-func CreateDiskImage(ctx context.Context, workspaceDir, containerImage string, creds container.PullCredentials) (string, error) {
+// CreateDiskImage pulls the image from the container registry and exports an
+// ext4 disk image based on the container image to the configured cache
+// directory.
+//
+// If dockerClient is non-nil, then Docker will be used to pull and export
+// the image. This ensures that image pulls are de-duped and that image layers
+// which are already in the local Docker cache can be reused.
+//
+// If the image is already cached, the image is not re-downloaded from the
+// registry, but the credentials are still authenticated with the remote
+// registry to ensure that the image can be accessed. The path to the disk image
+// is returned.
+func CreateDiskImage(ctx context.Context, dockerClient *dockerclient.Client, workspaceDir, containerImage string, creds container.PullCredentials) (string, error) {
 	existingPath, err := CachedDiskImagePath(ctx, workspaceDir, containerImage)
 	if err != nil {
 		return "", err
@@ -115,7 +124,7 @@ func CreateDiskImage(ctx context.Context, workspaceDir, containerImage string, c
 	containerImagesPath := filepath.Join(workspaceDir, "executor", hashedContainerName)
 
 	// container not found -- write one!
-	tmpImagePath, err := convertContainerToExt4FS(ctx, workspaceDir, containerImage, creds)
+	tmpImagePath, err := convertContainerToExt4FS(ctx, dockerClient, workspaceDir, containerImage, creds)
 	if err != nil {
 		return "", err
 	}
@@ -139,7 +148,7 @@ func CreateDiskImage(ctx context.Context, workspaceDir, containerImage string, c
 // image from an OCI container image reference.
 // NB: We use modern tools (not docker), that do not require root access. This
 // allows this binary to convert images even when not running as root.
-func convertContainerToExt4FS(ctx context.Context, workspaceDir, containerImage string, creds container.PullCredentials) (string, error) {
+func convertContainerToExt4FS(ctx context.Context, dockerClient *dockerclient.Client, workspaceDir, containerImage string, creds container.PullCredentials) (string, error) {
 	// Make a temp directory to work in. Delete it when this fuction returns.
 	rootUnpackDir, err := os.MkdirTemp(workspaceDir, "container-unpack-*")
 	if err != nil {
@@ -153,20 +162,53 @@ func convertContainerToExt4FS(ctx context.Context, workspaceDir, containerImage
 		return "", err
 	}
 
-	// in CLI-form, the commands below do this:
-	// skopeo copy docker://alpine:lotest oci:/tmp/image_unpack:latest
-	// umoci unpack --rootless --image /tmp/image_unpack /tmp/bundle
-	// /tmp/bundle/rootfs/ has the goods
-	dockerImageRef := fmt.Sprintf("docker://%s", containerImage)
+	// In CLI-form, the commands below do this:
+	//
+	// docker pull alpine:latest
+	// docker save alpine:latest --output /tmp/image_unpack/docker_image.tar
+	// skopeo copy docker-archive:/tmp/image_unpack/docker_image.tar oci:/tmp/image_unpack/oci_image:latest
+	// umoci unpack --rootless --image /tmp/image_unpack/oci_image /tmp/image_unpack/bundle
+	//
+	// If docker is not available then we use skopeo to pull the image, like so:
+	//
+	// skopeo copy docker://alpine:latest oci:/tmp/image_unpack/oci_image:latest
+	//
+	// After running these commands, /tmp/image_unpack/bundle/rootfs/ has the
+	// unpacked image contents.
+	//
+	// The reason we use docker pull instead of directly copying with skopeo is
+	// that docker handles de-duping for us, and we can make use of existing
+	// cached image layers.
+	//
+	// Also note that we use an intermediate `docker save` rather than using the
+	// `docker-daemon:` protocol to export the image directly, due to
+	// https://github.com/containers/image/issues/1049
+	srcRef := fmt.Sprintf("docker://%s", containerImage)
+	if dockerClient != nil {
+		log.Debugf("Pulling: %s", containerImage)
+		if err := docker.PullImage(ctx, dockerClient, containerImage, creds); err != nil {
+			return "", err
+		}
+		log.Debugf("Exporting image from docker daemon: %s", containerImage)
+		dockerImgPath := filepath.Join(rootUnpackDir, "docker_image.tar")
+		if err := docker.SaveImage(ctx, dockerClient, containerImage, dockerImgPath); err != nil {
+			return "", err
+		}
+		log.Debugf("Converting to OCI image: %s", containerImage)
+		srcRef = fmt.Sprintf("docker-archive:%s", dockerImgPath)
+	} else {
+		log.Debugf("Downloading image and converting to OCI format: %s", containerImage)
+	}
 	ociOutputRef := fmt.Sprintf("oci:%s:latest", ociImageDir)
-	skopeoArgs := []string{"copy", dockerImageRef, ociOutputRef}
+	skopeoArgs := []string{"copy", srcRef, ociOutputRef}
 	if srcCreds := creds.String(); srcCreds != "" {
 		skopeoArgs = append(skopeoArgs, "--src-creds", srcCreds)
 	}
 	if out, err := exec.CommandContext(ctx, "skopeo", skopeoArgs...).CombinedOutput(); err != nil {
 		return "", status.InternalErrorf("skopeo copy error: %q: %s", string(out), err)
 	}
 
+	log.Debugf("Unpacking OCI image: %s", containerImage)
 	// Make a directory to unpack the bundle to.
 	rootFSDir := filepath.Join(rootUnpackDir, "rootfs")
 	if err := disk.EnsureDirectoryExists(rootFSDir); err != nil {