github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/lib/circbufwriter/writer.go (about)

     1  package circbufwriter
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/armon/circbuf"
     9  )
    10  
    11  type circbufWriter struct {
    12  	// circle buffer for data to write
    13  	buf *circbuf.Buffer
    14  
    15  	// error to return from the writer
    16  	err error
    17  
    18  	// bufLock syncronizes access to err and buf
    19  	bufLock sync.Mutex
    20  
    21  	// wrapped writer
    22  	wr io.Writer
    23  
    24  	// signals to flush the buffer
    25  	flushCh chan struct{}
    26  }
    27  
    28  // New created a circle buffered writer that wraps the given writer. The
    29  // bufferSize is the amount of data that will be stored in memory before
    30  // overwriting.
    31  func New(w io.Writer, bufferSize int64) io.WriteCloser {
    32  	buf, _ := circbuf.NewBuffer(bufferSize)
    33  	c := &circbufWriter{
    34  		buf:     buf,
    35  		wr:      w,
    36  		flushCh: make(chan struct{}, 1),
    37  	}
    38  	go c.flushLoop()
    39  	return c
    40  }
    41  
    42  // Write will write the data to the buffer and attempt to flush the buffer to
    43  // the wrapped writer. If the wrapped writer blocks on write, subsequent write
    44  // will be written to the circle buffer.
    45  func (c *circbufWriter) Write(p []byte) (nn int, err error) {
    46  	// If the last write returned an error, return it here. Note there is a
    47  	// small chance of missing an error if multiple writes occure at the same
    48  	// time where the last write nils out the error before it can be returned
    49  	// here.
    50  	c.bufLock.Lock()
    51  	defer c.bufLock.Unlock()
    52  	if c.err != nil {
    53  		return nn, c.err
    54  	}
    55  
    56  	// Write to the buffer
    57  	nn, err = c.buf.Write(p)
    58  
    59  	// Signal to flush the buffer
    60  	select {
    61  	case c.flushCh <- struct{}{}:
    62  	default:
    63  		// flush is blocked
    64  	}
    65  	return nn, err
    66  }
    67  
    68  func (c *circbufWriter) Close() error {
    69  	// Guard against double closing channel
    70  	select {
    71  	case <-c.flushCh:
    72  	default:
    73  		close(c.flushCh)
    74  	}
    75  
    76  	// if the last write errored, it will return here
    77  	c.bufLock.Lock()
    78  	defer c.bufLock.Unlock()
    79  	return c.err
    80  }
    81  
    82  func (c *circbufWriter) flushLoop() {
    83  	// Check buffer every 100ms in case a flush from Write was missed
    84  	ticker := time.NewTicker(time.Millisecond * 100)
    85  	defer ticker.Stop()
    86  	for {
    87  		var err error
    88  		select {
    89  		case _, ok := <-c.flushCh:
    90  			if !ok {
    91  				// Close called, exiting loop
    92  				return
    93  			}
    94  			err = c.flush()
    95  		case <-ticker.C:
    96  			err = c.flush()
    97  		}
    98  
    99  		c.bufLock.Lock()
   100  		c.err = err
   101  		c.bufLock.Unlock()
   102  	}
   103  }
   104  
   105  func (c *circbufWriter) flush() error {
   106  	c.bufLock.Lock()
   107  	b := c.buf.Bytes()
   108  	c.buf.Reset()
   109  	c.bufLock.Unlock()
   110  
   111  	var err error
   112  	if len(b) > 0 {
   113  		_, err = c.wr.Write(b)
   114  	}
   115  	return err
   116  }