Add auto-annotation mutators on start up

main.go

@@ -4,6 +4,7 @@ package main
 
 import (
 	"crypto/tls"
+	"encoding/json"
 	"flag"
 	"fmt"
 	"os"
@@ -34,6 +35,7 @@ import (
 	"github.com/aws/amazon-cloudwatch-agent-operator/internal/webhook/podmutation"
 	"github.com/aws/amazon-cloudwatch-agent-operator/pkg/featuregate"
 	"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation"
+	"github.com/aws/amazon-cloudwatch-agent-operator/pkg/instrumentation/auto"
 	"github.com/aws/amazon-cloudwatch-agent-operator/pkg/sidecar"
 	// +kubebuilder:scaffold:imports
 )
@@ -88,6 +90,7 @@ func main() {
 		pprofAddr               string
 		agentImage              string
 		autoInstrumentationJava string
+		autoAnnotationConfigStr string
 		webhookPort             int
 		tlsOpt                  tlsConfig
 	)
@@ -97,6 +100,7 @@ func main() {
 	pflag.StringVar(&pprofAddr, "pprof-addr", "", "The address to expose the pprof server. Default is empty string which disables the pprof server.")
 	stringFlagOrEnv(&agentImage, "agent-image", "RELATED_IMAGE_COLLECTOR", fmt.Sprintf("%s:%s", cloudwatchAgentImageRepository, v.AmazonCloudWatchAgent), "The default CloudWatch Agent image. This image is used when no image is specified in the CustomResource.")
 	stringFlagOrEnv(&autoInstrumentationJava, "auto-instrumentation-java-image", "RELATED_IMAGE_AUTO_INSTRUMENTATION_JAVA", fmt.Sprintf("%s:%s", autoInstrumentationJavaImageRepository, v.AutoInstrumentationJava), "The default OpenTelemetry Java instrumentation image. This image is used when no image is specified in the CustomResource.")
+	pflag.StringVar(&autoAnnotationConfigStr, "auto-annotation-config", "", "The configuration for auto-annotation.")
 	pflag.Parse()
 
 	// set java instrumentation java image in environment variable to be used for default instrumentation
@@ -207,6 +211,23 @@ func main() {
 		os.Exit(1)
 	}
 
+	var autoAnnotationConfig auto.AnnotationConfig
+	if err := json.Unmarshal([]byte(autoAnnotationConfigStr), &autoAnnotationConfig); err != nil {
+		setupLog.Error(err, "unable to unmarshal auto-annotation config")
+	} else {
+		setupLog.Info("starting auto-annotator")
+		autoAnnotation := auto.NewAnnotationMutators(
+			mgr.GetClient(),
+			logger,
+			autoAnnotationConfig,
+			instrumentation.NewTypeSet(
+				instrumentation.TypeJava,
+				instrumentation.TypePython,
+			),
+		)
+		go autoAnnotation.MutateAll(ctx)
+	}
+
 	setupLog.Info("starting manager")
 	if err := mgr.Start(ctx); err != nil {
 		setupLog.Error(err, "problem running manager")

pkg/instrumentation/annotationmutator.go

@@ -0,0 +1,106 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+// AnnotationMutation is used to modify an annotation map.
+type AnnotationMutation interface {
+	// Mutate attempts to modify the annotations map. Returns whether the function changed the annotations.
+	Mutate(annotations map[string]string) bool
+}
+
+type insertAnnotationMutation struct {
+	insert     map[string]string
+	requireAll bool
+}
+
+func (m *insertAnnotationMutation) Mutate(annotations map[string]string) bool {
+	if m.requireAll && !m.validate(annotations) {
+		return false
+	}
+	var mutated bool
+	for key, value := range m.insert {
+		if _, ok := annotations[key]; !ok {
+			annotations[key] = value
+			mutated = true
+		}
+	}
+	return mutated
+}
+
+func (m *insertAnnotationMutation) validate(annotations map[string]string) bool {
+	for key := range m.insert {
+		if _, ok := annotations[key]; ok {
+			return false
+		}
+	}
+	return true
+}
+
+// NewInsertAnnotationMutation creates a new mutation that inserts annotations. If requireAll is enabled,
+// all provided annotation keys must be present for it to attempt to insert.
+func NewInsertAnnotationMutation(annotations map[string]string, requireAll bool) AnnotationMutation {
+	return &insertAnnotationMutation{insert: annotations, requireAll: requireAll}
+}
+
+type removeAnnotationMutation struct {
+	remove     []string
+	requireAll bool
+}
+
+func (m *removeAnnotationMutation) Mutate(annotations map[string]string) bool {
+	if m.requireAll && !m.validate(annotations) {
+		return false
+	}
+	var mutated bool
+	for _, key := range m.remove {
+		if _, ok := annotations[key]; ok {
+			delete(annotations, key)
+			mutated = true
+		}
+	}
+	return mutated
+}
+
+func (m *removeAnnotationMutation) validate(annotations map[string]string) bool {
+	for _, key := range m.remove {
+		if _, ok := annotations[key]; !ok {
+			return false
+		}
+	}
+	return true
+}
+
+// NewRemoveAnnotationMutation creates a new mutation that removes annotations. If requireAll is enabled,
+// all provided annotation keys must be present for it to attempt to remove them.
+func NewRemoveAnnotationMutation(annotations []string, requireAll bool) AnnotationMutation {
+	return &removeAnnotationMutation{remove: annotations, requireAll: requireAll}
+}
+
+type AnnotationMutator struct {
+	mutations []AnnotationMutation
+}
+
+// NewAnnotationMutator creates a mutator with the provided mutations that can mutate an Object's annotations.
+func NewAnnotationMutator(mutations []AnnotationMutation) AnnotationMutator {
+	return AnnotationMutator{mutations: mutations}
+}
+
+// Mutate modifies the object's annotations based on the mutator's mutations. Returns whether any of the
+// mutations changed the annotations.
+func (m *AnnotationMutator) Mutate(obj metav1.Object) bool {
+	annotations := obj.GetAnnotations()
+	if annotations == nil {
+		annotations = make(map[string]string)
+	}
+	var mutated bool
+	for _, mutation := range m.mutations {
+		mutated = mutated || mutation.Mutate(annotations)
+	}
+	obj.SetAnnotations(annotations)
+	return mutated
+}

pkg/instrumentation/annotationmutator_test.go

@@ -0,0 +1,127 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+)
+
+func TestMutateAnnotations(t *testing.T) {
+	testCases := map[string]struct {
+		annotations     map[string]string
+		mutations       []AnnotationMutation
+		wantAnnotations map[string]string
+		wantMutated     bool
+	}{
+		"TestInsert/Any": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewInsertAnnotationMutation(map[string]string{
+					"keyA": "3",
+					"keyC": "4",
+				}, false),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+				"keyC": "4",
+			},
+			wantMutated: true,
+		},
+		"TestInsert/All/Conflicts": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewInsertAnnotationMutation(map[string]string{
+					"keyA": "3",
+					"keyC": "4",
+				}, true),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			wantMutated: false,
+		},
+		"TestInsert/All/NoConflicts": {
+			annotations: nil,
+			mutations: []AnnotationMutation{
+				NewInsertAnnotationMutation(map[string]string{
+					"keyA": "3",
+					"keyC": "4",
+				}, true),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "3",
+				"keyC": "4",
+			},
+			wantMutated: true,
+		},
+		"TestRemove/Any": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewRemoveAnnotationMutation([]string{
+					"keyA",
+					"keyC",
+				}, false),
+			},
+			wantAnnotations: map[string]string{
+				"keyB": "2",
+			},
+			wantMutated: true,
+		},
+		"TestRemove/All/Conflicts": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewRemoveAnnotationMutation([]string{
+					"keyA",
+					"keyC",
+				}, true),
+			},
+			wantAnnotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			wantMutated: false,
+		},
+		"TestRemove/All/NoConflicts": {
+			annotations: map[string]string{
+				"keyA": "1",
+				"keyB": "2",
+			},
+			mutations: []AnnotationMutation{
+				NewRemoveAnnotationMutation([]string{
+					"keyA",
+					"keyB",
+				}, true),
+			},
+			wantAnnotations: map[string]string{},
+			wantMutated:     true,
+		},
+	}
+	for name, testCase := range testCases {
+		t.Run(name, func(t *testing.T) {
+			obj := metav1.ObjectMeta{
+				Annotations: testCase.annotations,
+			}
+			m := NewAnnotationMutator(testCase.mutations)
+			assert.Equal(t, testCase.wantMutated, m.Mutate(&obj))
+			assert.Equal(t, testCase.wantAnnotations, obj.GetAnnotations())
+		})
+	}
+}

pkg/instrumentation/annotationtype.go

@@ -0,0 +1,45 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+// Type is an enum for instrumentation types.
+type Type string
+
+// TypeSet is a map with Type keys.
+type TypeSet map[Type]any
+
+// NewTypeSet creates a new set of Type.
+func NewTypeSet(types ...Type) TypeSet {
+	s := make(TypeSet, len(types))
+	for _, t := range types {
+		s[t] = nil
+	}
+	return s
+}
+
+const (
+	TypeJava   Type = "java"
+	TypeNodeJS Type = "nodejs"
+	TypePython Type = "python"
+	TypeDotNet Type = "dotnet"
+	TypeGo     Type = "go"
+)
+
+// InjectAnnotationKey maps the instrumentation type to the inject annotation.
+func InjectAnnotationKey(instType Type) string {
+	switch instType {
+	case TypeJava:
+		return annotationInjectJava
+	case TypeNodeJS:
+		return annotationInjectNodeJS
+	case TypePython:
+		return annotationInjectPython
+	case TypeDotNet:
+		return annotationInjectDotNet
+	case TypeGo:
+		return annotationInjectGo
+	default:
+		return ""
+	}
+}

pkg/instrumentation/annotationtype_test.go

@@ -0,0 +1,38 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+package instrumentation
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestInjectAnnotationKey(t *testing.T) {
+	testCases := []struct {
+		instType Type
+		want     string
+	}{
+		{instType: TypeJava, want: annotationInjectJava},
+		{instType: TypeNodeJS, want: annotationInjectNodeJS},
+		{instType: TypePython, want: annotationInjectPython},
+		{instType: TypeDotNet, want: annotationInjectDotNet},
+		{instType: TypeGo, want: annotationInjectGo},
+		{instType: "unsupported", want: ""},
+	}
+	for _, testCase := range testCases {
+		assert.Equal(t, testCase.want, InjectAnnotationKey(testCase.instType))
+	}
+}
+
+func TestTypeSet(t *testing.T) {
+	types := []Type{TypeJava, TypeGo}
+	ts := NewTypeSet(types...)
+	_, ok := ts[TypeJava]
+	assert.True(t, ok)
+	_, ok = ts[TypeGo]
+	assert.True(t, ok)
+	_, ok = ts[TypePython]
+	assert.False(t, ok)
+}