pkg/ioutils/writeflusher.go

   1 package ioutils
   2
   3 import (
   4         "io"
   5         "sync"
   6 )
   7
   8 // WriteFlusher wraps the Write and Flush operation ensuring that every write
   9 // is a flush. In addition, the Close method can be called to intercept
  10 // Read/Write calls if the targets lifecycle has already ended.
  11 type WriteFlusher struct {
  12         w           io.Writer
  13         flusher     flusher
  14         flushed     chan struct{}
  15         flushedOnce sync.Once
  16         closed      chan struct{}
  17         closeLock   sync.Mutex
  18 }
  19
  20 type flusher interface {
  21         Flush()
  22 }
  23
  24 var errWriteFlusherClosed = io.EOF
  25
  26 func (wf *WriteFlusher) Write(b []byte) (n int, err error) {
  27         select {
  28         case <-wf.closed:
  29                 return 0, errWriteFlusherClosed
  30         default:
  31         }
  32
  33         n, err = wf.w.Write(b)
  34         wf.Flush() // every write is a flush.
  35         return n, err
  36 }
  37
  38 // Flush the stream immediately.
  39 func (wf *WriteFlusher) Flush() {
  40         select {
  41         case <-wf.closed:
  42                 return
  43         default:
  44         }
  45
  46         wf.flushedOnce.Do(func() {
  47                 close(wf.flushed)
  48         })
  49         wf.flusher.Flush()
  50 }
  51
  52 // Flushed returns the state of flushed.
  53 // If it's flushed, return true, or else it return false.
  54 func (wf *WriteFlusher) Flushed() bool {
  55         // BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to
  56         // be used to detect whether or a response code has been issued or not.
  57         // Another hook should be used instead.
  58         var flushed bool
  59         select {
  60         case <-wf.flushed:
  61                 flushed = true
  62         default:
  63         }
  64         return flushed
  65 }
  66
  67 // Close closes the write flusher, disallowing any further writes to the
  68 // target. After the flusher is closed, all calls to write or flush will
  69 // result in an error.
  70 func (wf *WriteFlusher) Close() error {
  71         wf.closeLock.Lock()
  72         defer wf.closeLock.Unlock()
  73
  74         select {
  75         case <-wf.closed:
  76                 return errWriteFlusherClosed
  77         default:
  78                 close(wf.closed)
  79         }
  80         return nil
  81 }
  82
  83 // NewWriteFlusher returns a new WriteFlusher.
  84 func NewWriteFlusher(w io.Writer) *WriteFlusher {
  85         var fl flusher
  86         if f, ok := w.(flusher); ok {
  87                 fl = f
  88         } else {
  89                 fl = &NopFlusher{}
  90         }
  91         return &WriteFlusher{w: w, flusher: fl, closed: make(chan struct{}), flushed: make(chan struct{})}
  92 }