volume: refactor volume.Set to be aware about GuestVolumeImageProvider

By knowing GuestVolumeImageProvider, callers can add any volumes
regardless of its lifecycle. For example, demux-backed volumes are
only accessible after starting a VM.

Signed-off-by: Kazuyoshi Kato <katokazu@amazon.com>

Author: Kazuyoshi Kato

snapshotter/volume_integ_test.go

@@ -61,7 +61,15 @@ func TestGuestVolumeFrom_Isolated(t *testing.T) {
 	remoteImage := volume.FromGuestImage(
 		client, vmID, al2stargz, "al2-snapshot", []string{"/etc/yum"},
 		volume.WithSnapshotter("demux"),
+		volume.WithPullOptions(containerd.WithImageHandlerWrapper(
+			source.AppendDefaultLabelsHandlerWrapper(al2stargz, 10*mib),
+		)),
 	)
+	err = vs.AddFrom(ctx, remoteImage)
+	require.NoError(t, err)
+
+	// PrepareDriveMount only copies images that are available before starting the VM.
+	// In this case, only postgres.
 	mount, err := vs.PrepareDriveMount(ctx, 10*mib)
 	require.NoError(t, err)
 
@@ -95,15 +103,9 @@ func TestGuestVolumeFrom_Isolated(t *testing.T) {
 	})
 	require.NoError(t, err, "Failed to configure VM metadata for registry resolution")
 
-	err = remoteImage.Pull(ctx,
-		containerd.WithImageHandlerWrapper(source.AppendDefaultLabelsHandlerWrapper(al2stargz, 10*mib)),
-	)
-	require.NoError(t, err)
-
-	err = remoteImage.Copy(ctx, "image")
-	require.NoError(t, err)
-
-	err = vs.AddFrom(ctx, remoteImage)
+	// PrepareGuestVolumes only copies images that are only available after starting the VM.
+	// In this case, only al2stargz.
+	err = vs.PrepareInGuest(ctx, "prepare-in-guest")
 	require.NoError(t, err)
 
 	image, err := client.Pull(ctx,

volume/guest_image.go

@@ -86,13 +86,13 @@ func (p *GuestVolumeImageProvider) Name() string {
 	return p.image
 }
 
-// Pull the image.
-func (p *GuestVolumeImageProvider) Pull(ctx context.Context, opts ...containerd.RemoteOpt) error {
+// pull the image.
+func (p *GuestVolumeImageProvider) pull(ctx context.Context) error {
 	remoteOpts := []containerd.RemoteOpt{
 		containerd.WithPullUnpack,
 		containerd.WithPullSnapshotter(p.config.snapshotter),
 	}
-	remoteOpts = append(remoteOpts, opts...)
+	remoteOpts = append(remoteOpts, p.config.pullOpts...)
 	image, err := p.client.Pull(ctx, p.image, remoteOpts...)
 	if err != nil {
 		return err
@@ -122,30 +122,30 @@ func mountHostDirs(image string) oci.SpecOpts {
 	}
 }
 
-// Copy files from the image by launching the container.
+// copy files from the image by launching the container.
 // The name must be unique within the VM and must be /[A-Z0-9a-z][A-Z0-9a-z._-]*/.
-func (p *GuestVolumeImageProvider) Copy(ctx context.Context, name string) error {
-	err := identifiers.Validate(name)
+func (p *GuestVolumeImageProvider) copy(ctx context.Context, containerName string) error {
+	err := identifiers.Validate(containerName)
 	if err != nil {
 		return err
 	}
 	container, err := p.client.NewContainer(ctx,
-		name,
+		containerName,
 		containerd.WithSnapshotter(p.config.snapshotter),
 		containerd.WithNewSnapshot(p.snapshot, p.cImage),
 		containerd.WithNewSpec(
 			firecrackeroci.WithVMID(p.vmID),
 			oci.WithProcessArgs(filepath.Join(containerBinDir, "volume-init")),
 			oci.WithProcessCwd("/"),
-			mountHostDirs(name),
+			mountHostDirs(containerName),
 		),
 	)
 	if err != nil {
-		return fmt.Errorf("failed to create container %q with %q: %w", name, p.config.snapshotter, err)
+		return fmt.Errorf("failed to create container %q with %q: %w", containerName, p.config.snapshotter, err)
 	}
 	defer container.Delete(ctx, containerd.WithSnapshotCleanup)
 
-	input := GuestVolumeImageInput{From: "/", To: filepath.Join(containerVolumesDir, name), Volumes: p.volumes}
+	input := GuestVolumeImageInput{From: "/", To: filepath.Join(containerVolumesDir, containerName), Volumes: p.volumes}
 	b, err := json.Marshal(&input)
 	if err != nil {
 		return err
@@ -197,7 +197,7 @@ func (p *GuestVolumeImageProvider) Copy(ctx context.Context, name string) error
 		return ctx.Err()
 	}
 
-	p.vmDir = filepath.Join(vmVolumesDir, name)
+	p.vmDir = filepath.Join(vmVolumesDir, containerName)
 	return nil
 }
 

volume/image.go

@@ -33,6 +33,7 @@ import (
 
 type imageProviderConfig struct {
 	snapshotter string
+	pullOpts    []containerd.RemoteOpt
 }
 
 type imageVolumeProvider struct {
@@ -53,6 +54,13 @@ func WithSnapshotter(ss string) ImageOpt {
 	}
 }
 
+// WithPullOptions sets the snapshotter's pull options.
+func WithPullOptions(opts ...containerd.RemoteOpt) ImageOpt {
+	return func(c *imageProviderConfig) {
+		c.pullOpts = opts
+	}
+}
+
 func getV1ImageConfig(ctx context.Context, image containerd.Image) (*v1.Image, error) {
 	var result v1.Image
 

volume/set.go

@@ -38,7 +38,8 @@ const (
 // Set is a set of volumes.
 type Set struct {
 	volumes   map[string]*Volume
-	providers map[string][]Mount
+	mounts    map[string][]Mount
+	providers map[string]Provider
 	tempDir   string
 	runtime   string
 	volumeDir string
@@ -65,7 +66,8 @@ func NewSet(runtime string) *Set {
 func NewSetWithTempDir(runtime, tempDir string) *Set {
 	return &Set{
 		runtime:   runtime,
-		providers: make(map[string][]Mount),
+		mounts:    make(map[string][]Mount),
+		providers: make(map[string]Provider),
 		volumes:   make(map[string]*Volume),
 		tempDir:   tempDir,
 	}
@@ -81,7 +83,12 @@ func (vs *Set) AddFrom(ctx context.Context, vp Provider) error {
 	if _, exists := vs.providers[vp.Name()]; exists {
 		return fmt.Errorf("failed to add %q: %w", vp.Name(), errdefs.ErrAlreadyExists)
 	}
-	dir, err := os.MkdirTemp(vs.tempDir, "AddFrom")
+	vs.providers[vp.Name()] = vp
+	return nil
+}
+
+func (vs *Set) copyToHostFromProvider(ctx context.Context, vp Provider) error {
+	dir, err := os.MkdirTemp(vs.tempDir, "copyToHostFromProvider")
 	if err != nil {
 		return err
 	}
@@ -100,7 +107,7 @@ func (vs *Set) AddFrom(ctx context.Context, vp Provider) error {
 
 		mounts = append(mounts, Mount{key: v.name, Destination: v.containerPath, ReadOnly: false})
 	}
-	vs.providers[vp.Name()] = mounts
+	vs.mounts[vp.Name()] = mounts
 	return nil
 }
 
@@ -147,7 +154,7 @@ func mountDiskImage(source, target string) error {
 
 // PrepareDirectory creates a directory that have volumes.
 func (vs *Set) PrepareDirectory(ctx context.Context) (retErr error) {
-	dir, err := os.MkdirTemp(vs.tempDir, "Prepare")
+	dir, err := os.MkdirTemp(vs.tempDir, "PrepareDirectory")
 	if err != nil {
 		retErr = err
 		return
@@ -161,19 +168,10 @@ func (vs *Set) PrepareDirectory(ctx context.Context) (retErr error) {
 		}
 	}()
 
-	for _, v := range vs.volumes {
-		path, err := fs.RootPath(dir, v.name)
-		if err != nil {
-			return err
-		}
-		if v.hostPath == "" {
-			continue
-		}
-		err = fs.CopyDir(path, v.hostPath)
-		if err != nil {
-			retErr = fmt.Errorf("failed to copy volume %q: %w", v.name, err)
-			return
-		}
+	err = vs.copyToHost(ctx, dir)
+	if err != nil {
+		retErr = err
+		return
 	}
 
 	vs.volumeDir = dir
@@ -221,30 +219,75 @@ func (vs *Set) PrepareDriveMount(ctx context.Context, size int64) (dm *proto.Fir
 		}
 	}()
 
+	err = vs.copyToHost(ctx, dir)
+	if err != nil {
+		retErr = err
+		return
+	}
+
+	dm = &proto.FirecrackerDriveMount{
+		HostPath:       path,
+		VMPath:         vmVolumePath,
+		FilesystemType: fsType,
+		IsWritable:     true,
+	}
+	vs.volumeDir = vmVolumePath
+	return
+}
+
+// PrepareInGuest prepares volumes inside the VM.
+func (vs *Set) PrepareInGuest(ctx context.Context, container string) error {
+	for _, provider := range vs.providers {
+		gp, guest := provider.(*GuestVolumeImageProvider)
+		if !guest {
+			continue
+		}
+
+		err := gp.pull(ctx)
+		if err != nil {
+			return err
+		}
+
+		err = gp.copy(ctx, container)
+		if err != nil {
+			return err
+		}
+
+		err = vs.copyToHostFromProvider(ctx, provider)
+		if err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+func (vs *Set) copyToHost(ctx context.Context, dir string) error {
+	for _, provider := range vs.providers {
+		_, guest := provider.(*GuestVolumeImageProvider)
+		if guest {
+			continue
+		}
+
+		err := vs.copyToHostFromProvider(ctx, provider)
+		if err != nil {
+			return err
+		}
+	}
+
 	for _, v := range vs.volumes {
 		path, err := fs.RootPath(dir, v.name)
 		if err != nil {
-			retErr = err
-			return
+			return err
 		}
 		if v.hostPath == "" {
 			continue
 		}
 		err = fs.CopyDir(path, v.hostPath)
 		if err != nil {
-			retErr = fmt.Errorf("failed to copy volume %q: %w", v.name, err)
-			return
+			return fmt.Errorf("failed to copy volume %q: %w", v.name, err)
 		}
 	}
-
-	dm = &proto.FirecrackerDriveMount{
-		HostPath:       path,
-		VMPath:         vmVolumePath,
-		FilesystemType: fsType,
-		IsWritable:     true,
-	}
-	vs.volumeDir = vmVolumePath
-	return
+	return nil
 }
 
 // Mount is used to expose volumes to containers.
@@ -303,5 +346,10 @@ func (vs *Set) WithMounts(mountpoints []Mount) (oci.SpecOpts, error) {
 
 // WithMountsFromProvider exposes volumes from the provider.
 func (vs *Set) WithMountsFromProvider(name string) (oci.SpecOpts, error) {
-	return vs.WithMounts(vs.providers[name])
+	_, exists := vs.providers[name]
+	if !exists {
+		return nil, fmt.Errorf("failed to find volume %q: %w", name, errdefs.ErrNotFound)
+	}
+
+	return vs.WithMounts(vs.mounts[name])
 }