From 994a554459c015ded229d26369ac3433e8a1aa06 Mon Sep 17 00:00:00 2001
From: astraw99 <wangchengiscool@gmail.com>
Date: Fri, 19 Nov 2021 10:32:11 +0800
Subject: [PATCH] fix logger unlock on panic

---
 klog.go      | 23 ++++++++++++++++++-----
 klog_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/klog.go b/klog.go
index 45efbb075..411deea87 100644
--- a/klog.go
+++ b/klog.go
@@ -917,7 +917,14 @@ func LogToStderr(stderr bool) {
 
 // output writes the data to the log files and releases the buffer.
 func (l *loggingT) output(s severity, log *logr.Logger, buf *buffer, depth int, file string, line int, alsoToStderr bool) {
+	var isLocked = true
 	l.mu.Lock()
+	defer func() {
+		if isLocked {
+			l.mu.Unlock() // Make sure to unlock if panic occurs below
+		}
+	}()
+
 	if l.traceLocation.isSet() {
 		if l.traceLocation.match(file, line) {
 			buf.Write(stacks(false))
@@ -928,7 +935,7 @@ func (l *loggingT) output(s severity, log *logr.Logger, buf *buffer, depth int,
 		// TODO: set 'severity' and caller information as structured log info
 		// keysAndValues := []interface{}{"severity", severityName[s], "file", file, "line", line}
 		if s == errorLog {
-			l.logr.WithCallDepth(depth+3).Error(nil, string(data))
+			l.logr.WithCallDepth(depth+3).Error(nil, string(data)) // If panic here will be captured by defer.
 		} else {
 			log.WithCallDepth(depth + 3).Info(string(data))
 		}
@@ -944,16 +951,20 @@ func (l *loggingT) output(s severity, log *logr.Logger, buf *buffer, depth int,
 			// will point to the same file, so just use one of them to write data.
 			if l.file[infoLog] == nil {
 				if err := l.createFiles(infoLog); err != nil {
-					os.Stderr.Write(data) // Make sure the message appears somewhere.
+					os.Stderr.Write(data)        // Make sure the message appears somewhere.
+					logExitFunc = func(error) {} // Set `logExitFunc` non-nil to unlock safely.
 					l.exit(err)
+					return
 				}
 			}
 			l.file[infoLog].Write(data)
 		} else {
 			if l.file[s] == nil {
 				if err := l.createFiles(s); err != nil {
-					os.Stderr.Write(data) // Make sure the message appears somewhere.
+					os.Stderr.Write(data)        // Make sure the message appears somewhere.
+					logExitFunc = func(error) {} // Set `logExitFunc` non-nil to unlock safely.
 					l.exit(err)
+					return
 				}
 			}
 
@@ -980,6 +991,7 @@ func (l *loggingT) output(s severity, log *logr.Logger, buf *buffer, depth int,
 		// If we got here via Exit rather than Fatal, print no stacks.
 		if atomic.LoadUint32(&fatalNoStacks) > 0 {
 			l.mu.Unlock()
+			isLocked = false
 			timeoutFlush(10 * time.Second)
 			os.Exit(1)
 		}
@@ -997,11 +1009,12 @@ func (l *loggingT) output(s severity, log *logr.Logger, buf *buffer, depth int,
 			}
 		}
 		l.mu.Unlock()
+		isLocked = false
 		timeoutFlush(10 * time.Second)
-		os.Exit(255) // C++ uses -1, which is silly because it's anded with 255 anyway.
+		os.Exit(255) // C++ uses -1, which is silly because it's anded(&) with 255 anyway.
 	}
 	l.putBuffer(buf)
-	l.mu.Unlock()
+
 	if stats := severityStats[s]; stats != nil {
 		atomic.AddInt64(&stats.lines, 1)
 		atomic.AddInt64(&stats.bytes, int64(len(data)))
diff --git a/klog_test.go b/klog_test.go
index 9a44c6be3..ad70040c6 100644
--- a/klog_test.go
+++ b/klog_test.go
@@ -1844,3 +1844,35 @@ func TestKObjs(t *testing.T) {
 		})
 	}
 }
+
+// Benchmark test for lock with/without defer
+type structWithLock struct {
+	m sync.Mutex
+	n int64
+}
+
+func BenchmarkWithoutDeferUnLock(b *testing.B) {
+	s := structWithLock{}
+	for i := 0; i < b.N; i++ {
+		s.addWithoutDefer()
+	}
+}
+
+func BenchmarkWithDeferUnLock(b *testing.B) {
+	s := structWithLock{}
+	for i := 0; i < b.N; i++ {
+		s.addWithDefer()
+	}
+}
+
+func (s *structWithLock) addWithoutDefer() {
+	s.m.Lock()
+	s.n++
+	s.m.Unlock()
+}
+
+func (s *structWithLock) addWithDefer() {
+	s.m.Lock()
+	defer s.m.Unlock()
+	s.n++
+}