From 22d6d3e0a7189817f6f041466abfe199c193834e Mon Sep 17 00:00:00 2001
From: kruskal <99559985+kruskall@users.noreply.github.com>
Date: Thu, 25 Jul 2024 10:56:31 +0200
Subject: [PATCH 1/2] fix: only flush logs if logs collection is enabled

---
 app/run.go | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/app/run.go b/app/run.go
index cc029590..3033c87b 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()

From c03196f5d4dcc9ebd07f39d3514467018d89a3ea Mon Sep 17 00:00:00 2001
From: kruskal <99559985+kruskall@users.noreply.github.com>
Date: Thu, 25 Jul 2024 13:47:13 +0200
Subject: [PATCH 2/2] test: add standard events chain test with logs collection
 disabled

There is no coverage/test for the app when logs collection is disabled
Add a testcase to ensure standard events chain work even if logs
collection is disabled.
---
 main_test.go | 35 +++++++++++++++++++++++++++++++++--
 1 file changed, 33 insertions(+), 2 deletions(-)

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() {