From 2840389b1a31ce7fa260704f6fbcad10a8e4aa2e Mon Sep 17 00:00:00 2001
From: Vincent Tsao <vtsao@netflix.com>
Date: Wed, 26 Jan 2022 12:37:10 -0500
Subject: [PATCH]

---
 logger.go      |  7 +++++++
 logger_test.go | 23 ++++++++++++++++++++++-
 2 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/logger.go b/logger.go
index 5ff0aef..4e94af7 100644
--- a/logger.go
+++ b/logger.go
@@ -387,6 +387,13 @@ func (logger *Logger) SetFormatter(formatter Formatter) {
 	logger.Formatter = formatter
 }
 
+// Output returns the IO writer for this logger in a concurrency safe manner.
+func (logger *Logger) Output() io.Writer {
+	logger.mu.Lock()
+	defer logger.mu.Unlock()
+	return logger.Out
+}
+
 // SetOutput sets the logger output.
 func (logger *Logger) SetOutput(output io.Writer) {
 	logger.mu.Lock()
diff --git a/logger_test.go b/logger_test.go
index 595c682..a6e28c6 100644
--- a/logger_test.go
+++ b/logger_test.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"fmt"
+	"sync"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -70,7 +71,7 @@ func TestWarninglnNotEqualToWarning(t *testing.T) {
 
 type testBufferPool struct {
 	buffers []*bytes.Buffer
-	get int
+	get     int
 }
 
 func (p *testBufferPool) Get() *bytes.Buffer {
@@ -95,3 +96,23 @@ func TestLogger_SetBufferPool(t *testing.T) {
 	assert.Equal(t, pool.get, 1, "Logger.SetBufferPool(): The BufferPool.Get() must be called")
 	assert.Len(t, pool.buffers, 1, "Logger.SetBufferPool(): The BufferPool.Put() must be called")
 }
+
+func TestOutputRace(t *testing.T) {
+	l := New()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		_ = l.Output()
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		out := &bytes.Buffer{}
+		l.SetOutput(out)
+	}()
+
+	wg.Wait()
+}