ateletpb: split the Restore snapshot source from the checkpoint destination

cmd/ateapi/internal/controlapi/functionaltest/actor_test.go

@@ -1910,12 +1910,120 @@ func TestResumeActorPassesLiteralEnv(t *testing.T) {
 	}
 }
 
-// TestResumeActor_NoWorkers tests that resuming an actor fails when no free workers are available.
-// Workflow:
-// 1. Creates a mock ActorTemplate.
-// 2. Creates an actor.
-// 3. Calls ResumeActor RPC without creating any workers.
-// 4. Verifies that ResumeActor fails with FailedPrecondition status.
+// createGoldenDataTemplate creates "tmpl1" like createTemplate, but with
+// onCommit DATA and onResume.fromData GOLDEN, so a resumed-after-suspend
+// actor takes the DATA_ON_GOLDEN path: its data snapshot combined with the
+// template's golden.
+func createGoldenDataTemplate(t *testing.T, tc *testContext, ns string) *ateapipb.ActorTemplate {
+	t.Helper()
+	ensureDefaultGvisorSandboxConfig(t, tc)
+	createWorkerPool(t, tc, ns, "pool1", map[string]string{poolLabelKey: ns})
+
+	created, err := tc.client.CreateActorTemplate(context.Background(), &ateapipb.CreateActorTemplateRequest{
+		ActorTemplate: &ateapipb.ActorTemplate{
+			Metadata: &ateapipb.ResourceMetadata{
+				Atespace: testAtespace,
+				Name:     "tmpl1",
+			},
+			SnapshotsConfig: &ateapipb.SnapshotsConfig{
+				StorageLocation: testStorageLocation,
+				OnPause:         ateapipb.SnapshotContentScope_SNAPSHOT_CONTENT_SCOPE_FULL,
+				OnCommit:        ateapipb.SnapshotContentScope_SNAPSHOT_CONTENT_SCOPE_DATA,
+				OnResume:        &ateapipb.OnResumeConfig{FromData: ateapipb.ResumeSource_RESUME_SOURCE_GOLDEN},
+			},
+			SandboxConfig: &ateapipb.SandboxConfig{
+				SandboxClass: ateapipb.SandboxClass_SANDBOX_CLASS_GVISOR,
+				ConfigName:   "gvisor-default",
+			},
+			Containers: []*ateapipb.Container{{
+				Name:    "main",
+				Image:   "main@sha256:abc",
+				Command: []string{"/main"},
+			}},
+			WorkerSelector: &ateapipb.Selector{
+				MatchLabels: map[string]string{poolLabelKey: ns},
+			},
+		},
+	})
+	if err != nil {
+		t.Fatalf("failed to create actor template: %v", err)
+	}
+	updated, err := tc.persistence.UpdateActorTemplate(context.Background(),
+		resources.ActorTemplateRefFromActorTemplate(created), store.PreconditionFrom(created),
+		func(dbTemplate *ateapipb.ActorTemplate) error {
+			dbTemplate.Status = &ateapipb.ActorTemplateStatus{
+				GoldenSnapshotStatus: &ateapipb.GoldenSnapshotStatus{
+					GoldenSnapshot: &ateapipb.ExternalSnapshot{SnapshotUri: goldenSnapshotURI(t), ContentScope: ateapipb.SnapshotContentScope_SNAPSHOT_CONTENT_SCOPE_FULL},
+				},
+			}
+			return nil
+		})
+	if err != nil {
+		t.Fatalf("failed to record the template's golden snapshot: %v", err)
+	}
+	return updated
+}
+
+// TestResumeActor_GoldenDataResumeSetsBaseConfig drives the DATA_ON_GOLDEN
+// resume end to end and pins the wire request's base snapshot fields: while
+// the golden_snapshot_uri -> base_config transition lasts, ateapi sets both
+// and they must agree, so ateapi and atelet can roll in either order.
+func TestResumeActor_GoldenDataResumeSetsBaseConfig(t *testing.T) {
+	ns := namespaceForTest("ns-resume-golden-data")
+	tc := setupTest(t, ns)
+	defer tc.cleanup()
+
+	tmpl := createGoldenDataTemplate(t, tc, ns)
+	workerName := createWorkerPod(t, tc, ns, "worker-1", "node1", "pool1")
+
+	const name = "id1"
+	actorRef := &ateapipb.ObjectRef{Atespace: testAtespace, Name: name}
+	if _, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{Actor: &ateapipb.Actor{
+		Metadata:      &ateapipb.ResourceMetadata{Atespace: testAtespace, Name: name},
+		ActorTemplate: &ateapipb.ObjectRef{Atespace: testAtespace, Name: "tmpl1"},
+	}}); err != nil {
+		t.Fatalf("CreateActor failed: %v", err)
+	}
+
+	// First resume runs fresh from the golden; the suspend then commits a
+	// DATA snapshot per onCommit.
+	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: actorRef}); err != nil {
+		t.Fatalf("ResumeActor (first) failed: %v", err)
+	}
+	suspended, err := tc.client.SuspendActor(context.Background(), &ateapipb.SuspendActorRequest{Actor: actorRef})
+	if err != nil {
+		t.Fatalf("SuspendActor failed: %v", err)
+	}
+	waitForWorkerAvailable(t, tc, workerName)
+	actorSnapshotURI := suspended.GetActor().GetStatus().GetExternalSnapshot().GetSnapshotUri()
+	if actorSnapshotURI == "" {
+		t.Fatal("SuspendActor recorded no external snapshot")
+	}
+
+	// Second resume: the actor's DATA snapshot rides on the template's
+	// golden.
+	if _, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{Actor: actorRef}); err != nil {
+		t.Fatalf("ResumeActor (second) failed: %v", err)
+	}
+	restoreReq := tc.fakeAtelet.lastRestoreRequest()
+	if restoreReq == nil {
+		t.Fatal("second resume sent no Restore request to atelet")
+	}
+	if got := restoreReq.GetScope(); got != ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN {
+		t.Fatalf("restore scope = %v, want SNAPSHOT_SCOPE_DATA_ON_GOLDEN", got)
+	}
+	if got := restoreReq.GetExternalConfig().GetSnapshotUri(); got != actorSnapshotURI {
+		t.Errorf("restore config snapshot uri = %q, want the actor's data snapshot %q", got, actorSnapshotURI)
+	}
+	golden := tmpl.GetStatus().GetGoldenSnapshotStatus().GetGoldenSnapshot().GetSnapshotUri()
+	if got := restoreReq.GetBaseConfig().GetSnapshotUri(); got != golden {
+		t.Errorf("restore base_config uri = %q, want the template's golden %q", got, golden)
+	}
+	if got := restoreReq.GetGoldenSnapshotUri(); got != golden {
+		t.Errorf("restore golden_snapshot_uri = %q, want %q (transitional dual-write must match base_config)", got, golden)
+	}
+}
+
 // TestResumeActor_NoWorkers tests that resuming an actor fails when no free workers are available.
 // Workflow:
 // 1. Creates a mock ActorTemplate.
@@ -2476,6 +2584,9 @@ func TestResumeActor_RepointTemplateBeforeResume(t *testing.T) {
 			if got := restoreReq.GetGoldenSnapshotUri(); got != "" {
 				t.Errorf("restore request to atelet had golden snapshot uri = %q, want empty", got)
 			}
+			if restoreReq.GetBaseConfig() != nil {
+				t.Errorf("restore request to atelet had base_config = %v, want unset", restoreReq.GetBaseConfig())
+			}
 		})
 	}
 }

cmd/ateapi/internal/controlapi/workflow_resume.go

@@ -700,6 +700,10 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 			req.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA
 		case !src.GoldenSnapshotURI.IsZero():
 			req.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			// Transitional dual-write: base_config supersedes
+			// golden_snapshot_uri, but an atelet from before it reads only
+			// the old field. Dropped once both components have rolled.
+			req.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: src.GoldenSnapshotURI.String()}
 			req.GoldenSnapshotUri = src.GoldenSnapshotURI.String()
 		default:
 			req.Scope = actorSnapshotContentScopeToAtelet(actorTemplate.GetSnapshotsConfig().GetOnPause())
@@ -718,11 +722,16 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 		}
 		var scope ateletpb.SnapshotScope
 		var goldenSnapshotURI string
+		var baseConfig *ateletpb.ExternalRestoreConfiguration
 		switch {
 		case src.TemplateReplaced:
 			scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA
 		case !src.GoldenSnapshotURI.IsZero():
 			scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			// Transitional dual-write: base_config supersedes
+			// golden_snapshot_uri, but an atelet from before it reads only
+			// the old field. Dropped once both components have rolled.
+			baseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: src.GoldenSnapshotURI.String()}
 			goldenSnapshotURI = src.GoldenSnapshotURI.String()
 		default:
 			scope = actorSnapshotContentScopeToAtelet(src.Scope)
@@ -737,12 +746,13 @@ func (w *ActorWorkflow) ensureAteletRestored(ctx context.Context, actorRef resou
 			Spec:                  workloadSpec,
 			Type:                  ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL,
 			Config: &ateletpb.RestoreRequest_ExternalConfig{
-				ExternalConfig: &ateletpb.ExternalCheckpointConfiguration{
+				ExternalConfig: &ateletpb.ExternalRestoreConfiguration{
 					SnapshotUri: src.SnapshotURI.String(),
 				},
 			},
 			Scope: scope,
-			// Empty unless this is a Golden data resume.
+			// Both empty unless this is a Golden data resume.
+			BaseConfig:        baseConfig,
 			GoldenSnapshotUri: goldenSnapshotURI,
 			ActorUid:          actor.GetMetadata().Uid,
 			EgressGateway:     egressGateway,

cmd/atelet/main.go

@@ -1072,9 +1072,10 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 	// and its pinned sandbox binaries are the ones that will run the restored
 	// guest (the golden snapshot's memory image must be resumed by the binaries
 	// that created it).
+	baseCfg := restoreBaseConfig(req)
 	var goldenRec *sandboxAssetsRecord
 	if req.GetScope() == ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN {
-		goldenURI, err := resources.ParseSnapshotURI(req.GetGoldenSnapshotUri())
+		goldenURI, err := resources.ParseSnapshotURI(baseCfg.GetSnapshotUri())
 		if err != nil {
 			return nil, ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonInvalidObjectURL)
 		}
@@ -1143,7 +1144,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 				if goldenRec == nil {
 					return fmt.Errorf("no golden snapshot record for a %s restore", req.GetScope())
 				}
-				if err := s.downloadCombinedCheckpoint(gctx, req.GetExternalConfig().GetSnapshotUri(), req.GetGoldenSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles); err != nil {
+				if err := s.downloadCombinedCheckpoint(gctx, req.GetExternalConfig().GetSnapshotUri(), baseCfg.GetSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles); err != nil {
 					return ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonFailedGetExternalObject, ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonTerminalFileSystemError)
 				}
 			} else if err := s.downloadExternalCheckpoint(gctx, req.GetExternalConfig().GetSnapshotUri(), checkpointDir, sandboxRec.SnapshotFiles); err != nil {
@@ -1166,7 +1167,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 			})
 			if combineWithGolden {
 				gLocal.Go(func() error {
-					if err := s.downloadExternalCheckpoint(gLocalCtx, req.GetGoldenSnapshotUri(), checkpointDir, goldenOnlyFiles(sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles)); err != nil {
+					if err := s.downloadExternalCheckpoint(gLocalCtx, baseCfg.GetSnapshotUri(), checkpointDir, goldenOnlyFiles(sandboxRec.SnapshotFiles, goldenRec.SnapshotFiles)); err != nil {
 						return ateerrors.CrashIfReason(ctx, err, ateerrors.ReasonFailedGetExternalObject, ateerrors.ReasonInvalidObjectURL, ateerrors.ReasonTerminalFileSystemError)
 					}
 					return nil
@@ -1242,7 +1243,7 @@ func (s *AteomHerder) Restore(ctx context.Context, req *ateletpb.RestoreRequest)
 		// Informational: for DATA_ON_GOLDEN the golden snapshot's files are
 		// already staged into the restore dir by the combined download above;
 		// ateom restores from the shared dir and never fetches this URI.
-		GoldenSnapshotUri: req.GetGoldenSnapshotUri(),
+		GoldenSnapshotUri: baseCfg.GetSnapshotUri(),
 	})
 	dAteom = time.Since(tAteom)
 	if err != nil {
@@ -1750,14 +1751,33 @@ func validateRestoreRequest(req *ateletpb.RestoreRequest) error {
 	}
 
 	// A DATA_ON_GOLDEN restore needs both halves: the actor's data snapshot
-	// (local pause checkpoint or external commit) and the golden snapshot,
-	// which is always external.
+	// (local pause checkpoint or external commit) and the base snapshot,
+	// which is always external. base_config supersedes golden_snapshot_uri;
+	// a transitional caller sets both, and they must agree.
+	base, legacy := req.GetBaseConfig().GetSnapshotUri(), req.GetGoldenSnapshotUri()
+	if base != "" && legacy != "" && base != legacy {
+		return fmt.Errorf("base_config.snapshot_uri %q and golden_snapshot_uri %q disagree", base, legacy)
+	}
 	if req.GetScope() == ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN {
-		if _, err := resources.ParseSnapshotURI(req.GetGoldenSnapshotUri()); err != nil {
-			return fmt.Errorf("invalid golden_snapshot_uri: %w", err)
+		if _, err := resources.ParseSnapshotURI(restoreBaseConfig(req).GetSnapshotUri()); err != nil {
+			return fmt.Errorf("invalid base snapshot URI: %w", err)
 		}
-	} else if req.GetGoldenSnapshotUri() != "" {
-		return fmt.Errorf("golden_snapshot_uri is only valid with snapshot scope %s", ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN)
+	} else if base != "" || legacy != "" {
+		return fmt.Errorf("a base snapshot (base_config or golden_snapshot_uri) is only valid with snapshot scope %s", ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN)
+	}
+	return nil
+}
+
+// restoreBaseConfig returns the base snapshot source of a DATA_ON_GOLDEN
+// restore, preferring base_config over the superseded golden_snapshot_uri
+// (still sent by callers that predate it). Nil when the request carries
+// neither; proto getters make that safe to read through.
+func restoreBaseConfig(req *ateletpb.RestoreRequest) *ateletpb.ExternalRestoreConfiguration {
+	if req.GetBaseConfig().GetSnapshotUri() != "" {
+		return req.GetBaseConfig()
+	}
+	if uri := req.GetGoldenSnapshotUri(); uri != "" {
+		return &ateletpb.ExternalRestoreConfiguration{SnapshotUri: uri}
 	}
 	return nil
 }

cmd/atelet/main_test.go

@@ -227,7 +227,7 @@ func validRestoreRequest() *ateletpb.RestoreRequest {
 		Spec:                  &ateletpb.WorkloadSpec{Containers: []*ateletpb.Container{{Name: "worker"}}},
 		Type:                  ateletpb.CheckpointType_CHECKPOINT_TYPE_EXTERNAL,
 		Config: &ateletpb.RestoreRequest_ExternalConfig{
-			ExternalConfig: &ateletpb.ExternalCheckpointConfiguration{
+			ExternalConfig: &ateletpb.ExternalRestoreConfiguration{
 				SnapshotUri: testSnapshotURI,
 			},
 		},
@@ -395,6 +395,29 @@ func TestValidateRestoreRequest(t *testing.T) {
 		{"golden uri with non-data-on-golden scope", makeReq(func(r *ateletpb.RestoreRequest) {
 			r.GoldenSnapshotUri = goldenSnapshotURI
 		}), true},
+		{"data-on-golden with base config only", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			r.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}
+		}), false},
+		// A transitional caller sets base_config and the superseded
+		// golden_snapshot_uri together; they must name one snapshot.
+		{"data-on-golden with agreeing base config and golden uri", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			r.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}
+			r.GoldenSnapshotUri = goldenSnapshotURI
+		}), false},
+		{"base config and golden uri disagree", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			r.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}
+			r.GoldenSnapshotUri = testSnapshotURI
+		}), true},
+		{"data-on-golden with bucketless base config", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.Scope = ateletpb.SnapshotScope_SNAPSHOT_SCOPE_DATA_ON_GOLDEN
+			r.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: "relative/path"}
+		}), true},
+		{"base config with non-data-on-golden scope", makeReq(func(r *ateletpb.RestoreRequest) {
+			r.BaseConfig = &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}
+		}), true},
 	}
 	for _, tc := range tests {
 		t.Run(tc.name, func(t *testing.T) {
@@ -405,6 +428,33 @@ func TestValidateRestoreRequest(t *testing.T) {
 	}
 }
 
+// TestRestoreBaseConfig pins the dual-read precedence during the
+// golden_snapshot_uri -> base_config transition: base_config wins when it
+// names a snapshot, the legacy field covers callers that predate it, and a
+// request with neither yields nil (safe through proto getters).
+func TestRestoreBaseConfig(t *testing.T) {
+	cases := []struct {
+		name    string
+		base    *ateletpb.ExternalRestoreConfiguration
+		legacy  string
+		wantURI string
+	}{
+		{"neither set", nil, "", ""},
+		{"base config only", &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}, "", goldenSnapshotURI},
+		{"legacy only", nil, goldenSnapshotURI, goldenSnapshotURI},
+		{"base config preferred over legacy", &ateletpb.ExternalRestoreConfiguration{SnapshotUri: goldenSnapshotURI}, testSnapshotURI, goldenSnapshotURI},
+		{"empty base config falls back to legacy", &ateletpb.ExternalRestoreConfiguration{}, goldenSnapshotURI, goldenSnapshotURI},
+	}
+	for _, tc := range cases {
+		t.Run(tc.name, func(t *testing.T) {
+			req := &ateletpb.RestoreRequest{BaseConfig: tc.base, GoldenSnapshotUri: tc.legacy}
+			if got := restoreBaseConfig(req).GetSnapshotUri(); got != tc.wantURI {
+				t.Errorf("restoreBaseConfig().GetSnapshotUri() = %q, want %q", got, tc.wantURI)
+			}
+		})
+	}
+}
+
 // Every valid atelet scope must map to its ateom counterpart; in particular
 // DATA_ON_GOLDEN must never silently degrade to FULL.
 func TestToAteomSnapshotScope(t *testing.T) {