From e98415ea2b6d3a9260cd26e7fce55416435e8f25 Mon Sep 17 00:00:00 2001
From: Robin Tang <rtang.cs@gmail.com>
Date: Thu, 7 Mar 2024 13:20:16 -0800
Subject: [PATCH] Update Logging handler (#238)

---
 lib/logger/logger.go | 22 ++++++++++++++++++----
 main.go              | 10 ++--------
 2 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/lib/logger/logger.go b/lib/logger/logger.go
index cada9916..dfac3df8 100644
--- a/lib/logger/logger.go
+++ b/lib/logger/logger.go
@@ -3,6 +3,7 @@ package logger
 import (
 	"log/slog"
 	"os"
+	"time"
 
 	"github.com/getsentry/sentry-go"
 	"github.com/lmittmann/tint"
@@ -12,10 +13,11 @@ import (
 	"github.com/artie-labs/reader/config"
 )
 
-func NewLogger(settings *config.Settings) (*slog.Logger, bool) {
+var handlersToTerminate []func()
+
+func NewLogger(settings *config.Settings) (*slog.Logger, func()) {
 	handler := tint.NewHandler(os.Stderr, &tint.Options{Level: slog.LevelInfo})
 
-	var loggingToSentry bool
 	if settings != nil && settings.Reporting != nil && settings.Reporting.Sentry != nil && settings.Reporting.Sentry.DSN != "" {
 		if err := sentry.Init(sentry.ClientOptions{Dsn: settings.Reporting.Sentry.DSN}); err != nil {
 			slog.New(handler).Warn("Failed to enable Sentry output", slog.Any("err", err))
@@ -24,19 +26,31 @@ func NewLogger(settings *config.Settings) (*slog.Logger, bool) {
 				handler,
 				slogsentry.Option{Level: slog.LevelError}.NewSentryHandler(),
 			)
-			loggingToSentry = true
+
+			slog.New(handler).Info("Sentry logger enabled")
+			handlersToTerminate = append(handlersToTerminate, func() {
+				sentry.Flush(2 * time.Second)
+			})
 		}
 	}
 
-	return slog.New(handler), loggingToSentry
+	return slog.New(handler), runHandlers
+}
+
+func runHandlers() {
+	for _, handlerToTerminate := range handlersToTerminate {
+		handlerToTerminate()
+	}
 }
 
 func Fatal(msg string, args ...any) {
 	slog.Error(msg, args...)
+	runHandlers()
 	os.Exit(1)
 }
 
 func Panic(msg string, args ...any) {
 	slog.Error(msg, args...)
+	runHandlers()
 	panic(msg)
 }
diff --git a/main.go b/main.go
index c92da243..e4cf99d3 100644
--- a/main.go
+++ b/main.go
@@ -5,10 +5,8 @@ import (
 	"flag"
 	"fmt"
 	"log/slog"
-	"time"
 
 	"github.com/artie-labs/transfer/lib/telemetry/metrics"
-	"github.com/getsentry/sentry-go"
 
 	"github.com/artie-labs/reader/config"
 	"github.com/artie-labs/reader/lib/kafkalib"
@@ -72,13 +70,9 @@ func main() {
 		logger.Fatal("Failed to read config file", slog.Any("err", err))
 	}
 
-	_logger, usingSentry := logger.NewLogger(cfg)
+	_logger, cleanUpHandlers := logger.NewLogger(cfg)
+	defer cleanUpHandlers()
 	slog.SetDefault(_logger)
-	if usingSentry {
-		defer sentry.Flush(2 * time.Second)
-		slog.Info("Sentry logger enabled")
-	}
-
 	ctx := context.Background()
 
 	statsD, err := setUpMetrics(cfg.Metrics)