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

@@ -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

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)
 	}

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

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)
 }

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"))