From 1b645101bb83c4a9c515b34593e3c6ba0adce55e Mon Sep 17 00:00:00 2001
From: Saket Chaudhary <saketchaudharynit@gmail.com>
Date: Wed, 17 Apr 2024 02:00:50 +0530
Subject: [PATCH] add timeout

---
 checks/runtime_check.go      |  4 ++--
 checks/sanity_check.go       | 11 +++++++++--
 checks/startup_check.go      |  2 +-
 checks/startup_check_test.go |  2 +-
 main.go                      |  2 +-
 5 files changed, 14 insertions(+), 7 deletions(-)

diff --git a/checks/runtime_check.go b/checks/runtime_check.go
index 8bb3917..fe7d96b 100644
--- a/checks/runtime_check.go
+++ b/checks/runtime_check.go
@@ -26,9 +26,9 @@ func init() {
 		CheckRedirect: func(req *http.Request, via []*http.Request) error {
 			return http.ErrUseLastResponse
 		},
-		Timeout: time.Second * 10,
+		Timeout: time.Second * 2,
 	}
-	githubClient = github.NewClient(&http.Client{Timeout: time.Second * 10})
+	githubClient = github.NewClient(&http.Client{Timeout: time.Second * 2})
 }
 
 func checkAndReturnRuntime() (runtimeConfig, error) {
diff --git a/checks/sanity_check.go b/checks/sanity_check.go
index 1c139aa..0cf238f 100644
--- a/checks/sanity_check.go
+++ b/checks/sanity_check.go
@@ -3,6 +3,7 @@ package checks
 import (
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/newrelic/newrelic-lambda-extension/config"
 	"github.com/newrelic/newrelic-lambda-extension/credentials"
@@ -28,11 +29,17 @@ func sanityCheck(ctx context.Context, conf *config.Configuration, res *api.Regis
 	}
 
 	envKeyExists := util.EnvVarExists("NEW_RELIC_LICENSE_KEY")
-	isSecretConfigured := credentials.IsSecretConfigured(ctx, conf)
+	var timeout = 1 * time.Second
+	ctxSecret, cancelSecret := context.WithTimeout(ctx, timeout)
+	defer cancelSecret()
+	isSecretConfigured := credentials.IsSecretConfigured(ctxSecret, conf)
 
+	ctxSMParameter, cancelSMParameter := context.WithTimeout(ctx, timeout)
+	defer cancelSMParameter()
+	
 	isSSMParameterConfigured := false
 	if conf.LicenseKeySSMParameterName != "" {
-		isSSMParameterConfigured = credentials.IsSSMParameterConfigured(ctx, conf)
+		isSSMParameterConfigured = credentials.IsSSMParameterConfigured(ctxSMParameter, conf)
 	}
 	
 
diff --git a/checks/startup_check.go b/checks/startup_check.go
index df600a7..0e94e66 100644
--- a/checks/startup_check.go
+++ b/checks/startup_check.go
@@ -40,7 +40,7 @@ func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.Registr
 func runCheck(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig, logSender LogSender, check checkFn) error {
 	err := check(ctx, conf, reg, r)
 	if err != nil {
-		errLog := fmt.Sprintf("Startup check failed: %v", err)
+		errLog := fmt.Sprintf("Startup check warning: %v", err)
 		util.Logln(errLog)
 
 		//Send a log line to NR as well
diff --git a/checks/startup_check_test.go b/checks/startup_check_test.go
index f66359c..8336e86 100644
--- a/checks/startup_check_test.go
+++ b/checks/startup_check_test.go
@@ -65,7 +65,7 @@ func TestRunCheckErr(t *testing.T) {
 	assert.Equal(t, true, tested)
 	assert.NotNil(t, result)
 
-	assert.Equal(t, "Startup check failed: Failure Test", string(logSender.sent[0].Content))
+	assert.Equal(t, "Startup check warning: Failure Test", string(logSender.sent[0].Content))
 }
 
 func TestRunChecks(t *testing.T) {
diff --git a/main.go b/main.go
index 454b90b..fa8e59b 100644
--- a/main.go
+++ b/main.go
@@ -86,7 +86,7 @@ func main() {
 	}
 
 	// Attempt to find the license key for telemetry sending
-	var timeout = 5 * time.Second
+	var timeout = 1 * time.Second
 	ctxLicenseKey, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 	licenseKey, err := credentials.GetNewRelicLicenseKey(ctxLicenseKey, conf)