From dfc9c86a144fd79f0a20ab4520a5d18052b37041 Mon Sep 17 00:00:00 2001
From: Arvindh <arvindh91@gmail.com>
Date: Fri, 26 Jan 2024 22:35:18 +0530
Subject: [PATCH] Fix: messaging handler logger

Signed-off-by: Arvindh <arvindh91@gmail.com>
---
 pkg/messaging/handler/logging.go | 54 +++++++++++++++++---------------
 1 file changed, 29 insertions(+), 25 deletions(-)

diff --git a/pkg/messaging/handler/logging.go b/pkg/messaging/handler/logging.go
index 6b92ca67b..4b1eceaf1 100644
--- a/pkg/messaging/handler/logging.go
+++ b/pkg/messaging/handler/logging.go
@@ -22,43 +22,51 @@ type loggingMiddleware struct {
 
 // AuthConnect implements session.Handler.
 func (lm *loggingMiddleware) AuthConnect(ctx context.Context) (err error) {
-	return lm.logAction(ctx, "AuthConnect", nil)
+	defer lm.logAction(ctx, "AuthConnect", nil, time.Now(), err)
+	return lm.svc.AuthConnect(ctx)
 }
 
 // AuthPublish implements session.Handler.
 func (lm *loggingMiddleware) AuthPublish(ctx context.Context, topic *string, payload *[]byte) (err error) {
-	return lm.logAction(ctx, "AuthPublish", &[]string{*topic})
+	defer lm.logAction(ctx, "AuthPublish", &[]string{*topic}, time.Now(), err)
+	return lm.svc.AuthPublish(ctx, topic, payload)
 }
 
 // AuthSubscribe implements session.Handler.
 func (lm *loggingMiddleware) AuthSubscribe(ctx context.Context, topics *[]string) (err error) {
-	return lm.logAction(ctx, "AuthSubscribe", topics)
+	defer lm.logAction(ctx, "AuthSubscribe", topics, time.Now(), err)
+	return lm.svc.AuthSubscribe(ctx, topics)
 }
 
 // Connect implements session.Handler.
 func (lm *loggingMiddleware) Connect(ctx context.Context) (err error) {
-	return lm.logAction(ctx, "Connect", nil)
+	defer lm.logAction(ctx, "Connect", nil, time.Now(), err)
+	return lm.svc.Connect(ctx)
 }
 
 // Disconnect implements session.Handler.
 func (lm *loggingMiddleware) Disconnect(ctx context.Context) (err error) {
-	return lm.logAction(ctx, "Disconnect", nil)
+	defer lm.logAction(ctx, "Disconnect", nil, time.Now(), err)
+	return lm.svc.Disconnect(ctx)
 }
 
 // Publish logs the publish request. It logs the time it took to complete the request.
 // If the request fails, it logs the error.
 func (lm *loggingMiddleware) Publish(ctx context.Context, topic *string, payload *[]byte) (err error) {
-	return lm.logAction(ctx, "Publish", &[]string{*topic})
+	defer lm.logAction(ctx, "Publish", &[]string{*topic}, time.Now(), err)
+	return lm.svc.Publish(ctx, topic, payload)
 }
 
 // Subscribe implements session.Handler.
 func (lm *loggingMiddleware) Subscribe(ctx context.Context, topics *[]string) (err error) {
-	return lm.logAction(ctx, "Subscribe", topics)
+	defer lm.logAction(ctx, "Subscribe", topics, time.Now(), err)
+	return lm.svc.Subscribe(ctx, topics)
 }
 
 // Unsubscribe implements session.Handler.
 func (lm *loggingMiddleware) Unsubscribe(ctx context.Context, topics *[]string) (err error) {
-	return lm.logAction(ctx, "Unsubscribe", topics)
+	defer lm.logAction(ctx, "Unsubscribe", topics, time.Now(), err)
+	return lm.svc.Unsubscribe(ctx, topics)
 }
 
 // LoggingMiddleware adds logging facilities to the adapter.
@@ -66,21 +74,17 @@ func LoggingMiddleware(svc session.Handler, logger *slog.Logger) session.Handler
 	return &loggingMiddleware{logger, svc}
 }
 
-func (lm *loggingMiddleware) logAction(ctx context.Context, action string, topics *[]string) (err error) {
-	defer func(begin time.Time) {
-		args := []any{
-			slog.String("duration", time.Since(begin).String()),
-		}
-		if topics != nil {
-			args = append(args, slog.Any("topics", *topics))
-		}
-		if err != nil {
-			args = append(args, slog.Any("error", err))
-			lm.logger.Warn(action+"() failed to complete successfully", args...)
-			return
-		}
-		lm.logger.Info(action+"() completed successfully", args...)
-	}(time.Now())
-
-	return nil
+func (lm *loggingMiddleware) logAction(ctx context.Context, action string, topics *[]string, t time.Time, err error) {
+	args := []any{
+		slog.String("duration", time.Since(t).String()),
+	}
+	if topics != nil {
+		args = append(args, slog.Any("topics", *topics))
+	}
+	if err != nil {
+		args = append(args, slog.Any("error", err))
+		lm.logger.Warn(action+" failed to complete successfully", args...)
+		return
+	}
+	lm.logger.Info(action+" completed successfully", args...)
 }