diff --git a/app/run.go b/app/run.go index 2c4903c2..4b021d64 100644 --- a/app/run.go +++ b/app/run.go @@ -96,8 +96,10 @@ func (app *App) Run(ctx context.Context) error { backgroundDataSendWg.Wait() if event.EventType == extension.Shutdown { app.logger.Infof("Exiting due to shutdown event with reason %s", event.ShutdownReason) - // Flush buffered logs if any - app.logsClient.FlushData(ctx, event.RequestID, event.InvokedFunctionArn, app.apmClient.LambdaDataChannel, true) + if app.logsClient != nil { + // Flush buffered logs if any + app.logsClient.FlushData(ctx, event.RequestID, event.InvokedFunctionArn, app.apmClient.LambdaDataChannel, true) + } // Since we have waited for the processEvent loop to finish we // already have received all the data we can from the agent. So, we // flush all the data to make sure that shutdown can correctly deduce @@ -127,8 +129,10 @@ func (app *App) Run(ctx context.Context) error { // that the underlying transport is reset for next invocation without // waiting for grace period if it got to unhealthy state. flushCtx, cancel := context.WithCancel(ctx) - // Flush buffered logs if any - app.logsClient.FlushData(ctx, event.RequestID, event.InvokedFunctionArn, app.apmClient.LambdaDataChannel, false) + if app.logsClient != nil { + // Flush buffered logs if any + app.logsClient.FlushData(ctx, event.RequestID, event.InvokedFunctionArn, app.apmClient.LambdaDataChannel, false) + } // Flush APM data now that the function invocation has completed app.apmClient.FlushAPMData(flushCtx) cancel() diff --git a/main_test.go b/main_test.go index 4bea2048..c71422a4 100644 --- a/main_test.go +++ b/main_test.go @@ -485,6 +485,29 @@ func TestStandardEventsChain(t *testing.T) { } } +// TestStandardEventsChainWithoutLogs checks a nominal sequence of events (fast APM server, only one standard event) +// with logs collection disabled +func TestStandardEventsChainWithoutLogs(t *testing.T) { + l, err := logger.New(logger.WithLevel(zapcore.DebugLevel)) + require.NoError(t, err) + + eventsChannel := newTestStructs(t) + apmServerInternals, _ := newMockApmServer(t, l) + logsapiAddr := randomAddr() + newMockLambdaServer(t, logsapiAddr, eventsChannel, l) + + eventsChain := []MockEvent{ + {Type: InvokeStandard, APMServerBehavior: TimelyResponse, ExecutionDuration: 1, Timeout: 5}, + } + eventQueueGenerator(eventsChain, eventsChannel) + select { + case <-runAppFull(t, logsapiAddr, true): + assert.Contains(t, apmServerInternals.Data, TimelyResponse) + case <-time.After(timeout): + t.Fatalf("timed out waiting for app to finish") + } +} + // TestFlush checks if the flushed param does not cause a panic or an unexpected behavior func TestFlush(t *testing.T) { l, err := logger.New(logger.WithLevel(zapcore.DebugLevel)) @@ -821,13 +844,21 @@ func TestMetrics(t *testing.T) { } func runApp(t *testing.T, logsapiAddr string) <-chan struct{} { + return runAppFull(t, logsapiAddr, false) +} + +func runAppFull(t *testing.T, logsapiAddr string, disableLogsAPI bool) <-chan struct{} { ctx, cancel := context.WithCancel(context.Background()) - app, err := app.New(ctx, + opts := []app.ConfigOption{ app.WithExtensionName("apm-lambda-extension"), app.WithLambdaRuntimeAPI(os.Getenv("AWS_LAMBDA_RUNTIME_API")), app.WithLogLevel("debug"), app.WithLogsapiAddress(logsapiAddr), - ) + } + if disableLogsAPI { + opts = append(opts, app.WithoutLogsAPI()) + } + app, err := app.New(ctx, opts...) require.NoError(t, err) go func() {