[FC] Allow skipping the resave of action snapshots with a platform property (#8808)

This only applies to action executions which were already created from a
snapshot. My hypothesis is that saving an updated snapshot doesn't
provide future executions any benefit in these cases.

The experiment part of this will be in a separate PR.

Author: vanja-p

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

@@ -546,6 +546,7 @@ type FirecrackerContainer struct {
 	createFromSnapshot      bool
 	supportsRemoteSnapshots bool
 	recyclingEnabled        bool
+	shouldSaveSnapshot      bool
 
 	// If set, the snapshot used to load the VM
 	snapshot *snaploader.Snapshot
@@ -718,6 +719,7 @@ func NewContainer(ctx context.Context, env environment.Env, task *repb.Execution
 		// breaking change is made to the vmexec API, the executor should not
 		// attempt to connect to snapshots that were created before the change.
 	}
+	c.shouldSaveSnapshot = c.supportsRemoteSnapshots || !c.createFromSnapshot || !platform.IsTrue(platform.FindEffectiveValue(task, platform.SkipResavingActionSnapshotsPropertyName))
 
 	return c, nil
 }
@@ -2584,31 +2586,35 @@ func (c *FirecrackerContainer) pause(ctx context.Context) error {
 
 	log.CtxInfof(ctx, "Pausing VM")
 
-	if c.isBalloonEnabled() && c.machineHasBalloon(ctx) {
+	if c.shouldSaveSnapshot && c.isBalloonEnabled() && c.machineHasBalloon(ctx) {
 		if err := c.reclaimMemoryWithBalloon(ctx); err != nil {
 			log.CtxErrorf(ctx, "Reclaiming memory with the balloon failed with: %s", err)
 		}
 	}
 
-	snapDetails := c.snapshotDetails(ctx)
+	snapDetails := c.snapshotDetails()
 
 	// Before taking a snapshot, set the workspace drive to point to an empty
 	// file, so that we don't have to persist the workspace contents.
-	if err := c.updateWorkspaceDriveToEmptyFile(ctx); err != nil {
-		return status.WrapError(err, "update workspace drive to empty file")
+	if c.shouldSaveSnapshot {
+		if err := c.updateWorkspaceDriveToEmptyFile(ctx); err != nil {
+			return status.WrapError(err, "update workspace drive to empty file")
+		}
 	}
 
 	if err := c.pauseVM(ctx); err != nil {
 		return err
 	}
 
-	// If an older snapshot is present -- nuke it since we're writing a new one.
-	if err := c.cleanupOldSnapshots(ctx, snapDetails); err != nil {
-		return err
-	}
+	if c.shouldSaveSnapshot {
+		// If an older snapshot is present -- nuke it since we're writing a new one.
+		if err := c.cleanupOldSnapshots(ctx, snapDetails); err != nil {
+			return err
+		}
 
-	if err := c.createSnapshot(ctx, snapDetails); err != nil {
-		return err
+		if err := c.createSnapshot(ctx, snapDetails); err != nil {
+			return err
+		}
 	}
 
 	// Stop the VM, UFFD page fault handler, and VBD servers to ensure nothing
@@ -2628,8 +2634,10 @@ func (c *FirecrackerContainer) pause(ctx context.Context) error {
 		return status.WrapError(err, "unmount vbds")
 	}
 
-	if err := c.saveSnapshot(ctx, snapDetails); err != nil {
-		return err
+	if c.shouldSaveSnapshot {
+		if err := c.saveSnapshot(ctx, snapDetails); err != nil {
+			return err
+		}
 	}
 
 	// Finish cleaning up VM resources
@@ -2747,7 +2755,7 @@ type snapshotDetails struct {
 	vmStateSnapshotName string
 }
 
-func (c *FirecrackerContainer) snapshotDetails(ctx context.Context) *snapshotDetails {
+func (c *FirecrackerContainer) snapshotDetails() *snapshotDetails {
 	if c.recycled {
 		return &snapshotDetails{
 			snapshotType:        diffSnapshotType,

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

@@ -679,6 +679,101 @@ func TestFirecracker_LocalSnapshotSharing(t *testing.T) {
 	err = c.Pause(ctx)
 	require.NoError(t, err)
 }
+func TestFirecracker_LocalSnapshotSharing_DontResave(t *testing.T) {
+	ctx := context.Background()
+	env := getTestEnv(ctx, t, envOpts{})
+	env.SetAuthenticator(testauth.NewTestAuthenticator(testauth.TestUsers("US1", "GR1")))
+	rootDir := testfs.MakeTempDir(t)
+	cfg := getExecutorConfig(t)
+
+	var containersToCleanup []*firecracker.FirecrackerContainer
+	t.Cleanup(func() {
+		for _, vm := range containersToCleanup {
+			err := vm.Remove(ctx)
+			assert.NoError(t, err)
+		}
+	})
+
+	workDir := testfs.MakeDirAll(t, rootDir, "work")
+	opts := firecracker.ContainerOpts{
+		ContainerImage:         busyboxImage,
+		ActionWorkingDirectory: workDir,
+		VMConfiguration: &fcpb.VMConfiguration{
+			NumCpus:           1,
+			MemSizeMb:         minMemSizeMB, // small to make snapshotting faster.
+			EnableNetworking:  false,
+			ScratchDiskSizeMb: 100,
+		},
+		ExecutorConfig: cfg,
+	}
+	task := &repb.ExecutionTask{
+		Command: &repb.Command{
+			// Note: platform must match in order to share snapshots
+			Platform: &repb.Platform{Properties: []*repb.Platform_Property{
+				{Name: "recycle-runner", Value: "true"},
+				{Name: platform.SkipResavingActionSnapshotsPropertyName, Value: "true"},
+			}},
+		},
+	}
+	baseVM, err := firecracker.NewContainer(ctx, env, task, opts)
+	require.NoError(t, err)
+	containersToCleanup = append(containersToCleanup, baseVM)
+	err = container.PullImageIfNecessary(ctx, env, baseVM, oci.Credentials{}, opts.ContainerImage)
+	require.NoError(t, err)
+	err = baseVM.Create(ctx, opts.ActionWorkingDirectory)
+	require.NoError(t, err)
+
+	// Create a snapshot. Data written to this snapshot should persist
+	// when other VMs reuse the snapshot
+	cmd := appendToLog("Base")
+	res := baseVM.Exec(ctx, cmd, nil /*=stdio*/)
+	require.NoError(t, res.Error)
+	require.Equal(t, "Base\n", string(res.Stdout))
+	err = baseVM.Pause(ctx)
+	require.NoError(t, err)
+
+	// Load the same base snapshot from multiple VMs - there should be no
+	// corruption or data transfer from snapshot sharing
+	for i := 0; i < 4; i++ {
+		workDir = testfs.MakeDirAll(t, rootDir, fmt.Sprintf("work-%d", i))
+		opts = firecracker.ContainerOpts{
+			ContainerImage:         busyboxImage,
+			ActionWorkingDirectory: workDir,
+			VMConfiguration: &fcpb.VMConfiguration{
+				NumCpus:           1,
+				MemSizeMb:         minMemSizeMB, // small to make snapshotting faster.
+				EnableNetworking:  false,
+				ScratchDiskSizeMb: 100,
+			},
+			ExecutorConfig: cfg,
+		}
+		forkedVM, err := firecracker.NewContainer(ctx, env, task, opts)
+		require.NoError(t, err)
+		containersToCleanup = append(containersToCleanup, forkedVM)
+
+		// The new VM should reuse the Base VM's snapshot, whether we call
+		// Create() or Unpause()
+		if i%2 == 0 {
+			err = forkedVM.Unpause(ctx)
+			require.NoError(t, err)
+		} else {
+			err = forkedVM.Create(ctx, workDir)
+			require.NoError(t, err)
+		}
+
+		// Write VM-specific data to the log
+		cmd = appendToLog(fmt.Sprintf("Fork-%d", i))
+		res = forkedVM.Exec(ctx, cmd, nil /*=stdio*/)
+		require.NoError(t, res.Error)
+		// The log should contain data written to the original snapshot
+		// and the current VM, but not from any of the other VMs sharing
+		// the same original snapshot
+		require.Equal(t, fmt.Sprintf("Base\nFork-%d\n", i), string(res.Stdout))
+		// This shouldn't save a snapshot
+		err = forkedVM.Pause(ctx)
+		require.NoError(t, err)
+	}
+}
 
 func TestFirecracker_RemoteSnapshotSharing(t *testing.T) {
 	ctx := context.Background()

enterprise/server/remote_execution/platform/platform.go

@@ -73,31 +73,32 @@ const (
 	// empty or unset.
 	unsetContainerImageVal = "none"
 
-	RecycleRunnerPropertyName            = "recycle-runner"
-	RunnerRecyclingMaxWaitPropertyName   = "runner-recycling-max-wait"
-	PreserveWorkspacePropertyName        = "preserve-workspace"
-	overlayfsWorkspacePropertyName       = "overlayfs-workspace"
-	cleanWorkspaceInputsPropertyName     = "clean-workspace-inputs"
-	persistentWorkerPropertyName         = "persistent-workers"
-	persistentWorkerKeyPropertyName      = "persistentWorkerKey"
-	persistentWorkerProtocolPropertyName = "persistentWorkerProtocol"
-	WorkflowIDPropertyName               = "workflow-id"
-	workloadIsolationPropertyName        = "workload-isolation-type"
-	initDockerdPropertyName              = "init-dockerd"
-	enableDockerdTCPPropertyName         = "enable-dockerd-tcp"
-	enableVFSPropertyName                = "enable-vfs"
-	HostedBazelAffinityKeyPropertyName   = "hosted-bazel-affinity-key"
-	useSelfHostedExecutorsPropertyName   = "use-self-hosted-executors"
-	disableMeasuredTaskSizePropertyName  = "debug-disable-measured-task-size"
-	disablePredictedTaskSizePropertyName = "debug-disable-predicted-task-size"
-	extraArgsPropertyName                = "extra-args"
-	EnvOverridesPropertyName             = "env-overrides"
-	EnvOverridesBase64PropertyName       = "env-overrides-base64"
-	IncludeSecretsPropertyName           = "include-secrets"
-	DefaultTimeoutPropertyName           = "default-timeout"
-	TerminationGracePeriodPropertyName   = "termination-grace-period"
-	SnapshotKeyOverridePropertyName      = "snapshot-key-override"
-	RetryPropertyName                    = "retry"
+	RecycleRunnerPropertyName               = "recycle-runner"
+	RunnerRecyclingMaxWaitPropertyName      = "runner-recycling-max-wait"
+	PreserveWorkspacePropertyName           = "preserve-workspace"
+	overlayfsWorkspacePropertyName          = "overlayfs-workspace"
+	cleanWorkspaceInputsPropertyName        = "clean-workspace-inputs"
+	persistentWorkerPropertyName            = "persistent-workers"
+	persistentWorkerKeyPropertyName         = "persistentWorkerKey"
+	persistentWorkerProtocolPropertyName    = "persistentWorkerProtocol"
+	WorkflowIDPropertyName                  = "workflow-id"
+	workloadIsolationPropertyName           = "workload-isolation-type"
+	initDockerdPropertyName                 = "init-dockerd"
+	enableDockerdTCPPropertyName            = "enable-dockerd-tcp"
+	enableVFSPropertyName                   = "enable-vfs"
+	HostedBazelAffinityKeyPropertyName      = "hosted-bazel-affinity-key"
+	useSelfHostedExecutorsPropertyName      = "use-self-hosted-executors"
+	disableMeasuredTaskSizePropertyName     = "debug-disable-measured-task-size"
+	disablePredictedTaskSizePropertyName    = "debug-disable-predicted-task-size"
+	extraArgsPropertyName                   = "extra-args"
+	EnvOverridesPropertyName                = "env-overrides"
+	EnvOverridesBase64PropertyName          = "env-overrides-base64"
+	IncludeSecretsPropertyName              = "include-secrets"
+	DefaultTimeoutPropertyName              = "default-timeout"
+	TerminationGracePeriodPropertyName      = "termination-grace-period"
+	SnapshotKeyOverridePropertyName         = "snapshot-key-override"
+	RetryPropertyName                       = "retry"
+	SkipResavingActionSnapshotsPropertyName = "skip-resaving-action-snapshots"
 
 	OperatingSystemPropertyName = "OSFamily"
 	LinuxOperatingSystemName    = "linux"