diff --git a/apm-lambda-extension/extension/client.go b/apm-lambda-extension/extension/client.go index dcbcb957..91170ef0 100644 --- a/apm-lambda-extension/extension/client.go +++ b/apm-lambda-extension/extension/client.go @@ -160,7 +160,7 @@ func (e *Client) InitError(ctx context.Context, errorType string) (*StatusRespon } defer httpRes.Body.Close() - if httpRes.StatusCode != 200 { + if httpRes.StatusCode > 299 { return nil, fmt.Errorf("initialization error request failed with status %s", httpRes.Status) } res := StatusResponse{} @@ -187,7 +187,7 @@ func (e *Client) ExitError(ctx context.Context, errorType string) (*StatusRespon } defer httpRes.Body.Close() - if httpRes.StatusCode != 200 { + if httpRes.StatusCode > 299 { return nil, fmt.Errorf("exit error request failed with status %s", httpRes.Status) } res := StatusResponse{} diff --git a/apm-lambda-extension/main.go b/apm-lambda-extension/main.go index a7b63555..c1aa95d5 100644 --- a/apm-lambda-extension/main.go +++ b/apm-lambda-extension/main.go @@ -39,6 +39,7 @@ var ( /* --- elastic vars --- */ func main() { + ctx, cancel := context.WithCancel(context.Background()) // Trigger ctx.Done() in all relevant goroutines when main ends @@ -59,7 +60,14 @@ func main() { // register extension with AWS Extension API res, err := extensionClient.Register(ctx, extensionName) if err != nil { - panic(err) + status, errRuntime := extensionClient.InitError(ctx, err.Error()) + if errRuntime != nil { + panic(errRuntime) + } + extension.Log.Errorf("Error: %s", err) + extension.Log.Infof("Init error signal sent to runtime : %s", status) + extension.Log.Infof("Exiting") + return } extension.Log.Debugf("Register response: %v", extension.PrettyPrint(res)) @@ -104,7 +112,13 @@ func main() { extension.Log.Infof("Waiting for next event...") event, err := extensionClient.NextEvent(ctx) if err != nil { - extension.Log.Errorf("Error: %v\n. Exiting.", err) + status, err := extensionClient.ExitError(ctx, err.Error()) + if err != nil { + panic(err) + } + extension.Log.Errorf("Error: %s", err) + extension.Log.Infof("Exit signal sent to runtime : %s", status) + extension.Log.Infof("Exiting") return } extension.Log.Debug("Received event.")