poc - agent controlled rolling restarts

Makefile

@@ -194,6 +194,9 @@ agent-image:
 agent-image-slow:
 	@ scripts/dev/run_python.sh scripts/release/pipeline_main.py --parallel-factor 1 agent
 
+agent-image-custom:
+	@ scripts/dev/run_python.sh scripts/release/pipeline_main.py agent --version "latest" --custom-agent-url "https://mciuploads.s3.amazonaws.com/mms-automation/mongodb-mms-build-agent/builds/patches/68caf1b06da1570007e898b4/automation-agent/local/mongodb-mms-automation-agent-13.41.0.9783-1.linux_x86_64.tar.gz"
+
 operator-image:
 	@ scripts/dev/run_python.sh scripts/release/pipeline_main.py operator
 

build_info.json

@@ -253,6 +253,34 @@
         ]
       }
     },
+    "agent-sidecar": {
+      "dockerfile-path": "docker/mongodb-kubernetes-agent-sidecar/Dockerfile",
+      "patch": {
+        "repositories": ["268558157000.dkr.ecr.us-east-1.amazonaws.com/dev/nnguyen-kops/agent-sidecar"],
+        "platforms": [
+          "linux/amd64"
+        ]
+      },
+      "staging": {
+        "sign": true,
+        "latest-tag": true,
+        "repositories": ["268558157000.dkr.ecr.us-east-1.amazonaws.com/staging/nnguyen-kops/agent-sidecar"],
+        "platforms": [
+          "linux/arm64",
+          "linux/amd64"
+        ]
+      },
+      "release": {
+        "version": "1.0.0",
+        "sign": true,
+        "olm-tag": true,
+        "repositories": ["quay.io/mongodb/mongodb-kubernetes-agent-sidecar"],
+        "platforms": [
+          "linux/arm64",
+          "linux/amd64"
+        ]
+      }
+    },
     "agent": {
       "dockerfile-path": "docker/mongodb-agent/Dockerfile",
       "patch": {

controllers/om/api/http.go

@@ -220,6 +220,11 @@ func (client *Client) authorizeRequest(method, hostname, path string, request *r
 		return err
 	}
 
+	if resp.StatusCode == http.StatusOK {
+		// No need to authorize, server didn't challenge us
+		return nil
+	}
+
 	if resp.StatusCode != http.StatusUnauthorized {
 		return apierror.New(
 			xerrors.Errorf(

controllers/om/automation_config_test.go

@@ -599,10 +599,10 @@ func TestAssigningListsReassignsInDeployment(t *testing.T) {
 
 func TestAutomationConfigEquality(t *testing.T) {
 	deployment1 := NewDeployment()
-	deployment1.setReplicaSets([]ReplicaSet{NewReplicaSet("1", "5.0.0")})
+	deployment1.SetReplicaSets([]ReplicaSet{NewReplicaSet("1", "5.0.0")})
 
 	deployment2 := NewDeployment()
-	deployment2.setReplicaSets([]ReplicaSet{NewReplicaSet("2", "5.0.0")})
+	deployment2.SetReplicaSets([]ReplicaSet{NewReplicaSet("2", "5.0.0")})
 
 	authConfig := Auth{
 		Users: []*MongoDBUser{
@@ -1060,7 +1060,7 @@ func TestApplyInto(t *testing.T) {
 
 func changeTypes(deployment Deployment) error {
 	rs := deployment.GetReplicaSets()
-	deployment.setReplicaSets(rs)
+	deployment.SetReplicaSets(rs)
 	return nil
 }
 

controllers/om/deployment.go

@@ -85,7 +85,7 @@ func BuildDeploymentFromBytes(jsonBytes []byte) (Deployment, error) {
 func NewDeployment() Deployment {
 	ans := Deployment{}
 	ans.setProcesses(make([]Process, 0))
-	ans.setReplicaSets(make([]ReplicaSet, 0))
+	ans.SetReplicaSets(make([]ReplicaSet, 0))
 	ans.setShardedClusters(make([]ShardedCluster, 0))
 	ans.setMonitoringVersions(make([]interface{}, 0))
 	ans.setBackupVersions(make([]interface{}, 0))
@@ -405,7 +405,7 @@ func (d Deployment) RemoveReplicaSetByName(name string, log *zap.SugaredLogger)
 		}
 	}
 
-	d.setReplicaSets(toKeep)
+	d.SetReplicaSets(toKeep)
 
 	members := rs.Members()
 	processNames := make([]string, len(members))
@@ -992,12 +992,12 @@ func (d Deployment) GetReplicaSets() []ReplicaSet {
 	}
 }
 
-func (d Deployment) setReplicaSets(replicaSets []ReplicaSet) {
+func (d Deployment) SetReplicaSets(replicaSets []ReplicaSet) {
 	d["replicaSets"] = replicaSets
 }
 
 func (d Deployment) addReplicaSet(rs ReplicaSet) {
-	d.setReplicaSets(append(d.GetReplicaSets(), rs))
+	d.SetReplicaSets(append(d.GetReplicaSets(), rs))
 }
 
 func (d Deployment) getShardedClusters() []ShardedCluster {

controllers/operator/agents/upgrade.go

@@ -48,29 +48,7 @@ type ClientSecret struct {
 // happens and all existing MongoDBs are required to get agents upgraded (otherwise the "You need to upgrade the
 // automation agent before publishing other changes" error happens for automation config pushes from the Operator)
 func UpgradeAllIfNeeded(ctx context.Context, cs ClientSecret, omConnectionFactory om.ConnectionFactory, watchNamespace []string, isMulti bool) {
-	mux.Lock()
-	defer mux.Unlock()
-
-	if !time.Now().After(nextScheduledTime) {
-		return
-	}
-	log := zap.S()
-	log.Info("Performing a regular upgrade of Agents for all the MongoDB resources in the cluster...")
-
-	allMDBs, err := readAllMongoDBs(ctx, cs.Client, watchNamespace, isMulti)
-	if err != nil {
-		log.Errorf("Failed to read MongoDB resources to ensure Agents have the latest version: %s", err)
-		return
-	}
-
-	err = doUpgrade(ctx, cs.Client, cs.SecretClient, omConnectionFactory, allMDBs)
-	if err != nil {
-		log.Errorf("Failed to perform upgrade of Agents: %s", err)
-	}
-
-	log.Info("The upgrade of Agents for all the MongoDB resources in the cluster is finished.")
-
-	nextScheduledTime = nextScheduledTime.Add(pause)
+	return
 }
 
 // ScheduleUpgrade allows to reset the timer to Now() which makes sure the next MongoDB reconciliation will ensure

controllers/operator/common_controller.go

@@ -965,7 +965,7 @@ type PrometheusConfiguration struct {
 	prometheusCertHash string
 }
 
-func ReconcileReplicaSetAC(ctx context.Context, d om.Deployment, spec mdbv1.DbCommonSpec, lastMongodConfig map[string]interface{}, resourceName string, rs om.ReplicaSetWithProcesses, caFilePath string, internalClusterPath string, pc *PrometheusConfiguration, log *zap.SugaredLogger) error {
+func ReconcileReplicaSetAC(ctx context.Context, d om.Deployment, spec mdbv1.DbCommonSpec, lastMongodConfig map[string]interface{}, resourceName string, rs om.ReplicaSetWithProcesses, caFilePath string, internalClusterPath string, pc *PrometheusConfiguration, statefulSetVersion string, log *zap.SugaredLogger) error {
 	// it is not possible to disable internal cluster authentication once enabled
 	if d.ExistingProcessesHaveInternalClusterAuthentication(rs.Processes) && spec.Security.GetInternalClusterAuthenticationMode() == "" {
 		return xerrors.Errorf("cannot disable x509 internal cluster authentication")
@@ -980,6 +980,20 @@ func ReconcileReplicaSetAC(ctx context.Context, d om.Deployment, spec mdbv1.DbCo
 	d.AddMonitoringAndBackup(log, spec.GetSecurity().IsTLSEnabled(), caFilePath)
 	d.ConfigureTLS(spec.GetSecurity(), caFilePath)
 	d.ConfigureInternalClusterAuthentication(rs.GetProcessNames(), spec.GetSecurity().GetInternalClusterAuthenticationMode(), internalClusterPath)
+	// Set StatefulSet version in replica set member tags for rolling restart coordination
+	if statefulSetVersion != "" {
+		replicaSets := d.GetReplicaSets()
+		for _, replicaSet := range replicaSets {
+				for _, member := range replicaSet.Members() {
+					tags := member.Tags()
+					// Add StatefulSet version tag
+					tags["kubeStatefulSetVersion"] = statefulSetVersion
+					member["tags"] = tags
+				}
+		}
+		d.SetReplicaSets(replicaSets)
+		log.Infof("Set StatefulSet version in replica set member tags: %s", statefulSetVersion)
+	}
 
 	// if we don't set up a prometheus connection, then we don't want to set up prometheus for instance because we do not support it yet.
 	if pc != nil {

controllers/operator/construct/database_construction.go

@@ -109,6 +109,7 @@ type DatabaseStatefulSetOptions struct {
 	DatabaseNonStaticImage string
 	MongodbImage           string
 	AgentImage             string
+	AgentSidecarImage      string
 
 	Annotations map[string]string
 	VaultConfig vault.VaultConfiguration
@@ -514,7 +515,7 @@ func buildDatabaseStatefulSetConfigurationFunction(mdb databaseStatefulSetSource
 	}
 	podTemplateModifications = append(podTemplateModifications, staticMods...)
 
-	return statefulset.Apply(
+	statefulSetModifications := []statefulset.Modification{
 		// StatefulSet metadata
 		statefulset.WithLabels(ssLabels),
 		statefulset.WithName(stsName),
@@ -524,10 +525,14 @@ func buildDatabaseStatefulSetConfigurationFunction(mdb databaseStatefulSetSource
 		statefulset.WithServiceName(opts.ServiceName),
 		statefulset.WithReplicas(opts.Replicas),
 		statefulset.WithOwnerReference(opts.OwnerReference),
+		// Set OnDelete update strategy for agent-controlled rolling restarts
+		statefulset.WithUpdateStrategyType(appsv1.OnDeleteStatefulSetStrategyType),
 		volumeClaimFuncs,
 		shareProcessNs,
 		statefulset.WithPodSpecTemplate(podtemplatespec.Apply(podTemplateModifications...)),
-	)
+	}
+
+	return statefulset.Apply(statefulSetModifications...)
 }
 
 func buildPersistentVolumeClaimsFuncs(opts DatabaseStatefulSetOptions) (map[string]persistentvolumeclaim.Modification, []corev1.VolumeMount) {
@@ -714,15 +719,55 @@ func buildStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions, mdb
 		configureContainerSecurityContext,
 	)}
 
+	// Hardcoded init database image for local development
+	hardcodedInitDatabaseImage := "268558157000.dkr.ecr.us-east-1.amazonaws.com/dev/nnguyen-kops/mongodb-kubernetes-init-database:latest"
+
 	agentUtilitiesHolderModifications := []func(*corev1.Container){container.Apply(
 		container.WithName(util.AgentContainerUtilitiesName),
 		container.WithArgs([]string{""}),
-		container.WithImage(opts.InitDatabaseImage),
+		container.WithImage(hardcodedInitDatabaseImage),
 		container.WithEnvs(databaseEnvVars(opts)...),
 		container.WithCommand([]string{"bash", "-c", "touch /tmp/agent-utilities-holder_marker && tail -F -n0 /tmp/agent-utilities-holder_marker"}),
 		configureContainerSecurityContext,
 	)}
 
+	// Agent sidecar container for health monitoring and pod deletion
+	agentSidecarModifications := []func(*corev1.Container){container.Apply(
+		container.WithName("agent-sidecar"),
+		container.WithImage(opts.AgentSidecarImage),
+		container.WithEnvs([]corev1.EnvVar{
+			{
+				Name: "POD_NAME",
+				ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{
+						FieldPath: "metadata.name",
+					},
+				},
+			},
+			{
+				Name: "POD_NAMESPACE",
+				ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{
+						FieldPath: "metadata.namespace",
+					},
+				},
+			},
+			{
+				Name:  "AGENT_STATUS_FILEPATH",
+				Value: "/var/log/mongodb-mms-automation/agent-health-status.json",
+			},
+		}...),
+		container.WithVolumeMounts([]corev1.VolumeMount{
+			{
+				Name:      util.PvcNameData,
+				MountPath: "/var/log/mongodb-mms-automation",
+				SubPath:   util.PvcNameLogs,
+				ReadOnly:  true,
+			},
+		}),
+		configureContainerSecurityContext,
+	)}
+
 	if opts.HostNameOverrideConfigmapName != "" {
 		volumes = append(volumes, statefulset.CreateVolumeFromConfigMap(opts.HostNameOverrideConfigmapName, opts.HostNameOverrideConfigmapName))
 		hostnameOverrideModification := container.WithVolumeMounts([]corev1.VolumeMount{
@@ -734,6 +779,7 @@ func buildStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions, mdb
 		agentContainerModifications = append(agentContainerModifications, hostnameOverrideModification)
 		mongodContainerModifications = append(mongodContainerModifications, hostnameOverrideModification)
 		agentUtilitiesHolderModifications = append(agentUtilitiesHolderModifications, hostnameOverrideModification)
+		agentSidecarModifications = append(agentSidecarModifications, hostnameOverrideModification)
 	}
 
 	mods := []podtemplatespec.Modification{
@@ -745,6 +791,11 @@ func buildStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions, mdb
 		podtemplatespec.WithContainerByIndex(2, agentUtilitiesHolderModifications...),
 	}
 
+	// Add agent sidecar container if image is provided
+	if opts.AgentSidecarImage != "" {
+		mods = append(mods, podtemplatespec.WithContainerByIndex(3, agentSidecarModifications...))
+	}
+
 	return podtemplatespec.Apply(mods...)
 }
 
@@ -774,6 +825,44 @@ func buildNonStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions,
 		container.WithEnvs(readinessEnvironmentVariablesToEnvVars(opts.AgentConfig.ReadinessProbe.EnvironmentVariables)...),
 	)}
 
+	// Agent sidecar container for health monitoring and pod deletion (non-static architecture)
+	_, configureContainerSecurityContext := podtemplatespec.WithDefaultSecurityContextsModifications()
+	agentSidecarModifications := []func(*corev1.Container){container.Apply(
+		container.WithName("agent-sidecar"),
+		container.WithImage(opts.AgentSidecarImage),
+		container.WithEnvs([]corev1.EnvVar{
+			{
+				Name: "POD_NAME",
+				ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{
+						FieldPath: "metadata.name",
+					},
+				},
+			},
+			{
+				Name: "POD_NAMESPACE",
+				ValueFrom: &corev1.EnvVarSource{
+					FieldRef: &corev1.ObjectFieldSelector{
+						FieldPath: "metadata.namespace",
+					},
+				},
+			},
+			{
+				Name:  "AGENT_STATUS_FILEPATH",
+				Value: "/var/log/mongodb-mms-automation/agent-health-status.json",
+			},
+		}...),
+		container.WithVolumeMounts([]corev1.VolumeMount{
+			{
+				Name:      util.PvcNameData,
+				MountPath: "/var/log/mongodb-mms-automation",
+				SubPath:   util.PvcNameLogs,
+				ReadOnly:  true,
+			},
+		}),
+		configureContainerSecurityContext,
+	)}
+
 	if opts.HostNameOverrideConfigmapName != "" {
 		volumes = append(volumes, statefulset.CreateVolumeFromConfigMap(opts.HostNameOverrideConfigmapName, opts.HostNameOverrideConfigmapName))
 		hostnameOverrideModification := container.WithVolumeMounts([]corev1.VolumeMount{
@@ -784,6 +873,7 @@ func buildNonStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions,
 		})
 		initContainerModifications = append(initContainerModifications, hostnameOverrideModification)
 		databaseContainerModifications = append(databaseContainerModifications, hostnameOverrideModification)
+		agentSidecarModifications = append(agentSidecarModifications, hostnameOverrideModification)
 	}
 
 	mods := []podtemplatespec.Modification{
@@ -795,6 +885,11 @@ func buildNonStaticArchitecturePodTemplateSpec(opts DatabaseStatefulSetOptions,
 		podtemplatespec.WithInitContainerByIndex(0, initContainerModifications...),
 	}
 
+	// Add agent sidecar container if image is provided (non-static architecture)
+	if opts.AgentSidecarImage != "" {
+		mods = append(mods, podtemplatespec.WithContainerByIndex(1, agentSidecarModifications...))
+	}
+
 	return podtemplatespec.Apply(mods...)
 }
 
@@ -954,9 +1049,12 @@ func databaseScriptsVolumeMount(readOnly bool) corev1.VolumeMount {
 func buildDatabaseInitContainer(initDatabaseImage string) container.Modification {
 	_, configureContainerSecurityContext := podtemplatespec.WithDefaultSecurityContextsModifications()
 
+	// Hardcoded init database image for local development
+	hardcodedInitDatabaseImage := "268558157000.dkr.ecr.us-east-1.amazonaws.com/dev/nnguyen-kops/mongodb-kubernetes-init-database:latest"
+
 	return container.Apply(
 		container.WithName(InitDatabaseContainerName),
-		container.WithImage(initDatabaseImage),
+		container.WithImage(hardcodedInitDatabaseImage),
 		container.WithVolumeMounts([]corev1.VolumeMount{
 			databaseScriptsVolumeMount(false),
 		}),
@@ -1001,6 +1099,24 @@ func databaseEnvVars(opts DatabaseStatefulSetOptions) []corev1.EnvVar {
 		)
 	}
 
+	vars = append(vars, corev1.EnvVar{
+		Name: "POD_NAMESPACE",
+		ValueFrom: &corev1.EnvVarSource{
+			FieldRef: &corev1.ObjectFieldSelector{
+				FieldPath: "metadata.namespace",
+			},
+		},
+	})
+
+	vars = append(vars, corev1.EnvVar{
+		Name: "POD_NAME",
+		ValueFrom: &corev1.EnvVarSource{
+			FieldRef: &corev1.ObjectFieldSelector{
+				FieldPath: "metadata.name",
+			},
+		},
+	})
+
 	// This is only used for debugging
 	if useDebugAgent := os.Getenv(util.EnvVarDebug); useDebugAgent != "" { // nolint:forbidigo
 		zap.S().Debugf("running the agent in debug mode")
@@ -1013,6 +1129,12 @@ func databaseEnvVars(opts DatabaseStatefulSetOptions) []corev1.EnvVar {
 		vars = append(vars, corev1.EnvVar{Name: util.EnvVarAgentVersion, Value: agentVersion})
 	}
 
+	// Support for custom agent URL
+	if customAgentURL := os.Getenv(util.EnvVarCustomAgentURL); customAgentURL != "" { // nolint:forbidigo
+		zap.S().Debugf("using a custom agent URL: %s", customAgentURL)
+		vars = append(vars, corev1.EnvVar{Name: util.EnvVarCustomAgentURL, Value: customAgentURL})
+	}
+
 	// append any additional env vars specified.
 	vars = append(vars, opts.ExtraEnvs...)