Add downward-node-labels propagation to Kibana, Logstash and Elastic …

pkg/apis/agent/v1alpha1/agent_types.go

@@ -14,6 +14,7 @@ import (
 
 	commonv1 "github.com/elastic/cloud-on-k8s/v3/pkg/apis/common/v1"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/version"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/utils/nodelabels"
 )
 
 const (
@@ -24,6 +25,9 @@ const (
 	AgentContainerName = "agent"
 	// FleetServerServiceAccount is the Elasticsearch service account to be used to authenticate.
 	FleetServerServiceAccount commonv1.ServiceAccountName = "fleet-server"
+	// DownwardNodeLabelsAnnotation holds an optional comma-separated list of expected node labels
+	// to be set as annotations on the Elastic Agent Pods.
+	DownwardNodeLabelsAnnotation = nodelabels.DownwardNodeLabelsAnnotation
 )
 
 // AgentSpec defines the desired state of the Agent
@@ -312,6 +316,17 @@ func (a *Agent) IsMarkedForDeletion() bool {
 	return !a.DeletionTimestamp.IsZero()
 }
 
+// DownwardNodeLabels returns the node labels to copy as annotations on the Agent Pods,
+// as declared via the DownwardNodeLabelsAnnotation annotation.
+func (a *Agent) DownwardNodeLabels() []string {
+	return nodelabels.FromAnnotations(a.Annotations)
+}
+
+// HasDownwardNodeLabels returns true if node labels are expected to be propagated to the Agent Pods.
+func (a *Agent) HasDownwardNodeLabels() bool {
+	return len(a.DownwardNodeLabels()) > 0
+}
+
 func (a *Agent) AssociationStatusMap(typ commonv1.AssociationType) commonv1.AssociationStatusMap {
 	switch typ {
 	case commonv1.ElasticsearchAssociationType:

pkg/apis/kibana/v1/kibana_types.go

@@ -13,6 +13,7 @@ import (
 
 	commonv1 "github.com/elastic/cloud-on-k8s/v3/pkg/apis/common/v1"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/version"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/utils/nodelabels"
 )
 
 const (
@@ -22,6 +23,9 @@ const (
 	Kind = "Kibana"
 	// KibanaServiceAccount is the Elasticsearch service account to be used to authenticate.
 	KibanaServiceAccount commonv1.ServiceAccountName = "kibana"
+	// DownwardNodeLabelsAnnotation holds an optional comma-separated list of expected node labels
+	// to be set as annotations on the Kibana Pods.
+	DownwardNodeLabelsAnnotation = nodelabels.DownwardNodeLabelsAnnotation
 )
 
 // +kubebuilder:object:root=true
@@ -154,6 +158,17 @@ func (k *Kibana) IsMarkedForDeletion() bool {
 	return !k.DeletionTimestamp.IsZero()
 }
 
+// DownwardNodeLabels returns the node labels to copy as annotations on the Kibana Pods,
+// as declared via the DownwardNodeLabelsAnnotation annotation.
+func (k *Kibana) DownwardNodeLabels() []string {
+	return nodelabels.FromAnnotations(k.Annotations)
+}
+
+// HasDownwardNodeLabels returns true if node labels are expected to be propagated to the Kibana Pods.
+func (k *Kibana) HasDownwardNodeLabels() bool {
+	return len(k.DownwardNodeLabels()) > 0
+}
+
 func (k *Kibana) SecureSettings() []commonv1.SecretSource {
 	return k.Spec.SecureSettings
 }

pkg/apis/logstash/v1alpha1/logstash_types.go

@@ -13,6 +13,7 @@ import (
 
 	commonv1 "github.com/elastic/cloud-on-k8s/v3/pkg/apis/common/v1"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/hash"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/utils/nodelabels"
 )
 
 type LogstashHealth string
@@ -35,6 +36,10 @@ const (
 	// 1) all Pods are Ready, and
 	// 2) any associations are configured and established
 	LogstashGreenHealth LogstashHealth = "green"
+
+	// DownwardNodeLabelsAnnotation holds an optional comma-separated list of expected node labels
+	// to be set as annotations on the Logstash Pods.
+	DownwardNodeLabelsAnnotation = nodelabels.DownwardNodeLabelsAnnotation
 )
 
 // LogstashSpec defines the desired state of Logstash
@@ -220,6 +225,17 @@ func (l *Logstash) IsMarkedForDeletion() bool {
 	return !l.DeletionTimestamp.IsZero()
 }
 
+// DownwardNodeLabels returns the node labels to copy as annotations on the Logstash Pods,
+// as declared via the DownwardNodeLabelsAnnotation annotation.
+func (l *Logstash) DownwardNodeLabels() []string {
+	return nodelabels.FromAnnotations(l.Annotations)
+}
+
+// HasDownwardNodeLabels returns true if node labels are expected to be propagated to the Logstash Pods.
+func (l *Logstash) HasDownwardNodeLabels() bool {
+	return len(l.DownwardNodeLabels()) > 0
+}
+
 // GetObservedGeneration will return the observedGeneration from the Elastic Logstash's status.
 func (l *Logstash) GetObservedGeneration() int64 {
 	return l.Status.ObservedGeneration

pkg/controller/agent/driver.go

@@ -24,6 +24,7 @@ import (
 	commonassociation "github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/association"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/certificates"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/defaults"
+	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/nodelabels"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/operator"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/reconciler"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/tracing"
@@ -173,7 +174,21 @@ func internalReconcile(params Params) (*reconciler.Results, agentv1alpha1.AgentS
 	if err != nil {
 		return results.WithError(err), params.Status
 	}
-	return reconcilePodVehicle(params, podTemplate)
+	podVehicleResults, status := reconcilePodVehicle(params, podTemplate)
+	results.WithResults(podVehicleResults)
+
+	// Patch the Pods to add the expected node labels as annotations. Record the error, if any, but do not stop the
+	// reconciliation loop as we don't want to prevent other updates from being applied.
+	results.WithResults(nodelabels.AnnotatePods(
+		params.Context,
+		params.Client,
+		params.Agent.Namespace,
+		map[string]string{NameLabelName: params.Agent.Name},
+		params.Agent.DownwardNodeLabels(),
+		params.Agent.Name,
+	))
+
+	return results, status
 }
 
 func reconcileService(params Params) (*corev1.Service, error) {

pkg/controller/agent/pod.go

@@ -30,6 +30,7 @@ import (
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/container"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/defaults"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/labels"
+	commonnodelabels "github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/nodelabels"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/reconciler"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/tracing"
 	"github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/version"
@@ -189,6 +190,12 @@ func buildPodTemplate(params Params, fleetCerts *certificates.CertificatesSecret
 	}
 	vols = append(vols, caAssocVols...)
 
+	// Changes to the downward-node-labels annotation must roll the Agent Pods so the new annotations
+	// are re-applied on scheduling.
+	if params.Agent.HasDownwardNodeLabels() {
+		_, _ = configHash.Write([]byte(params.Agent.Annotations[agentv1alpha1.DownwardNodeLabelsAnnotation]))
+	}
+
 	podMeta := params.Meta.Merge(metadata.Metadata{
 		Labels:      map[string]string{VersionLabelName: spec.Version},
 		Annotations: map[string]string{ConfigHashAnnotationName: fmt.Sprint(configHash.Sum32())},
@@ -207,6 +214,21 @@ func buildPodTemplate(params Params, fleetCerts *certificates.CertificatesSecret
 			}},
 		)
 
+	// Block Pod start on the operator patching the expected node-label annotations onto the
+	// Pod, so the Agent container does not start while those annotations are missing from
+	// the downward-API annotations file.
+	if params.Agent.HasDownwardNodeLabels() {
+		downwardAPIVolume := commonnodelabels.DownwardAPIVolume()
+		waitInit, err := commonnodelabels.WaitForAnnotationsInitContainer(params.Agent.DownwardNodeLabels())
+		if err != nil {
+			return corev1.PodTemplateSpec{}, err
+		}
+		builder = builder.
+			WithVolumes(downwardAPIVolume.Volume()).
+			WithInitContainers(waitInit).
+			WithInitContainerDefaults()
+	}
+
 	return builder.PodTemplate, nil
 }
 

pkg/controller/common/nodelabels/initcontainer.go

@@ -0,0 +1,74 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package nodelabels
+
+import (
+	"bytes"
+	"fmt"
+	"strings"
+	"text/template"
+
+	corev1 "k8s.io/api/core/v1"
+
+	commonvolume "github.com/elastic/cloud-on-k8s/v3/pkg/controller/common/volume"
+	esvolume "github.com/elastic/cloud-on-k8s/v3/pkg/controller/elasticsearch/volume"
+)
+
+// WaitForAnnotationsContainerName is the name of the init container that blocks Pod start
+// until the operator has patched the expected node-derived annotations onto the Pod.
+const WaitForAnnotationsContainerName = "elastic-internal-wait-for-node-labels"
+
+var waitScriptTemplate = template.Must(template.New("").Parse(
+	`#!/usr/bin/env bash
+set -eu
+expected_annotations=({{ .ExpectedAnnotations }})
+annotations_file={{ .AnnotationsFile }}
+function annotations_exist() {
+  for expected_annotation in "${expected_annotations[@]}"; do
+    if ! grep -qE "^${expected_annotation}=" "${annotations_file}"; then
+      return 1
+    fi
+  done
+  return 0
+}
+echo "Waiting for the following annotations to be set on Pod: {{ .ExpectedAnnotations }}"
+while ! annotations_exist; do sleep 2; done
+echo "All expected annotations are set."
+`))
+
+// DownwardAPIVolume returns the downward API volume that exposes the Pod annotations file
+// under the path polled by the wait-for-annotations init container.
+func DownwardAPIVolume() commonvolume.DownwardAPI {
+	return commonvolume.DownwardAPI{}.WithAnnotations(true)
+}
+
+// WaitForAnnotationsInitContainer builds an init container that blocks until the operator
+// has patched all expectedAnnotations onto the Pod's metadata.annotations. This mirrors the
+// behavior of the Elasticsearch prepare-fs init container and prevents the main container
+// from starting while the downward-API annotations file is still missing labels that users
+// may consume via `valueFrom.fieldRef`.
+//
+// The image and resources are left unset so they are inherited from the main container via
+// PodTemplateBuilder.WithInitContainerDefaults. Callers must also add the volume returned
+// by DownwardAPIVolume to the Pod.
+func WaitForAnnotationsInitContainer(expectedAnnotations []string) (corev1.Container, error) {
+	buf := bytes.Buffer{}
+	if err := waitScriptTemplate.Execute(&buf, struct {
+		ExpectedAnnotations string
+		AnnotationsFile     string
+	}{
+		ExpectedAnnotations: strings.Join(expectedAnnotations, " "),
+		AnnotationsFile:     fmt.Sprintf("%s/%s", esvolume.DownwardAPIMountPath, esvolume.AnnotationsFile),
+	}); err != nil {
+		return corev1.Container{}, err
+	}
+	return corev1.Container{
+		Name:    WaitForAnnotationsContainerName,
+		Command: []string{"bash", "-c", buf.String()},
+		VolumeMounts: []corev1.VolumeMount{
+			DownwardAPIVolume().VolumeMount(),
+		},
+	}, nil
+}

pkg/controller/common/nodelabels/initcontainer_test.go

@@ -0,0 +1,46 @@
+// Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+// or more contributor license agreements. Licensed under the Elastic License 2.0;
+// you may not use this file except in compliance with the Elastic License 2.0.
+
+package nodelabels
+
+import (
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestWaitForAnnotationsInitContainer(t *testing.T) {
+	c, err := WaitForAnnotationsInitContainer([]string{"topology.kubernetes.io/zone", "topology.kubernetes.io/region"})
+	require.NoError(t, err)
+	assert.Equal(t, WaitForAnnotationsContainerName, c.Name)
+	require.Len(t, c.VolumeMounts, 1)
+	assert.Equal(t, "downward-api", c.VolumeMounts[0].Name)
+	assert.Equal(t, "/mnt/elastic-internal/downward-api", c.VolumeMounts[0].MountPath)
+	assert.True(t, c.VolumeMounts[0].ReadOnly)
+
+	require.Len(t, c.Command, 3)
+	assert.Equal(t, "bash", c.Command[0])
+	assert.Equal(t, "-c", c.Command[1])
+	script := c.Command[2]
+	assert.Contains(t, script, "topology.kubernetes.io/zone topology.kubernetes.io/region")
+	assert.Contains(t, script, "/mnt/elastic-internal/downward-api/annotations")
+	// Each expected annotation is matched at the beginning of a line followed by '=':
+	assert.Contains(t, script, `grep -qE "^${expected_annotation}="`)
+	// No image/resources are set so they are inherited from the main container.
+	assert.Equal(t, "", c.Image)
+	assert.Empty(t, c.Resources.Limits)
+	assert.Empty(t, c.Resources.Requests)
+}
+
+func TestDownwardAPIVolume_IncludesAnnotations(t *testing.T) {
+	v := DownwardAPIVolume().Volume()
+	require.NotNil(t, v.VolumeSource.DownwardAPI)
+	paths := make([]string, 0, len(v.VolumeSource.DownwardAPI.Items))
+	for _, item := range v.VolumeSource.DownwardAPI.Items {
+		paths = append(paths, item.Path)
+	}
+	assert.Contains(t, strings.Join(paths, ","), "annotations")
+}