github.com/lmb/consul@v1.4.1/logger/gated_writer.go (about)

     1  package logger
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  // GatedWriter is an io.Writer implementation that buffers all of its
     9  // data into an internal buffer until it is told to let data through.
    10  type GatedWriter struct {
    11  	Writer io.Writer
    12  
    13  	buf   [][]byte
    14  	flush bool
    15  	lock  sync.RWMutex
    16  }
    17  
    18  // Flush tells the GatedWriter to flush any buffered data and to stop
    19  // buffering.
    20  func (w *GatedWriter) Flush() {
    21  	w.lock.Lock()
    22  	w.flush = true
    23  	w.lock.Unlock()
    24  
    25  	for _, p := range w.buf {
    26  		w.Write(p)
    27  	}
    28  	w.buf = nil
    29  }
    30  
    31  func (w *GatedWriter) Write(p []byte) (n int, err error) {
    32  	// Once we flush we no longer synchronize writers since there's
    33  	// no use of the internal buffer. This is the happy path.
    34  	w.lock.RLock()
    35  	if w.flush {
    36  		w.lock.RUnlock()
    37  		return w.Writer.Write(p)
    38  	}
    39  	w.lock.RUnlock()
    40  
    41  	// Now take the write lock.
    42  	w.lock.Lock()
    43  	defer w.lock.Unlock()
    44  
    45  	// Things could have changed between the locking operations, so we
    46  	// have to check one more time.
    47  	if w.flush {
    48  		return w.Writer.Write(p)
    49  	}
    50  
    51  	// Buffer up the written data.
    52  	p2 := make([]byte, len(p))
    53  	copy(p2, p)
    54  	w.buf = append(w.buf, p2)
    55  	return len(p), nil
    56  }