From bd840fe9ab5d7374b09cd6add903128dea5c27ce Mon Sep 17 00:00:00 2001
From: Jeff Mitchell <jeffrey.mitchell@gmail.com>
Date: Thu, 23 Jan 2020 13:57:18 -0500
Subject: [PATCH] Update gated-writer logic a bit (#8227)

This is to smooth some other changes coming once
https://github.com/hashicorp/go-hclog/pull/56 lands
---
 command/agent.go                   |  2 +-
 command/debug.go                   |  2 +-
 command/server.go                  |  2 +-
 helper/gated-writer/writer.go      | 31 +++++++++++++++---------------
 helper/gated-writer/writer_test.go |  2 +-
 5 files changed, 19 insertions(+), 20 deletions(-)

diff --git a/command/agent.go b/command/agent.go
index 3f9110512d..15d0a0517a 100644
--- a/command/agent.go
+++ b/command/agent.go
@@ -172,7 +172,7 @@ func (c *AgentCommand) Run(args []string) int {
 
 	// Create a logger. We wrap it in a gated writer so that it doesn't
 	// start logging too early.
-	c.logGate = &gatedwriter.Writer{Writer: os.Stderr}
+	c.logGate = gatedwriter.NewWriter(os.Stderr)
 	c.logWriter = c.logGate
 	if c.flagCombineLogs {
 		c.logWriter = os.Stdout
diff --git a/command/debug.go b/command/debug.go
index 8cec89f9c5..cbef2d6cca 100644
--- a/command/debug.go
+++ b/command/debug.go
@@ -223,7 +223,7 @@ func (c *DebugCommand) Run(args []string) int {
 	}
 
 	// Initialize the logger for debug output
-	logWriter := &gatedwriter.Writer{Writer: os.Stderr}
+	logWriter := gatedwriter.NewWriter(os.Stderr)
 	if c.logger == nil {
 		c.logger = logging.NewVaultLoggerWithWriter(logWriter, hclog.Trace)
 	}
diff --git a/command/server.go b/command/server.go
index b586e7a86d..17d5728b2a 100644
--- a/command/server.go
+++ b/command/server.go
@@ -710,7 +710,7 @@ func (c *ServerCommand) adjustLogLevel(config *server.Config, logLevelWasNotSet
 func (c *ServerCommand) processLogLevelAndFormat(config *server.Config) (log.Level, string, bool, logging.LogFormat, error) {
 	// Create a logger. We wrap it in a gated writer so that it doesn't
 	// start logging too early.
-	c.logGate = &gatedwriter.Writer{Writer: os.Stderr}
+	c.logGate = gatedwriter.NewWriter(os.Stderr)
 	c.logWriter = c.logGate
 	if c.flagCombineLogs {
 		c.logWriter = os.Stdout
diff --git a/helper/gated-writer/writer.go b/helper/gated-writer/writer.go
index 9c5aeba00f..59c9044405 100644
--- a/helper/gated-writer/writer.go
+++ b/helper/gated-writer/writer.go
@@ -1,6 +1,7 @@
 package gatedwriter
 
 import (
+	"bytes"
 	"io"
 	"sync"
 )
@@ -8,36 +9,34 @@ import (
 // Writer is an io.Writer implementation that buffers all of its
 // data into an internal buffer until it is told to let data through.
 type Writer struct {
-	Writer io.Writer
+	writer io.Writer
 
-	buf   [][]byte
+	buf   bytes.Buffer
 	flush bool
-	lock  sync.RWMutex
+	lock  sync.Mutex
+}
+
+func NewWriter(underlying io.Writer) *Writer {
+	return &Writer{writer: underlying}
 }
 
 // Flush tells the Writer to flush any buffered data and to stop
 // buffering.
 func (w *Writer) Flush() {
 	w.lock.Lock()
-	w.flush = true
-	w.lock.Unlock()
+	defer w.lock.Unlock()
 
-	for _, p := range w.buf {
-		w.Write(p)
-	}
-	w.buf = nil
+	w.flush = true
+	w.buf.WriteTo(w.writer)
 }
 
 func (w *Writer) Write(p []byte) (n int, err error) {
-	w.lock.RLock()
-	defer w.lock.RUnlock()
+	w.lock.Lock()
+	defer w.lock.Unlock()
 
 	if w.flush {
-		return w.Writer.Write(p)
+		return w.writer.Write(p)
 	}
 
-	p2 := make([]byte, len(p))
-	copy(p2, p)
-	w.buf = append(w.buf, p2)
-	return len(p), nil
+	return w.buf.Write(p)
 }
diff --git a/helper/gated-writer/writer_test.go b/helper/gated-writer/writer_test.go
index b007ef1fa8..31659a8add 100644
--- a/helper/gated-writer/writer_test.go
+++ b/helper/gated-writer/writer_test.go
@@ -12,7 +12,7 @@ func TestWriter_impl(t *testing.T) {
 
 func TestWriter(t *testing.T) {
 	buf := new(bytes.Buffer)
-	w := &Writer{Writer: buf}
+	w := NewWriter(buf)
 	w.Write([]byte("foo\n"))
 	w.Write([]byte("bar\n"))