diff --git a/ioutils/writeflusher.go b/ioutils/writeflusher.go index 2b35a26..52a4901 100644 --- a/ioutils/writeflusher.go +++ b/ioutils/writeflusher.go @@ -1,9 +1,7 @@ package ioutils import ( - "errors" "io" - "net/http" "sync" ) @@ -11,45 +9,43 @@ import ( // is a flush. In addition, the Close method can be called to intercept // Read/Write calls if the targets lifecycle has already ended. type WriteFlusher struct { - mu sync.Mutex - w io.Writer - flusher http.Flusher - flushed bool - closed error - - // TODO(stevvooe): Use channel for closed instead, remove mutex. Using a - // channel will allow one to properly order the operations. + w io.Writer + flusher flusher + flushed chan struct{} + flushedOnce sync.Once + closed chan struct{} + closeLock sync.Mutex } -var errWriteFlusherClosed = errors.New("writeflusher: closed") +type flusher interface { + Flush() +} + +var errWriteFlusherClosed = io.EOF func (wf *WriteFlusher) Write(b []byte) (n int, err error) { - wf.mu.Lock() - defer wf.mu.Unlock() - if wf.closed != nil { - return 0, wf.closed + select { + case <-wf.closed: + return 0, errWriteFlusherClosed + default: } n, err = wf.w.Write(b) - wf.flush() // every write is a flush. + wf.Flush() // every write is a flush. return n, err } // Flush the stream immediately. func (wf *WriteFlusher) Flush() { - wf.mu.Lock() - defer wf.mu.Unlock() - - wf.flush() -} - -// flush the stream immediately without taking a lock. Used internally. -func (wf *WriteFlusher) flush() { - if wf.closed != nil { + select { + case <-wf.closed: return + default: } - wf.flushed = true + wf.flushedOnce.Do(func() { + close(wf.flushed) + }) wf.flusher.Flush() } @@ -59,34 +55,38 @@ func (wf *WriteFlusher) Flushed() bool { // BUG(stevvooe): Remove this method. Its use is inherently racy. Seems to // be used to detect whether or a response code has been issued or not. // Another hook should be used instead. - wf.mu.Lock() - defer wf.mu.Unlock() - - return wf.flushed + var flushed bool + select { + case <-wf.flushed: + flushed = true + default: + } + return flushed } // Close closes the write flusher, disallowing any further writes to the // target. After the flusher is closed, all calls to write or flush will // result in an error. func (wf *WriteFlusher) Close() error { - wf.mu.Lock() - defer wf.mu.Unlock() + wf.closeLock.Lock() + defer wf.closeLock.Unlock() - if wf.closed != nil { - return wf.closed + select { + case <-wf.closed: + return errWriteFlusherClosed + default: + close(wf.closed) } - - wf.closed = errWriteFlusherClosed return nil } // NewWriteFlusher returns a new WriteFlusher. func NewWriteFlusher(w io.Writer) *WriteFlusher { - var flusher http.Flusher - if f, ok := w.(http.Flusher); ok { - flusher = f + var fl flusher + if f, ok := w.(flusher); ok { + fl = f } else { - flusher = &NopFlusher{} + fl = &NopFlusher{} } - return &WriteFlusher{w: w, flusher: flusher} + return &WriteFlusher{w: w, flusher: fl, closed: make(chan struct{}), flushed: make(chan struct{})} }