From 9410bbf07307e4bf0fc3171ab07d3423ed7acc0e Mon Sep 17 00:00:00 2001
From: Charles Cheng <ctycheng@amazon.com>
Date: Wed, 27 Nov 2024 21:31:02 +0000
Subject: [PATCH] Add metrics to acs for eni provisioning workflow monitoring

---
 ecs-agent/acs/session/session.go       | 3 +--
 ecs-agent/api/ecs/client/ecs_client.go | 8 +++++++-
 ecs-agent/api/ecs/interface.go         | 4 ++++
 ecs-agent/metrics/constants.go         | 6 +++---
 go.mod                                 | 3 +++
 5 files changed, 18 insertions(+), 6 deletions(-)
 create mode 100644 go.mod

diff --git a/ecs-agent/acs/session/session.go b/ecs-agent/acs/session/session.go
index 04f0ecc36dc..4cc8dc7c922 100644
--- a/ecs-agent/acs/session/session.go
+++ b/ecs-agent/acs/session/session.go
@@ -242,9 +242,7 @@ func (s *session) startSessionOnce(ctx context.Context) error {
 		s.firstDiscoverPollEndpointTime = time.Now()
 	}
 
-	discoverPollEndpointMetric := s.metricsFactory.New(metrics.ACSDiscoverPollEndpointDurationName)
 	acsEndpoint, err := s.ecsClient.DiscoverPollEndpoint(s.containerInstanceARN)
-	discoverPollEndpointMetric.Done(err)
 
 	if err != nil {
 		logger.Error("ACS: Unable to discover poll endpoint", logger.Fields{
@@ -253,6 +251,7 @@ func (s *session) startSessionOnce(ctx context.Context) error {
 		})
 		return err
 	}
+	s.metricsFactory.New(metrics.DiscoverPollEndpointDurationName).WithGauge(s.ecsClient.GetDiscoverPollEndpointDuration()).Done(nil)
 
 	client := s.clientFactory.New(
 		s.acsURL(acsEndpoint),
diff --git a/ecs-agent/api/ecs/client/ecs_client.go b/ecs-agent/api/ecs/client/ecs_client.go
index 6e57ba6b27e..eee36f04a89 100644
--- a/ecs-agent/api/ecs/client/ecs_client.go
+++ b/ecs-agent/api/ecs/client/ecs_client.go
@@ -77,6 +77,7 @@ type ecsClient struct {
 	shouldExcludeIPv6PortBinding     bool
 	sascCustomRetryBackoff           func(func() error) error
 	stscAttachmentCustomRetryBackoff func(func() error) error
+	discoverPollEndpointDuration     time.Duration
 }
 
 // NewECSClient creates a new ECSClient interface object.
@@ -747,7 +748,7 @@ func (client *ecsClient) discoverPollEndpoint(containerInstanceArn string,
 			}
 		}
 	}
-
+	discoverPollEndpointStartTime := time.Now()
 	// Cache miss or expired, invoke the ECS DiscoverPollEndpoint API.
 	logger.Debug("Invoking DiscoverPollEndpoint", logger.Fields{
 		field.ContainerInstanceARN: containerInstanceArn,
@@ -777,6 +778,7 @@ func (client *ecsClient) discoverPollEndpoint(containerInstanceArn string,
 		return nil, err
 	}
 
+	client.discoverPollEndpointDuration = time.Since(discoverPollEndpointStartTime)
 	// Cache the response from ECS.
 	client.pollEndpointCache.Set(containerInstanceArn, output)
 	return output, nil
@@ -870,3 +872,7 @@ func trimString(inputString string, maxLen int) string {
 		return inputString
 	}
 }
+
+func (client *ecsClient) GetDiscoverPollEndpointDuration() time.Duration {
+	return client.discoverPollEndpointDuration
+}
diff --git a/ecs-agent/api/ecs/interface.go b/ecs-agent/api/ecs/interface.go
index fb3b3fab53f..6c962dc2b90 100644
--- a/ecs-agent/api/ecs/interface.go
+++ b/ecs-agent/api/ecs/interface.go
@@ -14,6 +14,8 @@
 package ecs
 
 import (
+	"time"
+
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/aws/request"
 
@@ -62,6 +64,8 @@ type ECSClient interface {
 	UpdateContainerInstancesState(instanceARN, status string) error
 	// GetHostResources retrieves a map that map the resource name to the corresponding resource
 	GetHostResources() (map[string]*ecs.Resource, error)
+	// GetDiscoverPollEndpointDuration retrieves the time it takes for the DiscoverPollEndpoint call
+	GetDiscoverPollEndpointDuration() time.Duration
 }
 
 // ECSSDK is an interface that specifies the subset of the AWS Go SDK's ECS
diff --git a/ecs-agent/metrics/constants.go b/ecs-agent/metrics/constants.go
index 021dacd944a..0be004b5deb 100644
--- a/ecs-agent/metrics/constants.go
+++ b/ecs-agent/metrics/constants.go
@@ -47,9 +47,9 @@ const (
 	TCSDisconnectTimeoutMetricName = agentAvailabilityNamespace + ".TCSDisconnectTimeout"
 
 	// ACS Session Metrics
-	acsStartSessionNamespace            = "ACSStartSession"
-	ACSDiscoverPollEndpointDurationName = acsStartSessionNamespace + ".DiscoverPollEndpointDuration"
-	ACSConnectionMetricDurationName     = acsStartSessionNamespace + ".ACSConnectEndpointDuration"
+	acsStartSessionNamespace         = "ACSStartSession"
+	DiscoverPollEndpointDurationName = acsStartSessionNamespace + ".DiscoverPollEndpointDuration"
+	ACSConnectionMetricDurationName  = acsStartSessionNamespace + ".ACSConnectEndpointDuration"
 
 	dbClientMetricNamespace                 = "Data"
 	GetNetworkConfigurationByTaskMetricName = dbClientMetricNamespace + ".GetNetworkConfigurationByTask"
diff --git a/go.mod b/go.mod
new file mode 100644
index 00000000000..bd833852a62
--- /dev/null
+++ b/go.mod
@@ -0,0 +1,3 @@
+module github.com/aws/amazon-ecs-agent/ecs-agent
+
+go 1.22.7