Use Context for Timeout Handling (#72)

* Use Context for Timeout Handling

* Cancel inflight request during shutdown/timeout

* Refactoring timeout handling to use context timeouts

* Adding lifecycle unit tests

* Fix blocking logs channel issue

* Adding more lifecycle tests

* Tidy up go modules

* Stop sharing HandlerFunc across log servers (breaks tests)

* Move log server host override to config

* Add tests for client error handling

* Adding tests to improve coverage

* Add InitError and ExitError tests

* Coverage is almost there, fewmore tests

* Panic on 500 Internal Server Error, otherwise continue

checks/agent_version_check.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"io/ioutil"
@@ -18,7 +19,7 @@ type LayerAgentVersion struct {
 
 // We are only returning an error message when an out of date agent version is detected.
 // All other errors will result in a nil return value.
-func agentVersionCheck(conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig) error {
+func agentVersionCheck(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig) error {
 	if r.AgentVersion == "" {
 		return nil
 	}

checks/agent_version_check_test.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 	"testing"
@@ -14,9 +15,10 @@ func TestAgentVersion(t *testing.T) {
 	conf := config.Configuration{}
 	reg := api.RegistrationResponse{}
 	r := runtimeConfig{}
+	ctx := context.Background()
 
 	// No version set
-	err := agentVersionCheck(&conf, &reg, r)
+	err := agentVersionCheck(ctx, &conf, &reg, r)
 	assert.Nil(t, err)
 
 	// Error
@@ -33,11 +35,11 @@ func TestAgentVersion(t *testing.T) {
 	f, _ := os.Create(filepath.Join(testFile, r.agentVersionFile))
 	f.WriteString("10.1.0")
 
-	err = agentVersionCheck(&conf, &reg, r)
+	err = agentVersionCheck(ctx, &conf, &reg, r)
 	assert.EqualError(t, err, "Agent version out of date: v10.1.0, in order access up to date features please upgrade to the latest New Relic python layer that includes agent version v10.1.2")
 
 	// Success
 	r.AgentVersion = "10.1.0"
-	err = agentVersionCheck(&conf, &reg, r)
+	err = agentVersionCheck(ctx, &conf, &reg, r)
 	assert.Nil(t, err)
 }

checks/handler_check.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 	"strings"
 
@@ -16,7 +17,7 @@ type handlerConfigs struct {
 
 var handlerPath = "/var/task"
 
-func checkHandler(conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig) error {
+func handlerCheck(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig) error {
 	if r.language != "" {
 		h := handlerConfigs{
 			handlerName: reg.Handler,

checks/handler_check_test.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"os"
 	"path/filepath"
 	"testing"
@@ -57,15 +58,16 @@ func TestHandlerCheck(t *testing.T) {
 	conf := config.Configuration{}
 	reg := api.RegistrationResponse{}
 	r := runtimeConfigs[Node]
+	ctx := context.Background()
 
 	// No Runtime
-	err := checkHandler(&conf, &reg, runtimeConfig{})
+	err := handlerCheck(ctx, &conf, &reg, runtimeConfig{})
 	assert.Nil(t, err)
 
 	// Error
 	reg.Handler = testHandler
 	conf.NRHandler = config.EmptyNRWrapper
-	err = checkHandler(&conf, &reg, r)
+	err = handlerCheck(ctx, &conf, &reg, r)
 	assert.EqualError(t, err, "Missing handler file path/to/app.handler (NEW_RELIC_LAMBDA_HANDLER=Undefined)")
 
 	// Success
@@ -79,6 +81,6 @@ func TestHandlerCheck(t *testing.T) {
 
 	reg.Handler = testHandler
 	conf.NRHandler = config.EmptyNRWrapper
-	err = checkHandler(&conf, &reg, r)
+	err = handlerCheck(ctx, &conf, &reg, r)
 	assert.Nil(t, err)
 }

checks/sanity_check.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/newrelic/newrelic-lambda-extension/config"
@@ -21,12 +22,12 @@ var (
 )
 
 // sanityCheck checks for configuration that is either misplaced or in conflict
-func sanityCheck(conf *config.Configuration, res *api.RegistrationResponse, _ runtimeConfig) error {
+func sanityCheck(ctx context.Context, conf *config.Configuration, res *api.RegistrationResponse, _ runtimeConfig) error {
 	if util.AnyEnvVarsExist(awsLogIngestionEnvVars) {
 		return fmt.Errorf("Environment varaible '%s' is used by aws-log-ingestion and has no effect here. Recommend unsetting this environment variable within this function.", util.AnyEnvVarsExistString(awsLogIngestionEnvVars))
 	}
 
-	if credentials.IsSecretConfigured(conf) && util.EnvVarExists("NEW_RELIC_LICENSE_KEY") {
+	if credentials.IsSecretConfigured(ctx, conf) && util.EnvVarExists("NEW_RELIC_LICENSE_KEY") {
 		return fmt.Errorf("There is both a AWS Secrets Manager secret and a NEW_RELIC_LICENSE_KEY environment variable set. Recommend removing the NEW_RELIC_LICENSE_KEY environment variable and using the AWS Secrets Manager secret.")
 	}
 

checks/sanity_check_test.go

@@ -1,11 +1,13 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 	"os"
 	"testing"
 
 	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/aws/request"
 	"github.com/aws/aws-sdk-go/service/secretsmanager"
 	"github.com/aws/aws-sdk-go/service/secretsmanager/secretsmanageriface"
 	"github.com/newrelic/newrelic-lambda-extension/config"
@@ -19,7 +21,7 @@ type mockSecretManager struct {
 	secretsmanageriface.SecretsManagerAPI
 }
 
-func (mockSecretManager) GetSecretValue(*secretsmanager.GetSecretValueInput) (*secretsmanager.GetSecretValueOutput, error) {
+func (mockSecretManager) GetSecretValueWithContext(context.Context, *secretsmanager.GetSecretValueInput, ...request.Option) (*secretsmanager.GetSecretValueOutput, error) {
 	return &secretsmanager.GetSecretValueOutput{
 		SecretString: aws.String(`{"LicenseKey": "foo"}`),
 	}, nil
@@ -29,30 +31,32 @@ type mockSecretManagerErr struct {
 	secretsmanageriface.SecretsManagerAPI
 }
 
-func (mockSecretManagerErr) GetSecretValue(*secretsmanager.GetSecretValueInput) (*secretsmanager.GetSecretValueOutput, error) {
+func (mockSecretManagerErr) GetSecretValueWithContext(context.Context, *secretsmanager.GetSecretValueInput, ...request.Option) (*secretsmanager.GetSecretValueOutput, error) {
 	return nil, fmt.Errorf("Something went wrong")
 }
 
 func TestSanityCheck(t *testing.T) {
+	ctx := context.Background()
+
 	if util.AnyEnvVarsExist(awsLogIngestionEnvVars) {
-		assert.Error(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+		assert.Error(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 	} else {
-		assert.Nil(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+		assert.Nil(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 	}
 
 	os.Setenv("DEBUG_LOGGING_ENABLED", "1")
-	assert.Error(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+	assert.Error(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 	os.Unsetenv("DEBUG_LOGGING_ENABLED")
 
 	os.Unsetenv("NEW_RELIC_LICENSE_KEY")
 	credentials.OverrideSecretsManager(&mockSecretManager{})
-	assert.Nil(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+	assert.Nil(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 
 	os.Setenv("NEW_RELIC_LICENSE_KEY", "foobar")
 	defer os.Unsetenv("NEW_RELIC_LICENSE_KEY")
 	credentials.OverrideSecretsManager(&mockSecretManager{})
-	assert.Error(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+	assert.Error(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 
 	credentials.OverrideSecretsManager(&mockSecretManagerErr{})
-	assert.Nil(t, sanityCheck(&config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
+	assert.Nil(t, sanityCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, runtimeConfig{}))
 }

checks/startup_check.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -10,41 +11,40 @@ import (
 	"github.com/newrelic/newrelic-lambda-extension/util"
 )
 
-type checkFn func(*config.Configuration, *api.RegistrationResponse, runtimeConfig) error
+type checkFn func(context.Context, *config.Configuration, *api.RegistrationResponse, runtimeConfig) error
 
 type LogSender interface {
-	SendFunctionLogs(lines []logserver.LogLine) error
+	SendFunctionLogs(ctx context.Context, lines []logserver.LogLine) error
 }
 
 /// Register checks here
 var checks = []checkFn{
 	agentVersionCheck,
-	checkHandler,
+	handlerCheck,
 	sanityCheck,
 	vendorCheck,
 }
 
-func RunChecks(conf *config.Configuration, reg *api.RegistrationResponse, logSender LogSender) {
+func RunChecks(ctx context.Context, conf *config.Configuration, reg *api.RegistrationResponse, logSender LogSender) {
 	runtimeConfig, err := checkAndReturnRuntime()
 	if err != nil {
 		errLog := fmt.Sprintf("There was an issue querying for the latest agent version: %v", err)
 		util.Logln(errLog)
 	}
 
 	for _, check := range checks {
-		runCheck(conf, reg, runtimeConfig, logSender, check)
+		runCheck(ctx, conf, reg, runtimeConfig, logSender, check)
 	}
 }
 
-func runCheck(conf *config.Configuration, reg *api.RegistrationResponse, r runtimeConfig, logSender LogSender, check checkFn) error {
-	err := check(conf, reg, r)
-
+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)
 		util.Logln(errLog)
 
 		//Send a log line to NR as well
-		logSender.SendFunctionLogs([]logserver.LogLine{
+		logSender.SendFunctionLogs(ctx, []logserver.LogLine{
 			{
 				Time:      time.Now(),
 				RequestID: "0",

checks/startup_check_test.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 	"testing"
 
@@ -14,7 +15,7 @@ type TestLogSender struct {
 	sent []logserver.LogLine
 }
 
-func (c *TestLogSender) SendFunctionLogs(lines []logserver.LogLine) error {
+func (c *TestLogSender) SendFunctionLogs(ctx context.Context, lines []logserver.LogLine) error {
 	c.sent = append(c.sent, lines...)
 	return nil
 }
@@ -24,14 +25,15 @@ func TestRunCheck(t *testing.T) {
 	resp := api.RegistrationResponse{}
 	r := runtimeConfig{}
 	client := TestLogSender{}
+	ctx := context.Background()
 
 	tested := false
-	testCheck := func(conf *config.Configuration, resp *api.RegistrationResponse, r runtimeConfig) error {
+	testCheck := func(ctx context.Context, conf *config.Configuration, resp *api.RegistrationResponse, r runtimeConfig) error {
 		tested = true
 		return nil
 	}
 
-	result := runCheck(&conf, &resp, r, &client, testCheck)
+	result := runCheck(ctx, &conf, &resp, r, &client, testCheck)
 
 	assert.Equal(t, true, tested)
 	assert.Nil(t, result)
@@ -42,14 +44,15 @@ func TestRunCheckErr(t *testing.T) {
 	resp := api.RegistrationResponse{}
 	r := runtimeConfig{}
 	logSender := TestLogSender{}
+	ctx := context.Background()
 
 	tested := false
-	testCheck := func(conf *config.Configuration, resp *api.RegistrationResponse, r runtimeConfig) error {
+	testCheck := func(ctx context.Context, conf *config.Configuration, resp *api.RegistrationResponse, r runtimeConfig) error {
 		tested = true
 		return fmt.Errorf("Failure Test")
 	}
 
-	result := runCheck(&conf, &resp, r, &logSender, testCheck)
+	result := runCheck(ctx, &conf, &resp, r, &logSender, testCheck)
 
 	assert.Equal(t, true, tested)
 	assert.NotNil(t, result)
@@ -64,5 +67,6 @@ func TestRunChecks(t *testing.T) {
 
 	client = &mockClientError{}
 
-	RunChecks(c, r, l)
+	ctx := context.Background()
+	RunChecks(ctx, c, r, l)
 }

checks/vendor_check.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/newrelic/newrelic-lambda-extension/config"
@@ -10,7 +11,7 @@ import (
 
 // vendorCheck checks to see if the user included a vendored copy of the agent along
 // with their function while also using a layer that includes the agent
-func vendorCheck(_ *config.Configuration, _ *api.RegistrationResponse, r runtimeConfig) error {
+func vendorCheck(ctx context.Context, _ *config.Configuration, _ *api.RegistrationResponse, r runtimeConfig) error {
 
 	if util.PathExists(r.vendorAgentPath) && util.AnyPathsExist(r.layerAgentPaths) {
 		return fmt.Errorf("Vendored agent found at '%s', a layer already includes this agent at '%s'. Recommend using the layer agent to avoid unexpected agent behavior.", r.vendorAgentPath, util.AnyPathsExistString(r.layerAgentPaths))

checks/vendor_check_test.go

@@ -1,6 +1,7 @@
 package checks
 
 import (
+	"context"
 	"testing"
 
 	"github.com/newrelic/newrelic-lambda-extension/config"
@@ -10,24 +11,24 @@ import (
 )
 
 func TestVendorCheck(t *testing.T) {
-
 	n := runtimeConfigs[Node]
+	ctx := context.Background()
 
 	if !util.AnyPathsExist(n.layerAgentPaths) && !util.PathExists(n.vendorAgentPath) {
-		assert.Nil(t, vendorCheck(&config.Configuration{}, &api.RegistrationResponse{}, n))
+		assert.Nil(t, vendorCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, n))
 	}
 
 	if util.PathExists(n.layerAgentPaths[0]) && util.PathExists(n.vendorAgentPath) {
-		assert.Error(t, vendorCheck(&config.Configuration{}, &api.RegistrationResponse{}, n))
+		assert.Error(t, vendorCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, n))
 	}
 
 	p := runtimeConfigs[Python]
 
 	if !util.AnyPathsExist(p.layerAgentPaths) && !util.PathExists(p.vendorAgentPath) {
-		assert.Nil(t, vendorCheck(&config.Configuration{}, &api.RegistrationResponse{}, n))
+		assert.Nil(t, vendorCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, n))
 	}
 
 	if util.AnyPathsExist(p.layerAgentPaths) && util.PathExists(p.vendorAgentPath) {
-		assert.Error(t, vendorCheck(&config.Configuration{}, &api.RegistrationResponse{}, n))
+		assert.Error(t, vendorCheck(ctx, &config.Configuration{}, &api.RegistrationResponse{}, n))
 	}
 }