Add a feature flag to user agent for container insights (#342)

* Add a feature flag to user agent for container insights * Update agentinfo.OutputPlugins directly instead of parsing. * Add TestUserAgent and fix test cases. * Add function to reset user agent string. * Reset the user agent string before constructing for container insights. Refactor and add test cases. * Move the regex for container insights to top level to reuse it. * Refactor to store userAgent string in map with group name as key. * Move test cases for container_insights flag to agentinfo. * Update test for container_insights and minior refacotr. Co-authored-by: Thomas Yang <yangth@amazon.com>
aws · Jan 25, 2022 · 0aad31d · 0aad31d
1 parent 63ef9c4
commit 0aad31d
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 25 deletions.
diff --git a/cfg/agentinfo/info.go b/cfg/agentinfo/info.go
@@ -8,16 +8,19 @@ import (
 	"io/ioutil"
 	"os"
 	"path/filepath"
+	"regexp"
 	"runtime"
 	"strings"
 
 	"github.com/aws/amazon-cloudwatch-agent/cfg/envconfig"
 )
 
-const versionFilename = "CWAGENT_VERSION"
-
-// We will fall back to a major version if no valid version file is found
-const fallbackVersion = "1"
+const (
+	containerInsightRegexp = "^/aws/.*containerinsights/.*/(performance|prometheus)$"
+	versionFilename        = "CWAGENT_VERSION"
+	// We will fall back to a major version if no valid version file is found
+	fallbackVersion = "1"
+)
 
 var isRunningAsRoot = func() bool {
 	return os.Getuid() == 0
@@ -29,7 +32,8 @@ var (
 	InputPlugins  []string
 	OutputPlugins []string
 
-	userAgent string
+	userAgentMap        = make(map[string]string)
+	ciCompiledRegexp, _ = regexp.Compile(containerInsightRegexp)
 )
 
 func Version() string {
@@ -50,26 +54,30 @@ func Build() string {
 	return BuildStr
 }
 
-func Plugins() string {
+func Plugins(groupName string) string {
 	outputs := strings.Join(OutputPlugins, " ")
 	inputs := strings.Join(InputPlugins, " ")
 
 	if !isRunningAsRoot() {
 		inputs += " run_as_user" // `inputs` is never empty, or agent will not start
 	}
+	if ciCompiledRegexp.MatchString(groupName) && !strings.Contains(outputs, "container_insights") {
+		outputs += " container_insights"
+	}
 
 	return fmt.Sprintf("inputs:(%s) outputs:(%s)", inputs, outputs)
 }
 
-func UserAgent() string {
-	if userAgent == "" {
-		ua := os.Getenv(envconfig.CWAGENT_USER_AGENT)
+func UserAgent(groupName string) string {
+	ua, found := userAgentMap[groupName]
+	if !found {
+		ua = os.Getenv(envconfig.CWAGENT_USER_AGENT)
 		if ua == "" {
-			ua = fmt.Sprintf("%s %s", FullVersion(), Plugins())
+			ua = fmt.Sprintf("%s %s", FullVersion(), Plugins(groupName))
 		}
-		userAgent = ua
+		userAgentMap[groupName] = ua
 	}
-	return userAgent
+	return ua
 }
 
 func FullVersion() string {

diff --git a/cfg/agentinfo/info_test.go b/cfg/agentinfo/info_test.go
@@ -13,6 +13,7 @@ import (
 	"testing"
 
 	"github.com/aws/amazon-cloudwatch-agent/cfg/envconfig"
+	"github.com/stretchr/testify/assert"
 )
 
 func TestVersionUseInjectedIfAvailable(t *testing.T) {
@@ -80,42 +81,78 @@ func TestPlugins(t *testing.T) {
 	OutputPlugins = []string{"x", "y", "z"}
 
 	isRunningAsRoot = func() bool { return true }
-	plugins := Plugins()
-	expected := fmt.Sprintf("inputs:(a b c) outputs:(x y z)")
+	plugins := Plugins("")
+	expected := "inputs:(a b c) outputs:(x y z)"
+	if plugins != expected {
+		t.Errorf("wrong plugins string constructed '%v', expecting '%v'", plugins, expected)
+	}
+
+	plugins = Plugins("/aws/ecs/containerinsights/ecs-cluster-name/performance")
+	expected = "inputs:(a b c) outputs:(x y z container_insights)"
 	if plugins != expected {
 		t.Errorf("wrong plugins string constructed '%v', expecting '%v'", plugins, expected)
 	}
 
 	isRunningAsRoot = func() bool { return false }
-	plugins = Plugins()
-	expected = fmt.Sprintf("inputs:(a b c run_as_user) outputs:(x y z)")
+	plugins = Plugins("")
+	expected = "inputs:(a b c run_as_user) outputs:(x y z)"
 	if plugins != expected {
 		t.Errorf("wrong plugins string constructed '%v', expecting '%v'", plugins, expected)
 	}
 }
 
 func TestUserAgent(t *testing.T) {
-	userAgent = ""
 	VersionStr = "VSTR"
 	BuildStr = "BSTR"
 	InputPlugins = []string{"a", "b", "c"}
 	OutputPlugins = []string{"x", "y", "z"}
 
 	isRunningAsRoot = func() bool { return true }
 
-	ua := UserAgent()
-	expected := fmt.Sprintf("CWAgent/VSTR (%v; %v; %v) BSTR inputs:(a b c) outputs:(x y z)", runtime.Version(), runtime.GOOS, runtime.GOARCH)
-	if ua != expected {
-		t.Errorf("wrong UserAgent string constructed '%v', expecting '%v'", ua, expected)
+	tests := []struct {
+		name         string
+		logGroupName string
+		expected     string
+		errorMessage string
+	}{
+		{
+			"non container insights",
+			"test-group",
+			fmt.Sprintf("CWAgent/VSTR (%v; %v; %v) BSTR inputs:(a b c) outputs:(x y z)", runtime.Version(), runtime.GOOS, runtime.GOARCH),
+			"wrong UserAgent string constructed",
+		},
+		{
+			"container insights EKS",
+			"/aws/containerinsights/eks-cluster-name/performance",
+			fmt.Sprintf("CWAgent/VSTR (%v; %v; %v) BSTR inputs:(a b c) outputs:(x y z container_insights)", runtime.Version(), runtime.GOOS, runtime.GOARCH),
+			"\"container_insights\" flag shoould be in the outputs plugin list in container insights mode",
+		},
+		{
+			"container insights ECS",
+			"/aws/ecs/containerinsights/ecs-cluster-name/performance",
+			fmt.Sprintf("CWAgent/VSTR (%v; %v; %v) BSTR inputs:(a b c) outputs:(x y z container_insights)", runtime.Version(), runtime.GOOS, runtime.GOARCH),
+			"\"container_insights\" flag shoould be in the outputs plugin list in container insights mode",
+		},
+		{
+			"container insights prometheus",
+			"/aws/containerinsights/cluster-name/prometheus",
+			fmt.Sprintf("CWAgent/VSTR (%v; %v; %v) BSTR inputs:(a b c) outputs:(x y z container_insights)", runtime.Version(), runtime.GOOS, runtime.GOARCH),
+			"\"container_insights\" flag shoould be in the outputs plugin list in container insights mode",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			assert.Equal(t, tc.expected, UserAgent(tc.logGroupName), tc.errorMessage)
+		})
 	}
 }
 
 func TestUserAgentEnvOverride(t *testing.T) {
-	userAgent = ""
 	os.Setenv(envconfig.CWAGENT_USER_AGENT, "CUSTOM CWAGENT USER AGENT")
 	expected := "CUSTOM CWAGENT USER AGENT"
 
-	ua := UserAgent()
+	ua := UserAgent("TestUserAgentEnvOverride")
 	if ua != expected {
 		t.Errorf("UserAgent should use value configured in environment variable CWAGENT_USER_AGENT, but '%v' found, expecting '%v'", ua, expected)
 	}

diff --git a/plugins/outputs/cloudwatch/cloudwatch.go b/plugins/outputs/cloudwatch/cloudwatch.go
@@ -145,7 +145,7 @@ func (c *CloudWatch) Connect() error {
 		})
 
 	svc.Handlers.Build.PushBackNamed(handlers.NewRequestCompressionHandler([]string{opPutLogEvents, opPutMetricData}))
-	svc.Handlers.Build.PushBackNamed(handlers.NewCustomHeaderHandler("User-Agent", agentinfo.UserAgent()))
+	svc.Handlers.Build.PushBackNamed(handlers.NewCustomHeaderHandler("User-Agent", agentinfo.UserAgent("")))
 
 	//Format unique roll up list
 	c.RollupDimensions = GetUniqueRollupList(c.RollupDimensions)

diff --git a/plugins/outputs/cloudwatchlogs/cloudwatchlogs.go b/plugins/outputs/cloudwatchlogs/cloudwatchlogs.go
@@ -127,7 +127,7 @@ func (c *CloudWatchLogs) getDest(t Target) *cwDest {
 		},
 	)
 	client.Handlers.Build.PushBackNamed(handlers.NewRequestCompressionHandler([]string{"PutLogEvents"}))
-	client.Handlers.Build.PushBackNamed(handlers.NewCustomHeaderHandler("User-Agent", agentinfo.UserAgent()))
+	client.Handlers.Build.PushBackNamed(handlers.NewCustomHeaderHandler("User-Agent", agentinfo.UserAgent(t.Group)))
 
 	pusher := NewPusher(t, client, c.ForceFlushInterval.Duration, maxRetryTimeout, c.Log)
 	cwd := &cwDest{pusher: pusher, retryer: logThrottleRetryer}