Prototype prepared Agent surge with capacity fallback

go.mod

@@ -29,6 +29,7 @@ require (
 	k8s.io/apimachinery v0.35.3
 	k8s.io/cli-runtime v0.35.3
 	k8s.io/client-go v0.35.3
+	k8s.io/component-helpers v0.35.3 // indirect
 	k8s.io/klog/v2 v2.140.0
 	k8s.io/kube-aggregator v0.35.3
 	k8s.io/kube-openapi v0.0.0-20251125145642-4e65d59e963e
@@ -291,7 +292,6 @@ require (
 	k8s.io/apiserver v0.35.3 // indirect
 	k8s.io/cloud-provider v0.35.0 // indirect
 	k8s.io/component-base v0.35.3 // indirect
-	k8s.io/component-helpers v0.35.3 // indirect
 	k8s.io/csi-translation-lib v0.35.0 // indirect
 	k8s.io/gengo/v2 v2.0.0-20250922181213-ec3ebc5fd46b // indirect
 	modernc.org/sqlite v1.34.1 // indirect

internal/controller/datadogagentinternal/controller_reconcile_agent.go

@@ -12,9 +12,11 @@ import (
 
 	edsv1alpha1 "github.com/DataDog/extendeddaemonset/api/v1alpha1"
 	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	ctrl "sigs.k8s.io/controller-runtime"
+	"sigs.k8s.io/controller-runtime/pkg/client"
 	"sigs.k8s.io/controller-runtime/pkg/reconcile"
 
 	datadoghqv1alpha1 "github.com/DataDog/datadog-operator/api/datadoghq/v1alpha1"
@@ -289,7 +291,48 @@ func (r *Reconciler) reconcileV2Agent(ctx context.Context, requiredComponents fe
 		return reconcile.Result{}, nil
 	}
 
-	return r.createOrUpdateDaemonset(ctx, ddai, daemonset, newStatus, updateDSStatusV2WithAgent)
+	rolloutBudget := preparedRolloutBudget(ddai, &r.options.ExtendedDaemonsetOptions)
+	rolloutEnabled := preparedRolloutEnabled(ddai)
+	var currentDaemonSet *appsv1.DaemonSet
+	if rolloutEnabled {
+		reader := r.apiReader
+		if reader == nil {
+			reader = r.client
+		}
+		currentDaemonSet = &appsv1.DaemonSet{}
+		if getErr := reader.Get(ctx, client.ObjectKeyFromObject(daemonset), currentDaemonSet); getErr != nil {
+			if !errors.IsNotFound(getErr) {
+				return reconcile.Result{}, getErr
+			}
+			currentDaemonSet = nil
+		}
+	}
+	affinityMigration, prepareErr := configurePreparedRollout(ddai, daemonset, currentDaemonSet, rolloutBudget)
+	if prepareErr != nil {
+		objLogger.Error(prepareErr, "Prepared Agent rollout request is incompatible with the rendered Pod template")
+		if r.recorder != nil {
+			r.recorder.Eventf(ddai, corev1.EventTypeWarning, "AgentPreparedRolloutRejected", "Prepared Agent rollout is disabled for this template: %v", prepareErr)
+		}
+		return reconcile.Result{}, prepareErr
+	}
+	result, err := r.createOrUpdateDaemonset(ctx, ddai, daemonset, newStatus, updateDSStatusV2WithAgent)
+	if err != nil || !rolloutEnabled {
+		return result, err
+	}
+	if affinityMigration {
+		if result.RequeueAfter == 0 || result.RequeueAfter > time.Second {
+			result.RequeueAfter = time.Second
+		}
+		return result, nil
+	}
+	fallbackResult, fallbackErr := r.reconcilePreparedRollout(ctx, ddai, daemonset, rolloutBudget)
+	if fallbackErr != nil {
+		return reconcile.Result{}, fallbackErr
+	}
+	if fallbackResult.RequeueAfter > 0 && (result.RequeueAfter == 0 || fallbackResult.RequeueAfter < result.RequeueAfter) {
+		result.RequeueAfter = fallbackResult.RequeueAfter
+	}
+	return result, nil
 }
 
 func updateDSStatusV2WithAgent(dsName string, ds *appsv1.DaemonSet, newStatus *datadoghqv1alpha1.DatadogAgentInternalStatus, updateTime metav1.Time, status metav1.ConditionStatus, reason, message string) {

internal/controller/datadogagentinternal/controller_reconcile_agent_test.go

@@ -1,23 +1,105 @@
 package datadogagentinternal
 
 import (
+	"context"
 	"testing"
 
+	appsv1 "k8s.io/api/apps/v1"
+	corev1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/runtime"
+	"k8s.io/apimachinery/pkg/util/intstr"
+	"k8s.io/client-go/tools/record"
 	"k8s.io/utils/ptr"
+	"sigs.k8s.io/controller-runtime/pkg/client/fake"
 
+	apicommon "github.com/DataDog/datadog-operator/api/datadoghq/common"
 	datadoghqv1alpha1 "github.com/DataDog/datadog-operator/api/datadoghq/v1alpha1"
 	datadoghqv2alpha1 "github.com/DataDog/datadog-operator/api/datadoghq/v2alpha1"
 	"github.com/DataDog/datadog-operator/internal/controller/datadogagent/component"
+	"github.com/DataDog/datadog-operator/internal/controller/datadogagent/defaults"
+	"github.com/DataDog/datadog-operator/internal/controller/datadogagent/feature"
+	"github.com/DataDog/datadog-operator/internal/controller/datadogagent/store"
 	"github.com/DataDog/datadog-operator/pkg/constants"
 	"github.com/DataDog/datadog-operator/pkg/kubernetes"
+	pkgtestutils "github.com/DataDog/datadog-operator/pkg/testutils"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 )
 
 const defaultProvider = kubernetes.DefaultProvider
 const gkeCosProvider = kubernetes.GKECloudProvider + "-" + kubernetes.GKECosType
 
+func TestReconcileV2AgentCreatesPreparedSurgeDaemonSet(t *testing.T) {
+	r, ddai := newPreparedRolloutReconciler(t, true)
+	status := &datadoghqv1alpha1.DatadogAgentInternalStatus{}
+
+	result, err := r.reconcileV2Agent(
+		context.Background(),
+		preparedRolloutRequiredComponents(),
+		nil,
+		ddai,
+		feature.NewResourceManagers(store.NewStore(ddai, nil)),
+		status,
+		defaultProvider,
+	)
+
+	require.NoError(t, err)
+	assert.Equal(t, resourceFallbackPollInterval, result.RequeueAfter)
+	daemonSets := &appsv1.DaemonSetList{}
+	require.NoError(t, r.client.List(context.Background(), daemonSets))
+	require.Len(t, daemonSets.Items, 1)
+	ds := &daemonSets.Items[0]
+	assert.Equal(t, preparedRolloutModeV1, ds.Spec.Template.Annotations[preparedRolloutModeAnnotation])
+	require.NotNil(t, ds.Spec.UpdateStrategy.RollingUpdate)
+	assert.Equal(t, intstr.FromInt(1), *ds.Spec.UpdateStrategy.RollingUpdate.MaxSurge)
+	assert.Equal(t, intstr.FromInt(0), *ds.Spec.UpdateStrategy.RollingUpdate.MaxUnavailable)
+}
+
+func newPreparedRolloutReconciler(t *testing.T, hostNetwork bool) (*Reconciler, *datadoghqv1alpha1.DatadogAgentInternal) {
+	t.Helper()
+	scheme := runtime.NewScheme()
+	require.NoError(t, corev1.AddToScheme(scheme))
+	require.NoError(t, appsv1.AddToScheme(scheme))
+	require.NoError(t, datadoghqv1alpha1.AddToScheme(scheme))
+	c := fake.NewClientBuilder().WithScheme(scheme).WithStatusSubresource(&appsv1.DaemonSet{}).Build()
+	one := intstr.FromInt(1)
+	ddai := pkgtestutils.NewDatadogAgentInternal("datadog-agent", "agent", nil)
+	ddai.UID = "ddai-uid"
+	ddai.Annotations = map[string]string{preparedRolloutModeAnnotation: preparedRolloutModeV1}
+	ddai.Spec.Features = &datadoghqv2alpha1.DatadogFeatures{}
+	ddai.Spec.Override = map[datadoghqv2alpha1.ComponentName]*datadoghqv2alpha1.DatadogAgentComponentOverride{
+		datadoghqv2alpha1.NodeAgentComponentName: {
+			HostNetwork: ptr.To(hostNetwork),
+			UpdateStrategy: &apicommon.UpdateStrategy{
+				Type: string(appsv1.RollingUpdateDaemonSetStrategyType),
+				RollingUpdate: &apicommon.RollingUpdate{
+					MaxSurge:       ptr.To(one),
+					MaxUnavailable: ptr.To(one),
+				},
+			},
+		},
+	}
+	defaults.DefaultDatadogAgentSpec(&ddai.Spec)
+	return &Reconciler{
+		client:    c,
+		apiReader: c,
+		scheme:    scheme,
+		recorder:  record.NewFakeRecorder(10),
+	}, ddai
+}
+
+func preparedRolloutRequiredComponents() feature.RequiredComponents {
+	return feature.RequiredComponents{Agent: feature.RequiredComponent{
+		IsRequired: ptr.To(true),
+		Containers: []apicommon.AgentContainerName{
+			apicommon.CoreAgentContainerName,
+			apicommon.TraceAgentContainerName,
+		},
+	}}
+}
+
 // func Test_getValidDaemonSetNames(t *testing.T) {
 // 	testCases := []struct {
 // 		name       string