Fix race condition for start up.

aws · Feb 21, 2024 · 02c98e9 · 02c98e9
1 parent 9b7274d
commit 02c98e9
Show file tree

Hide file tree

Showing 4 changed files with 99 additions and 22 deletions.
diff --git a/main.go b/main.go
@@ -3,13 +3,15 @@
 package main
 
 import (
+	"context"
 	"crypto/tls"
 	"encoding/json"
 	"flag"
 	"fmt"
 	"os"
 	"runtime"
 	"strings"
+	"time"
 
 	routev1 "github.com/openshift/api/route/v1"
 	monitoringv1 "github.com/prometheus-operator/prometheus-operator/pkg/apis/monitoring/v1"
@@ -211,8 +213,15 @@ func main() {
 			mgr.GetWebhookServer().Register("/mutate-v1-namespace", &webhook.Admission{
 				Handler: namespacemutation.NewWebhookHandler(decoder, autoAnnotationMutators),
 			})
-			setupLog.Info("Starting auto-annotation")
-			go autoAnnotationMutators.MutateAndPatchAll(ctx)
+			setupLog.Info("Auto-annotation is enabled")
+			go waitForWebhookServerStart(
+				ctx,
+				mgr.GetWebhookServer().StartedChecker(),
+				func(ctx context.Context) {
+					setupLog.Info("Applying auto-annotation")
+					autoAnnotationMutators.MutateAndPatchAll(ctx)
+				},
+			)
 		}
 	}
 
@@ -253,6 +262,23 @@ func main() {
 	}
 }
 
+func waitForWebhookServerStart(ctx context.Context, checker healthz.Checker, callback func(context.Context)) {
+	ticker := time.NewTicker(time.Second)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ticker.C:
+			if err := checker(nil); err == nil {
+				setupLog.Info("Webhook server has started")
+				callback(ctx)
+				return
+			}
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 // This function get the option from command argument (tlsConfig), check the validity through k8sapiflag
 // and set the config for webhook server.
 // refer to https://pkg.go.dev/k8s.io/component-base/cli/flag

diff --git a/pkg/instrumentation/annotationmutator.go b/pkg/instrumentation/annotationmutator.go
@@ -95,10 +95,11 @@ func (m *AnnotationMutator) Mutate(obj metav1.Object) bool {
 	if annotations == nil {
 		annotations = make(map[string]string)
 	}
-	var mutated bool
+	var anyMutated bool
 	for _, mutation := range m.mutations {
-		mutated = mutated || mutation.Mutate(annotations)
+		mutated := mutation.Mutate(annotations)
+		anyMutated = anyMutated || mutated
 	}
 	obj.SetAnnotations(annotations)
-	return mutated
+	return anyMutated
 }
diff --git a/pkg/instrumentation/annotationmutator_test.go b/pkg/instrumentation/annotationmutator_test.go
@@ -48,6 +48,22 @@ func TestMutateAnnotations(t *testing.T) {
 			},
 			wantMutated: true,
 		},
+		"TestInsert/Multiple": {
+			annotations: nil,
+			mutations: []AnnotationMutation{
+				NewInsertAnnotationMutation(map[string]string{
+					"keyA": "3",
+				}),
+				NewInsertAnnotationMutation(map[string]string{
+					"keyC": "4",
+				}),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "3",
+				"keyC": "4",
+			},
+			wantMutated: true,
+		},
 		"TestRemove/Conflicts": {
 			annotations: map[string]string{
 				"keyA": "1",
@@ -79,6 +95,40 @@ func TestMutateAnnotations(t *testing.T) {
 			wantAnnotations: map[string]string{},
 			wantMutated:     true,
 		},
+		"TestRemove/Multiple": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewRemoveAnnotationMutation([]string{
+					"keyA",
+				}),
+				NewRemoveAnnotationMutation([]string{
+					"keyB",
+				}),
+			},
+			wantAnnotations: map[string]string{},
+			wantMutated:     true,
+		},
+		"TestBoth": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewRemoveAnnotationMutation([]string{
+					"keyA",
+				}),
+				NewInsertAnnotationMutation(map[string]string{
+					"keyA": "3",
+				}),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "3",
+			},
+			wantMutated: true,
+		},
 	}
 	for name, testCase := range testCases {
 		t.Run(name, func(t *testing.T) {

diff --git a/pkg/instrumentation/defaultinstrumentation.go b/pkg/instrumentation/defaultinstrumentation.go
@@ -19,24 +19,24 @@ const (
 	defaultNamespace      = "default"
 	defaultKind           = "Instrumentation"
 
-	otelSampleEnabledKey                   = "OTEL_SMP_ENABLED" //TODO: remove in favor of new name once safe
-	otelSampleEnabledDefaultValue          = "true" //TODO: remove in favor of new name once safe
-	otelAppSignalsEnabledKey               = "OTEL_AWS_APP_SIGNALS_ENABLED"
-	otelAppSignalsEnabledDefaultValue      = "true"
-	otelTracesSamplerArgKey                = "OTEL_TRACES_SAMPLER_ARG"
-	otelTracesSamplerArgDefaultValue       = "endpoint=http://cloudwatch-agent.amazon-cloudwatch:2000"
-	otelTracesSamplerKey                   = "OTEL_TRACES_SAMPLER"
-	otelTracesSamplerDefaultValue          = "xray"
-	otelExporterOtlpProtocolKey            = "OTEL_EXPORTER_OTLP_PROTOCOL"
-	otelExporterOtlpProtocolValue          = "http/protobuf"
-	otelExporterTracesEndpointKey          = "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"
-	otelExporterTracesEndpointDefaultValue = "http://cloudwatch-agent.amazon-cloudwatch:4316/v1/traces"
-	otelExporterSmpEndpointKey             = "OTEL_AWS_SMP_EXPORTER_ENDPOINT" //TODO: remove in favor of new name once safe
-	otelExporterSmpEndpointDefaultValue    = "http://cloudwatch-agent.amazon-cloudwatch:4315" //TODO: remove in favor of new name once safe
-	otelExporterAppSignalsEndpointKey      = "OTEL_AWS_APP_SIGNALS_EXPORTER_ENDPOINT"
+	otelSampleEnabledKey                       = "OTEL_SMP_ENABLED" //TODO: remove in favor of new name once safe
+	otelSampleEnabledDefaultValue              = "true"             //TODO: remove in favor of new name once safe
+	otelAppSignalsEnabledKey                   = "OTEL_AWS_APP_SIGNALS_ENABLED"
+	otelAppSignalsEnabledDefaultValue          = "true"
+	otelTracesSamplerArgKey                    = "OTEL_TRACES_SAMPLER_ARG"
+	otelTracesSamplerArgDefaultValue           = "endpoint=http://cloudwatch-agent.amazon-cloudwatch:2000"
+	otelTracesSamplerKey                       = "OTEL_TRACES_SAMPLER"
+	otelTracesSamplerDefaultValue              = "xray"
+	otelExporterOtlpProtocolKey                = "OTEL_EXPORTER_OTLP_PROTOCOL"
+	otelExporterOtlpProtocolValue              = "http/protobuf"
+	otelExporterTracesEndpointKey              = "OTEL_EXPORTER_OTLP_TRACES_ENDPOINT"
+	otelExporterTracesEndpointDefaultValue     = "http://cloudwatch-agent.amazon-cloudwatch:4316/v1/traces"
+	otelExporterSmpEndpointKey                 = "OTEL_AWS_SMP_EXPORTER_ENDPOINT"                 //TODO: remove in favor of new name once safe
+	otelExporterSmpEndpointDefaultValue        = "http://cloudwatch-agent.amazon-cloudwatch:4315" //TODO: remove in favor of new name once safe
+	otelExporterAppSignalsEndpointKey          = "OTEL_AWS_APP_SIGNALS_EXPORTER_ENDPOINT"
 	otelExporterAppSignalsEndpointDefaultValue = "http://cloudwatch-agent.amazon-cloudwatch:4315"
-	otelExporterMetricKey                  = "OTEL_METRICS_EXPORTER"
-	otelExporterMetricDefaultValue         = "none"
+	otelExporterMetricKey                      = "OTEL_METRICS_EXPORTER"
+	otelExporterMetricDefaultValue             = "none"
 
 	otelPythonDistro                   = "OTEL_PYTHON_DISTRO"
 	otelPythonDistroDefaultValue       = "aws_distro"