From 3645243e01bd1f7147fcabe63fe3b4a085e55be7 Mon Sep 17 00:00:00 2001
From: 8naama <naama.bendalak@logz.io>
Date: Mon, 23 Sep 2024 15:58:16 +0300
Subject: [PATCH] changes following review

---
 logzio-lambda-extensions-logs/main.go | 1 +
 1 file changed, 1 insertion(+)

diff --git a/logzio-lambda-extensions-logs/main.go b/logzio-lambda-extensions-logs/main.go
index 3018294..5692a38 100644
--- a/logzio-lambda-extensions-logs/main.go
+++ b/logzio-lambda-extensions-logs/main.go
@@ -122,6 +122,7 @@ func main() {
 		select {
 		case <-ctx.Done():
 			logger.Info(printPrefix, "Received context done event")
+			logsApiLogger.Drain()
 			logsApiAgent.Shutdown()
 			logger.Info(printPrefix, "Exiting")
 			return