github.com/nsqio/nsq@v1.3.0/internal/writers/spread_writer.go (about)

     1  package writers
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  )
     7  
     8  type SpreadWriter struct {
     9  	w        io.Writer
    10  	interval time.Duration
    11  	buf      [][]byte
    12  	exitCh   chan int
    13  }
    14  
    15  func NewSpreadWriter(w io.Writer, interval time.Duration, exitCh chan int) *SpreadWriter {
    16  	return &SpreadWriter{
    17  		w:        w,
    18  		interval: interval,
    19  		buf:      make([][]byte, 0),
    20  		exitCh:   exitCh,
    21  	}
    22  }
    23  
    24  func (s *SpreadWriter) Write(p []byte) (int, error) {
    25  	b := make([]byte, len(p))
    26  	copy(b, p)
    27  	s.buf = append(s.buf, b)
    28  	return len(p), nil
    29  }
    30  
    31  func (s *SpreadWriter) Flush() {
    32  	if len(s.buf) == 0 {
    33  		// nothing to write, just wait
    34  		select {
    35  		case <-time.After(s.interval):
    36  		case <-s.exitCh:
    37  		}
    38  		return
    39  	}
    40  	sleep := s.interval / time.Duration(len(s.buf))
    41  	ticker := time.NewTicker(sleep)
    42  	for _, b := range s.buf {
    43  		s.w.Write(b)
    44  		select {
    45  		case <-ticker.C:
    46  		case <-s.exitCh: // skip sleeps finish writes
    47  		}
    48  	}
    49  	ticker.Stop()
    50  	s.buf = s.buf[:0]
    51  }