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