OTLP/ContainerInsights - KSM Node-metadata Enrichment (#2072)

Co-authored-by: Mitali Salvi <44349099+mitali-salvi@users.noreply.github.com>
Co-authored-by: Kaushik Surya <108111936+sky333999@users.noreply.github.com>

Author: 3 people

.gitignore

@@ -10,3 +10,4 @@ CWAGENT_VERSION
 terraform.*
 **/.terraform/*
 coverage.txt
+AGENTS.md

cfg/envconfig/envconfig.go

@@ -38,6 +38,7 @@ const (
 	CWAgentMergedOtelConfig     = "CWAGENT_MERGED_OTEL_CONFIG"
 	CWAgentLogsBackpressureMode = "CWAGENT_LOGS_BACKPRESSURE_MODE"
 	SystemMetricsEnabled        = "SYSTEM_METRICS_ENABLED"
+	OtelCIVersion               = "OTEL_CI_VERSION"
 
 	// confused deputy prevention related headers
 	AmzSourceAccount = "AMZ_SOURCE_ACCOUNT" // populates the "x-amz-source-account" header

extension/entitystore/ec2Info.go

@@ -25,8 +25,12 @@ const (
 )
 
 type EC2Info struct {
-	InstanceID string
-	AccountID  string
+	InstanceID       string
+	AccountID        string
+	InstanceType     string
+	ImageID          string
+	AvailabilityZone string
+	Hostname         string
 
 	// region is used while making call to describeTags Ec2 API for AutoScalingGroup
 	Region string
@@ -42,7 +46,7 @@ func (ei *EC2Info) initEc2Info() {
 		return
 	}
 	ei.logger.Debug("Initializing EC2Info")
-	if err := ei.setInstanceIDAccountID(); err != nil {
+	if err := ei.setEC2Metadata(); err != nil {
 		return
 	}
 	ei.logger.Debug("Finished initializing EC2Info")
@@ -60,11 +64,35 @@ func (ei *EC2Info) GetAccountID() string {
 	return ei.AccountID
 }
 
-func (ei *EC2Info) setInstanceIDAccountID() error {
+func (ei *EC2Info) GetInstanceType() string {
+	ei.mutex.RLock()
+	defer ei.mutex.RUnlock()
+	return ei.InstanceType
+}
+
+func (ei *EC2Info) GetImageID() string {
+	ei.mutex.RLock()
+	defer ei.mutex.RUnlock()
+	return ei.ImageID
+}
+
+func (ei *EC2Info) GetAvailabilityZone() string {
+	ei.mutex.RLock()
+	defer ei.mutex.RUnlock()
+	return ei.AvailabilityZone
+}
+
+func (ei *EC2Info) GetHostname() string {
+	ei.mutex.RLock()
+	defer ei.mutex.RUnlock()
+	return ei.Hostname
+}
+
+func (ei *EC2Info) setEC2Metadata() error {
 	for {
 		metadataDoc, err := ei.metadataProvider.Get(context.Background())
 		if err != nil {
-			ei.logger.Debug("Failed to get Instance ID / Account ID through metadata provider", zap.Error(err))
+			ei.logger.Debug("Failed to get EC2 metadata through metadata provider", zap.Error(err))
 			wait := time.NewTimer(1 * time.Minute)
 			select {
 			case <-ei.done:
@@ -74,14 +102,24 @@ func (ei *EC2Info) setInstanceIDAccountID() error {
 				continue
 			}
 		}
-		ei.logger.Debug("Successfully retrieved Instance ID and Account ID")
+		ei.logger.Debug("Successfully retrieved EC2 metadata")
+
+		hostname, err := ei.metadataProvider.Hostname(context.Background())
+		if err != nil {
+			ei.logger.Warn("Failed to get hostname from metadata provider, proceeding without it", zap.Error(err))
+		}
+
 		ei.mutex.Lock()
 		ei.InstanceID = metadataDoc.InstanceID
 		if idLength := len(ei.InstanceID); idLength > instanceIdSizeMax {
 			ei.logger.Warn("InstanceId length exceeds characters limit and will be ignored", zap.Int("length", idLength), zap.Int("character limit", instanceIdSizeMax))
 			ei.InstanceID = ""
 		}
 		ei.AccountID = metadataDoc.AccountID
+		ei.InstanceType = metadataDoc.InstanceType
+		ei.ImageID = metadataDoc.ImageID
+		ei.AvailabilityZone = metadataDoc.AvailabilityZone
+		ei.Hostname = hostname
 		ei.mutex.Unlock()
 		return nil
 	}

extension/entitystore/ec2Info_test.go

@@ -17,26 +17,28 @@ import (
 )
 
 var mockedInstanceIdentityDoc = &ec2metadata.EC2InstanceIdentityDocument{
-	InstanceID:   "i-01d2417c27a396e44",
-	AccountID:    "874389809020",
-	Region:       "us-east-1",
-	InstanceType: "m5ad.large",
-	ImageID:      "ami-09edd32d9b0990d49",
+	InstanceID:       "i-01d2417c27a396e44",
+	AccountID:        "874389809020",
+	Region:           "us-east-1",
+	InstanceType:     "m5ad.large",
+	ImageID:          "ami-09edd32d9b0990d49",
+	AvailabilityZone: "us-east-1a",
 }
 
 var mockedInstanceIdentityDocWithLargeInstanceId = &ec2metadata.EC2InstanceIdentityDocument{
-	InstanceID:   "i-01d2417c27a396e44394824728",
-	AccountID:    "874389809020",
-	Region:       "us-east-1",
-	InstanceType: "m5ad.large",
-	ImageID:      "ami-09edd32d9b0990d49",
+	InstanceID:       "i-01d2417c27a396e44394824728",
+	AccountID:        "874389809020",
+	Region:           "us-east-1",
+	InstanceType:     "m5ad.large",
+	ImageID:          "ami-09edd32d9b0990d49",
+	AvailabilityZone: "us-east-1a",
 }
 
 var (
 	tagVal3 = "ASG-1"
 )
 
-func TestSetInstanceIDAccountID(t *testing.T) {
+func TestSetEC2Metadata(t *testing.T) {
 	type args struct {
 		metadataProvider ec2metadataprovider.MetadataProvider
 	}
@@ -53,8 +55,12 @@ func TestSetInstanceIDAccountID(t *testing.T) {
 			},
 			wantErr: false,
 			want: EC2Info{
-				InstanceID: mockedInstanceIdentityDoc.InstanceID,
-				AccountID:  mockedInstanceIdentityDoc.AccountID,
+				InstanceID:       mockedInstanceIdentityDoc.InstanceID,
+				AccountID:        mockedInstanceIdentityDoc.AccountID,
+				InstanceType:     mockedInstanceIdentityDoc.InstanceType,
+				ImageID:          mockedInstanceIdentityDoc.ImageID,
+				AvailabilityZone: mockedInstanceIdentityDoc.AvailabilityZone,
+				Hostname:         "MockHostName",
 			},
 		},
 		{
@@ -64,8 +70,12 @@ func TestSetInstanceIDAccountID(t *testing.T) {
 			},
 			wantErr: false,
 			want: EC2Info{
-				InstanceID: "",
-				AccountID:  mockedInstanceIdentityDocWithLargeInstanceId.AccountID,
+				InstanceID:       "",
+				AccountID:        mockedInstanceIdentityDocWithLargeInstanceId.AccountID,
+				InstanceType:     mockedInstanceIdentityDocWithLargeInstanceId.InstanceType,
+				ImageID:          mockedInstanceIdentityDocWithLargeInstanceId.ImageID,
+				AvailabilityZone: mockedInstanceIdentityDocWithLargeInstanceId.AvailabilityZone,
+				Hostname:         "MockHostName",
 			},
 		},
 	}
@@ -76,11 +86,15 @@ func TestSetInstanceIDAccountID(t *testing.T) {
 				metadataProvider: tt.args.metadataProvider,
 				logger:           logger,
 			}
-			if err := ei.setInstanceIDAccountID(); (err != nil) != tt.wantErr {
-				t.Errorf("setInstanceIDAccountID() error = %v, wantErr %v", err, tt.wantErr)
+			if err := ei.setEC2Metadata(); (err != nil) != tt.wantErr {
+				t.Errorf("setEC2Metadata() error = %v, wantErr %v", err, tt.wantErr)
 			}
 			assert.Equal(t, tt.want.InstanceID, ei.GetInstanceID())
 			assert.Equal(t, tt.want.AccountID, ei.GetAccountID())
+			assert.Equal(t, tt.want.InstanceType, ei.GetInstanceType())
+			assert.Equal(t, tt.want.ImageID, ei.GetImageID())
+			assert.Equal(t, tt.want.AvailabilityZone, ei.GetAvailabilityZone())
+			assert.Equal(t, tt.want.Hostname, ei.GetHostname())
 		})
 	}
 }
@@ -157,3 +171,32 @@ func TestNotInitIfMetadataProviderIsEmpty(t *testing.T) {
 		})
 	}
 }
+
+func TestGettersReturnEmptyBeforeInit(t *testing.T) {
+	ei := &EC2Info{}
+	assert.Equal(t, "", ei.GetInstanceID())
+	assert.Equal(t, "", ei.GetAccountID())
+	assert.Equal(t, "", ei.GetInstanceType())
+	assert.Equal(t, "", ei.GetImageID())
+	assert.Equal(t, "", ei.GetAvailabilityZone())
+	assert.Equal(t, "", ei.GetHostname())
+}
+
+func TestHostnameFailureProceedsWithoutIt(t *testing.T) {
+	logger, _ := zap.NewDevelopment()
+	ei := &EC2Info{
+		metadataProvider: &mockMetadataProvider{
+			InstanceIdentityDocument: mockedInstanceIdentityDoc,
+			HostnameError:            true,
+		},
+		logger: logger,
+	}
+	err := ei.setEC2Metadata()
+	assert.NoError(t, err, "should succeed even when Hostname() fails")
+	// Hostname is empty but all other fields are populated
+	assert.Equal(t, "", ei.GetHostname())
+	assert.Equal(t, mockedInstanceIdentityDoc.InstanceID, ei.GetInstanceID())
+	assert.Equal(t, mockedInstanceIdentityDoc.InstanceType, ei.GetInstanceType())
+	assert.Equal(t, mockedInstanceIdentityDoc.ImageID, ei.GetImageID())
+	assert.Equal(t, mockedInstanceIdentityDoc.AvailabilityZone, ei.GetAvailabilityZone())
+}

extension/entitystore/extension.go

@@ -5,6 +5,7 @@ package entitystore
 
 import (
 	"context"
+	"os"
 	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
@@ -16,8 +17,11 @@ import (
 	"go.opentelemetry.io/collector/extension"
 	"go.uber.org/atomic"
 	"go.uber.org/zap"
+	"k8s.io/client-go/kubernetes"
+	"k8s.io/client-go/rest"
 
 	configaws "github.com/aws/amazon-cloudwatch-agent/cfg/aws"
+	"github.com/aws/amazon-cloudwatch-agent/cfg/envconfig"
 	"github.com/aws/amazon-cloudwatch-agent/internal/ec2metadataprovider"
 	"github.com/aws/amazon-cloudwatch-agent/internal/retryer"
 	"github.com/aws/amazon-cloudwatch-agent/plugins/processors/awsentity/entityattributes"
@@ -76,6 +80,10 @@ type EntityStore struct {
 	metadataprovider ec2metadataprovider.MetadataProvider
 
 	podTerminationCheckInterval time.Duration
+
+	// leaseWriter creates and renews a Kubernetes Lease with IMDS metadata
+	// for KSM node metadata enrichment
+	leaseWriter *LeaseWriter
 }
 
 var _ extension.Extension = (*EntityStore)(nil)
@@ -102,6 +110,12 @@ func (e *EntityStore) Start(ctx context.Context, host component.Host) error {
 		// https://github.com/kubernetes/cloud-provider-aws/issues/762
 		if e.kubernetesMode == "" {
 			go e.serviceprovider.startServiceProvider()
+		} else if getEnv(envconfig.OtelCIVersion) != "" {
+			e.startLeaseWriter()
+		} else {
+			e.logger.Debug("Skipping LeaseWriter - OTEL_CI_VERSION not set",
+				zap.String("OTEL_CI_VERSION", getEnv(envconfig.OtelCIVersion)),
+			)
 		}
 	}
 	if e.kubernetesMode != "" {
@@ -114,6 +128,9 @@ func (e *EntityStore) Start(ctx context.Context, host component.Host) error {
 }
 
 func (e *EntityStore) Shutdown(_ context.Context) error {
+	if e.leaseWriter != nil {
+		e.leaseWriter.Stop()
+	}
 	close(e.done)
 	if e.eksInfo != nil && e.eksInfo.podToServiceEnvMap != nil {
 		e.eksInfo.podToServiceEnvMap.Stop()
@@ -260,21 +277,60 @@ func (e *EntityStore) createServiceKeyAttributes(serviceAttr ServiceAttribute) m
 	return serviceKeyAttr
 }
 
-var getMetaDataProvider = func() ec2metadataprovider.MetadataProvider {
-	mdCredentialConfig := &configaws.CredentialConfig{}
-	return ec2metadataprovider.NewMetadataProvider(mdCredentialConfig.Credentials(), retryer.GetDefaultRetryNumber())
-}
+// startLeaseWriter initializes and starts the LeaseWriter for publishing
+// IMDS metadata as a Kubernetes Lease. Must be called after ec2Info is
+// initialized (the LeaseWriter's waitForEC2Info handles the race).
+func (e *EntityStore) startLeaseWriter() {
+	if e.leaseWriter != nil {
+		return
+	}
+	nodeName := getEnv("K8S_NODE_NAME")
+	if nodeName == "" {
+		e.logger.Error("K8S_NODE_NAME env var not set, skipping LeaseWriter startup")
+		return
+	}
+	namespace := getEnv("K8S_NAMESPACE")
+	if namespace == "" {
+		namespace = "amazon-cloudwatch"
+	}
 
-var getEC2Provider = func(region string, ec2CredentialConfig *configaws.CredentialConfig) ec2iface.EC2API {
-	ec2CredentialConfig.Region = region
-	return ec2.New(
-		ec2CredentialConfig.Credentials(),
-		&aws.Config{
-			LogLevel: configaws.SDKLogLevel(),
-			Logger:   configaws.SDKLogger{},
-		})
+	k8sConfig, err := getK8sConfig()
+	if err != nil {
+		e.logger.Error("Failed to create in-cluster K8s config for LeaseWriter", zap.Error(err))
+		return
+	}
+	clientset, err := kubernetes.NewForConfig(k8sConfig)
+	if err != nil {
+		e.logger.Error("Failed to create K8s clientset for LeaseWriter", zap.Error(err))
+		return
+	}
+
+	lw := NewLeaseWriter(&e.ec2Info, nodeName, namespace, clientset.CoordinationV1(), e.logger)
+	lw.Start()
+	e.leaseWriter = lw
 }
 
+var (
+	// Package-level vars for testability.
+	getEnv       = os.Getenv
+	getK8sConfig = rest.InClusterConfig
+
+	getMetaDataProvider = func() ec2metadataprovider.MetadataProvider {
+		mdCredentialConfig := &configaws.CredentialConfig{}
+		return ec2metadataprovider.NewMetadataProvider(mdCredentialConfig.Credentials(), retryer.GetDefaultRetryNumber())
+	}
+
+	getEC2Provider = func(region string, ec2CredentialConfig *configaws.CredentialConfig) ec2iface.EC2API {
+		ec2CredentialConfig.Region = region
+		return ec2.New(
+			ec2CredentialConfig.Credentials(),
+			&aws.Config{
+				LogLevel: configaws.SDKLogLevel(),
+				Logger:   configaws.SDKLogger{},
+			})
+	}
+)
+
 func addNonEmptyToMap(m map[string]*string, key, value string) {
 	if value != "" {
 		m[key] = aws.String(value)