From 9a44ef1c9c6bd8ff21276b4c86a41153f12344d3 Mon Sep 17 00:00:00 2001
From: Leonard Wen <leonard.wen@thecarousell.com>
Date: Fri, 12 Aug 2022 18:03:04 +0800
Subject: [PATCH] add RWlock for logger

---
 utils/log/loggers/fields.go                | 42 ++++++++++++++++++----
 utils/log/loggers/test/benchmark_test.go   |  3 ++
 utils/log/loggers/test/parallelism_test.go |  6 ++--
 3 files changed, 41 insertions(+), 10 deletions(-)

diff --git a/utils/log/loggers/fields.go b/utils/log/loggers/fields.go
index 41105977..8855219a 100644
--- a/utils/log/loggers/fields.go
+++ b/utils/log/loggers/fields.go
@@ -2,6 +2,7 @@ package loggers
 
 import (
 	"context"
+	"sync"
 )
 
 type logsContext string
@@ -12,6 +13,10 @@ var (
 
 //LogFields contains all fields that have to be added to logs
 type LogFields map[string]interface{}
+type ProtectedLogFields struct {
+	Content LogFields
+	mtx     sync.RWMutex
+}
 
 // Add or modify log fields
 func (o LogFields) Add(key string, value interface{}) {
@@ -28,14 +33,19 @@ func (o LogFields) Del(key string) {
 //AddToLogContext adds log fields to context.
 // Any info added here will be added to all logs using this context
 func AddToLogContext(ctx context.Context, key string, value interface{}) context.Context {
-	data := FromContext(ctx)
+	data := fromContext(ctx)
+	//Initialize if key doesn't exist
 	if data == nil {
-		ctx = context.WithValue(ctx, contextKey, make(LogFields))
-		data = FromContext(ctx)
+		ctx = context.WithValue(ctx, contextKey, &ProtectedLogFields{Content: make(LogFields)})
+		data = fromContext(ctx)
 	}
 	m := ctx.Value(contextKey)
-	if data, ok := m.(LogFields); ok {
-		data.Add(key, value)
+	if data, ok := m.(*ProtectedLogFields); ok {
+		data.mtx.Lock()
+		defer data.mtx.Unlock()
+		// d := data.Content
+		// fmt.Printf("Address %p\n", d)
+		data.Content.Add(key, value)
 	}
 	return ctx
 }
@@ -46,8 +56,26 @@ func FromContext(ctx context.Context) LogFields {
 		return nil
 	}
 	if h := ctx.Value(contextKey); h != nil {
-		if logData, ok := h.(LogFields); ok {
-			return logData
+		if plf, ok := h.(*ProtectedLogFields); ok {
+			plf.mtx.RLock()
+			defer plf.mtx.RUnlock()
+			content := make(LogFields)
+			for k, v := range plf.Content {
+				content[k] = v
+			}
+			return content
+		}
+	}
+	return nil
+}
+
+func fromContext(ctx context.Context) *ProtectedLogFields {
+	if ctx == nil {
+		return nil
+	}
+	if h := ctx.Value(contextKey); h != nil {
+		if plf, ok := h.(*ProtectedLogFields); ok {
+			return plf
 		}
 	}
 	return nil
diff --git a/utils/log/loggers/test/benchmark_test.go b/utils/log/loggers/test/benchmark_test.go
index fa3b5750..935ce3dc 100644
--- a/utils/log/loggers/test/benchmark_test.go
+++ b/utils/log/loggers/test/benchmark_test.go
@@ -11,6 +11,9 @@ import (
 
 func BenchmarkFromContext(b *testing.B) {
 	ctx := context.Background()
+	for i := 0; i < 10000; i++ {
+		s.AddToLogContext(ctx, fmt.Sprintf("key%d", i), "good value")
+	}
 	for i := 0; i < b.N; i++ {
 		s.FromContext(ctx)
 	}
diff --git a/utils/log/loggers/test/parallelism_test.go b/utils/log/loggers/test/parallelism_test.go
index d89ebc3d..2c586367 100644
--- a/utils/log/loggers/test/parallelism_test.go
+++ b/utils/log/loggers/test/parallelism_test.go
@@ -13,8 +13,8 @@ import (
 	"github.com/stretchr/testify/assert"
 )
 
-const readWorkerCount = 5
-const writeWorkerCount = 5
+const readWorkerCount = 50
+const writeWorkerCount = 50
 
 func readWorker(idx int, ctx context.Context) {
 	lf := s.FromContext(ctx)
@@ -64,7 +64,7 @@ func TestParallelWrite(t *testing.T) {
 	wg.Wait()
 
 	lf := s.FromContext(ctx)
-	fmt.Println(lf)
+	fmt.Println("lf", lf)
 
 	assert.Contains(t, lf, "test-key")
 	for i := 1; i <= writeWorkerCount; i++ {