Run Extension checks conditionally (#235)

* Add env var to ignore extension checks
newrelic · Oct 15, 2024 · 2682333 · 2682333
1 parent ac7ad99
commit 2682333
Show file tree

Hide file tree

Showing 5 changed files with 125 additions and 6 deletions.
diff --git a/checks/startup_check.go b/checks/startup_check.go
@@ -18,11 +18,11 @@ type LogSender interface {
 }
 
 /// Register checks here
-var checks = []checkFn{
-	agentVersionCheck,
-	handlerCheck,
-	sanityCheck,
-	vendorCheck,
+var checks = map[string]checkFn{
+    "agent":         agentVersionCheck, 
+    "handler":       handlerCheck,
+    "sanity":        sanityCheck, 
+    "vendor":        vendorCheck, 
 }
 
 func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, logSender LogSender) {
@@ -32,7 +32,10 @@ func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.Registr
 		util.Logln(errLog)
 	}
 
-	for _, check := range checks {
+	for checkName, check := range checks {
+		if conf.IgnoreExtensionChecks[checkName] {
+			continue
+		}
 		runCheck(ctx, conf, reg, runtimeConfig, logSender, check)
 	}
 }

diff --git a/checks/startup_check_test.go b/checks/startup_check_test.go
@@ -78,3 +78,25 @@ func TestRunChecks(t *testing.T) {
 	ctx := context.Background()
 	RunChecks(ctx, c, r, l)
 }
+
+func TestRunChecksIgnoreExtensionChecks(t *testing.T) {
+	c := &config.Configuration{IgnoreExtensionChecks: map[string]bool{"agent": true}}
+	r := &api.RegistrationResponse{}
+	l := &TestLogSender{}
+
+	client = &mockClientError{}
+
+	ctx := context.Background()
+	RunChecks(ctx, c, r, l)
+}
+
+func TestRunChecksIgnoreExtensionChecksAll(t *testing.T) {
+	c := &config.Configuration{IgnoreExtensionChecks: map[string]bool{"all": true}}
+	r := &api.RegistrationResponse{}
+	l := &TestLogSender{}
+
+	client = &mockClientError{}
+
+	ctx := context.Background()
+	RunChecks(ctx, c, r, l)
+}
diff --git a/config/config.go b/config/config.go
@@ -21,6 +21,7 @@ var EmptyNRWrapper = "Undefined"
 type Configuration struct {
 	TestingOverride            bool // ignores envioronment specific details when running unit tests
 	ExtensionEnabled           bool
+	IgnoreExtensionChecks      map[string]bool
 	LogsEnabled                bool
 	SendFunctionLogs           bool
 	CollectTraceID             bool
@@ -37,9 +38,42 @@ type Configuration struct {
 	ClientTimeout              time.Duration
 }
 
+func parseIgnoredExtensionChecks(nrIgnoreExtensionChecksOverride bool, nrIgnoreExtensionChecksStr string) map[string]bool {
+	ignoredChecks := make(map[string]bool)
+
+	if !nrIgnoreExtensionChecksOverride || nrIgnoreExtensionChecksStr == "" {
+		return nil
+	}
+
+	validChecks := map[string]bool{
+		"agent":   true,
+		"handler": true,
+		"sanity":  true,
+		"vendor":  true,
+	}
+
+	ignoredChecksStr := strings.ToLower(nrIgnoreExtensionChecksStr)
+
+	if ignoredChecksStr == "all" {
+		ignoredChecks["all"] = true
+		return ignoredChecks
+	}
+
+	checks := strings.Split(ignoredChecksStr, ",")
+	for _, check := range checks {
+		trimmedCheck := strings.TrimSpace(check)
+		if trimmedCheck != "" && validChecks[trimmedCheck] {
+			ignoredChecks[trimmedCheck] = true
+		}
+	}
+
+	return ignoredChecks
+}
+
 func ConfigurationFromEnvironment() *Configuration {
 	nrEnabledStr, nrEnabledOverride := os.LookupEnv("NEW_RELIC_ENABLED")
 	nrEnabledRubyStr, nrEnabledRubyOverride := os.LookupEnv("NEW_RELIC_AGENT_ENABLED")
+	nrIgnoreExtensionChecksStr, nrIgnoreExtensionChecksOverride := os.LookupEnv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
 	enabledStr, extensionEnabledOverride := os.LookupEnv("NEW_RELIC_LAMBDA_EXTENSION_ENABLED")
 	licenseKey, lkOverride := os.LookupEnv("NEW_RELIC_LICENSE_KEY")
 	licenseKeySecretId, lkSecretOverride := os.LookupEnv("NEW_RELIC_LICENSE_KEY_SECRET")
@@ -93,6 +127,8 @@ func ConfigurationFromEnvironment() *Configuration {
 		ret.LicenseKeySSMParameterName = licenseKeySSMParameterName
 	}
 
+	ret.IgnoreExtensionChecks = parseIgnoredExtensionChecks(nrIgnoreExtensionChecksOverride, nrIgnoreExtensionChecksStr)
+
 	if nrOverride {
 		ret.NRHandler = nrHandler
 	} else {

diff --git a/config/config_test.go b/config/config_test.go
@@ -99,6 +99,60 @@ func TestConfigurationFromEnvironmentNRAgentEnabled(t *testing.T) {
 	assert.Equal(t, conf.ExtensionEnabled, false)
 }
 
+func TestConfigurationFromEnvironmentExtensionChecks(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "agent,handler,dummy")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, conf.IgnoreExtensionChecks["agent"], true)
+	assert.Equal(t, conf.IgnoreExtensionChecks["handler"], true)
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 2)
+}
+
+func TestConfigurationFromEnvironmentExtensionChecksAll(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "ALL")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, conf.IgnoreExtensionChecks["agent"], false)
+	assert.Equal(t, conf.IgnoreExtensionChecks["handler"], false)
+	assert.Equal(t, conf.IgnoreExtensionChecks["sanity"], false)
+	assert.Equal(t, conf.IgnoreExtensionChecks["vendor"], false)
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 1)
+}
+
+func TestConfigurationFromEnvironmentExtensionChecksIncorrectString(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "incorrect,valuess,...,,")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 0)
+}
+
+func TestConfigurationFromEnvironmentExtensionChecksIncorrectStringWithAll(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "incorrect,valuess,...,ALL,")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 0)
+}
+
+func TestConfigurationFromEnvironmentExtensionChecksIncorrectStringAll(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "All,ALL,...,ALL,")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 0)
+}
+
+func TestConfigurationFromEnvironmentExtensionChecksEmptyString(t *testing.T) {
+	os.Setenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS", "")
+	defer os.Unsetenv("NEW_RELIC_IGNORE_EXTENSION_CHECKS")
+
+	conf := ConfigurationFromEnvironment()
+	assert.Equal(t, len(conf.IgnoreExtensionChecks), 0)
+}
+
 func TestConfigurationFromEnvironmentSecretId(t *testing.T) {
 	os.Setenv("NEW_RELIC_LICENSE_KEY_SECRET", "secretId")
 	defer os.Unsetenv("NEW_RELIC_LICENSE_KEY_SECRET")

diff --git a/main.go b/main.go
@@ -137,6 +137,10 @@ func main() {
 
 	// Run startup checks
 	go func() {
+		if conf.IgnoreExtensionChecks["all"] {
+			util.Debugf("Ignoring all extension checks")
+			return
+		}
 		checks.RunChecks(ctx, conf, registrationResponse, telemetryClient)
 	}()