From 2e4cb758d99a34d374ebcf2222c2249ab48ae9b1 Mon Sep 17 00:00:00 2001 From: Scot Schuchert-Wells Date: Wed, 11 Nov 2026 15:02:24 +0000 Subject: [PATCH] fix: route restart --location through the Workload A location-scoped restart wrote the restart annotation straight onto the workload's WorkloadDeployments. Project users may update Workloads but not WorkloadDeployments, so the command failed as forbidden. Even with that permission the edit would not hold, because the workload controller owns deployment specs and overwrites them on its next reconcile. The command now records the restart as a per-location annotation on the Workload, and the workload controller copies it onto the template of the deployments at that location only. A location restart needs the same permission as a full restart and rolls only that location. Key changes: - add location-restart.compute.datumapis.com/ on Workloads and the locationRestartedAt template annotation it maps to - stamp the template annotation on matching deployments in the controller - update the Workload from restart --location instead of its deployments --- api/v1alpha/annotations.go | 20 ++++ internal/cmd/compute/restart/restart.go | 51 +++++---- internal/cmd/compute/restart/restart_test.go | 102 ++++++++++++++++++ internal/controller/workload_controller.go | 6 ++ .../controller/workload_controller_test.go | 56 ++++++++++ 5 files changed, 213 insertions(+), 22 deletions(-) create mode 100644 internal/cmd/compute/restart/restart_test.go diff --git a/api/v1alpha/annotations.go b/api/v1alpha/annotations.go index edc3860a..7bb7706c 100644 --- a/api/v1alpha/annotations.go +++ b/api/v1alpha/annotations.go @@ -20,6 +20,26 @@ const ( // template hash and triggers the existing ordered in-place roll. RestartedAtAnnotation = AnnotationNamespace + "/restartedAt" + // LocationRestartAnnotationPrefix prefixes a Workload annotation that asks + // for a rolling restart of the workload at one location. The annotation + // name after the prefix is the location name and the value is an RFC3339 + // timestamp, for example + // location-restart.compute.datumapis.com/us-east-1: 2026-01-02T15:04:05Z. + // + // The workload controller copies the value into + // LocationRestartedAtAnnotation on the template of the deployments at that + // location only, so a restart there needs only permission to update the + // Workload. Deployments are owned by the workload controller, and a direct + // edit to their template is overwritten on its next reconcile. + LocationRestartAnnotationPrefix = "location-restart." + AnnotationNamespace + "/" + + // LocationRestartedAtAnnotation is set by the workload controller on the + // template of a WorkloadDeployment from the Workload's + // LocationRestartAnnotationPrefix annotation for its location. Like + // RestartedAtAnnotation, it is part of the template hash and triggers the + // ordered in-place roll. + LocationRestartedAtAnnotation = AnnotationNamespace + "/locationRestartedAt" + // ReferencedDataGateStartAnnotation is stamped on an Instance by the cell // InstanceReconciler the first time it observes the ReferencedData scheduling // gate. Its value is an RFC3339 timestamp. Used to compute gate-wait duration diff --git a/internal/cmd/compute/restart/restart.go b/internal/cmd/compute/restart/restart.go index e15380d9..a589f041 100644 --- a/internal/cmd/compute/restart/restart.go +++ b/internal/cmd/compute/restart/restart.go @@ -3,12 +3,15 @@ package restart import ( "context" "fmt" + "slices" + "strings" "time" "github.com/spf13/cobra" k8serrors "k8s.io/apimachinery/pkg/api/errors" "k8s.io/apimachinery/pkg/labels" "k8s.io/apimachinery/pkg/types" + "k8s.io/apimachinery/pkg/util/validation" "sigs.k8s.io/controller-runtime/pkg/client" computev1alpha "go.datum.net/compute/api/v1alpha" @@ -57,7 +60,6 @@ func runRestart(cmd *cobra.Command, args []string, location string) error { out := cmd.OutOrStdout() if location == "" { - // Restart all placements by annotating the workload template. if workload.Spec.Template.Annotations == nil { workload.Spec.Template.Annotations = make(map[string]string) } @@ -74,7 +76,23 @@ func runRestart(cmd *cobra.Command, args []string, location string) error { return nil } - // Restart only deployments in the given location. + if err := restartLocation(ctx, c, &workload, location, restartedAt); err != nil { + return err + } + + fmt.Fprintf(out, + "Restarting workload %q in %s — rolling restart initiated.\nRun 'datumctl compute rollout %s' to watch progress.\n", + workloadName, location, workloadName, + ) + return nil +} + +func restartLocation(ctx context.Context, c client.Client, workload *computev1alpha.Workload, location, restartedAt string) error { + annotation := computev1alpha.LocationRestartAnnotationPrefix + location + if errs := validation.IsQualifiedName(annotation); len(errs) > 0 { + return fmt.Errorf("location %q cannot be restarted on its own: %s", location, strings.Join(errs, "; ")) + } + selector := labels.SelectorFromSet(labels.Set{ computev1alpha.WorkloadUIDLabel: string(workload.UID), }) @@ -86,31 +104,20 @@ func runRestart(cmd *cobra.Command, args []string, location string) error { return fmt.Errorf("listing deployments: %w", err) } - var matched []computev1alpha.WorkloadDeployment - for _, d := range deployList.Items { - if d.Spec.LocationRef.Name == location { - matched = append(matched, d) - } + if !slices.ContainsFunc(deployList.Items, func(d computev1alpha.WorkloadDeployment) bool { + return d.Spec.LocationRef.Name == location + }) { + return fmt.Errorf("no deployment found for workload %q in location %q", workload.Name, location) } - if len(matched) == 0 { - return fmt.Errorf("no deployment found for workload %q in location %q", workloadName, location) + if workload.Annotations == nil { + workload.Annotations = make(map[string]string) } + workload.Annotations[annotation] = restartedAt - for i := range matched { - if matched[i].Spec.Template.Annotations == nil { - matched[i].Spec.Template.Annotations = make(map[string]string) - } - matched[i].Spec.Template.Annotations[computev1alpha.RestartedAtAnnotation] = restartedAt - - if err := c.Update(ctx, &matched[i]); err != nil { - return fmt.Errorf("updating deployment in %s: %w", location, err) - } + if err := c.Update(ctx, workload); err != nil { + return fmt.Errorf("updating workload: %w", err) } - fmt.Fprintf(out, - "Restarting workload %q in %s — rolling restart initiated.\nRun 'datumctl compute rollout %s' to watch progress.\n", - workloadName, location, workloadName, - ) return nil } diff --git a/internal/cmd/compute/restart/restart_test.go b/internal/cmd/compute/restart/restart_test.go new file mode 100644 index 00000000..0f229a03 --- /dev/null +++ b/internal/cmd/compute/restart/restart_test.go @@ -0,0 +1,102 @@ +package restart + +import ( + "context" + "errors" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + metav1 "k8s.io/apimachinery/pkg/apis/meta/v1" + "k8s.io/apimachinery/pkg/runtime" + "k8s.io/apimachinery/pkg/types" + "sigs.k8s.io/controller-runtime/pkg/client" + "sigs.k8s.io/controller-runtime/pkg/client/fake" + "sigs.k8s.io/controller-runtime/pkg/client/interceptor" + + locationsv1alpha1 "go.miloapis.com/locations/api/v1alpha1" + + computev1alpha "go.datum.net/compute/api/v1alpha" + "go.datum.net/compute/internal/cmd/compute/util" +) + +const ( + testWorkload = "web" + testLocation = "us-east-1" + testRestartedAt = "2026-01-02T15:04:05Z" +) + +func newFixture(t *testing.T) (client.Client, *computev1alpha.Workload) { + t.Helper() + + s := runtime.NewScheme() + require.NoError(t, computev1alpha.AddToScheme(s)) + + workload := &computev1alpha.Workload{ + ObjectMeta: metav1.ObjectMeta{ + Name: testWorkload, + Namespace: util.ResourceNamespace, + UID: types.UID("workload-uid"), + }, + } + deployment := &computev1alpha.WorkloadDeployment{ + ObjectMeta: metav1.ObjectMeta{ + Name: testWorkload + "-default-" + testLocation, + Namespace: util.ResourceNamespace, + Labels: map[string]string{computev1alpha.WorkloadUIDLabel: string(workload.UID)}, + }, + Spec: computev1alpha.WorkloadDeploymentSpec{ + LocationRef: locationsv1alpha1.LocationReference{Name: testLocation}, + }, + } + + errDeploymentWrite := errors.New("workloaddeployments are not writable by project users") + c := fake.NewClientBuilder(). + WithScheme(s). + WithObjects(workload, deployment). + WithInterceptorFuncs(interceptor.Funcs{ + Update: func(ctx context.Context, c client.WithWatch, obj client.Object, opts ...client.UpdateOption) error { + if _, ok := obj.(*computev1alpha.WorkloadDeployment); ok { + return errDeploymentWrite + } + return c.Update(ctx, obj, opts...) + }, + Patch: func(ctx context.Context, c client.WithWatch, obj client.Object, patch client.Patch, opts ...client.PatchOption) error { + if _, ok := obj.(*computev1alpha.WorkloadDeployment); ok { + return errDeploymentWrite + } + return c.Patch(ctx, obj, patch, opts...) + }, + }). + Build() + + var stored computev1alpha.Workload + require.NoError(t, c.Get(context.Background(), client.ObjectKeyFromObject(workload), &stored)) + return c, &stored +} + +func TestRestartLocation_AnnotatesWorkloadOnly(t *testing.T) { + t.Parallel() + ctx := context.Background() + c, workload := newFixture(t) + + require.NoError(t, restartLocation(ctx, c, workload, testLocation, testRestartedAt)) + + var got computev1alpha.Workload + require.NoError(t, c.Get(ctx, client.ObjectKeyFromObject(workload), &got)) + assert.Equal(t, testRestartedAt, got.Annotations[computev1alpha.LocationRestartAnnotationPrefix+testLocation]) + assert.NotContains(t, got.Spec.Template.Annotations, computev1alpha.RestartedAtAnnotation) +} + +func TestRestartLocation_UnknownLocation(t *testing.T) { + t.Parallel() + ctx := context.Background() + c, workload := newFixture(t) + + err := restartLocation(ctx, c, workload, "eu-west-1", testRestartedAt) + require.ErrorContains(t, err, `no deployment found for workload "web" in location "eu-west-1"`) + + var got computev1alpha.Workload + require.NoError(t, c.Get(ctx, client.ObjectKeyFromObject(workload), &got)) + assert.Empty(t, got.Annotations) +} diff --git a/internal/controller/workload_controller.go b/internal/controller/workload_controller.go index 7eff89dc..1d8fef6b 100644 --- a/internal/controller/workload_controller.go +++ b/internal/controller/workload_controller.go @@ -675,6 +675,12 @@ func (r *WorkloadReconciler) getDeploymentsForWorkload( if template.Spec.Runtime.Resources.InstanceType == "" { template.Spec.Runtime.Resources.InstanceType = instancetype.D1Standard2 } + if restartedAt, ok := workload.Annotations[computev1alpha.LocationRestartAnnotationPrefix+locationName]; ok { + if template.Annotations == nil { + template.Annotations = map[string]string{} + } + template.Annotations[computev1alpha.LocationRestartedAtAnnotation] = restartedAt + } return template }(), ScaleSettings: placement.ScaleSettings, diff --git a/internal/controller/workload_controller_test.go b/internal/controller/workload_controller_test.go index 3521c982..96283799 100644 --- a/internal/controller/workload_controller_test.go +++ b/internal/controller/workload_controller_test.go @@ -809,3 +809,59 @@ func TestGetDeploymentsForWorkload_RequiresComputeAvailability(t *testing.T) { assert.Equal(t, testLocationName, deployment.Spec.LocationRef.Name) } } + +func TestGetDeploymentsForWorkload_LocationRestartReachesOnlyThatLocation(t *testing.T) { + t.Parallel() + + const otherLocation = "ord" + const restartedAt = "2026-01-02T15:04:05Z" + + workload := &computev1alpha.Workload{ + ObjectMeta: metav1.ObjectMeta{ + Name: rdTestWorkloadName, + Namespace: testDefaultNamespace, + UID: types.UID("workload-uid"), + Annotations: map[string]string{ + computev1alpha.LocationRestartAnnotationPrefix + testLocationName: restartedAt, + }, + }, + Spec: computev1alpha.WorkloadSpec{ + Placements: []computev1alpha.WorkloadPlacement{ + { + Name: testDefaultPlacement, + Locations: []locationsv1alpha1.LocationReference{ + {Name: testLocationName}, + {Name: otherLocation}, + }, + ScaleSettings: computev1alpha.HorizontalScaleSettings{MinReplicas: 1}, + }, + }, + }, + } + + s := newNetworkingScheme() + require.NoError(t, locationsv1alpha1.AddToScheme(s)) + cl := fake.NewClientBuilder(). + WithScheme(s). + WithObjects( + newLocationsServiceLocation(testLocationName, "DFW"), + newLocationsServiceLocation(otherLocation, "ORD"), + newTestComputeAvailability(testLocationName), + newTestComputeAvailability(otherLocation), + ). + WithIndex(&computev1alpha.WorkloadDeployment{}, deploymentWorkloadUIDIndex, deploymentWorkloadUIDIndexFunc). + Build() + r := &WorkloadReconciler{} + + desired, _, err := r.getDeploymentsForWorkload(context.Background(), cl, workload) + require.NoError(t, err) + require.Len(t, desired, 2) + + byLocation := map[string]computev1alpha.WorkloadDeployment{} + for _, d := range desired { + byLocation[d.Spec.LocationRef.Name] = d + } + assert.Equal(t, restartedAt, byLocation[testLocationName].Spec.Template.Annotations[computev1alpha.LocationRestartedAtAnnotation]) + assert.NotContains(t, byLocation[otherLocation].Spec.Template.Annotations, computev1alpha.LocationRestartedAtAnnotation) + assert.Empty(t, workload.Spec.Template.Annotations, "the workload's own template must not change") +}