fix: route restart --location through the Workload

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

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
 }

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)
+}

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,

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")
+}