@@ -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

+	w           io.Writer

+	flusher     flusher

+	flushed     chan struct{}

+	flushedOnce sync.Once

+	closed      chan struct{}

+	closeLock   sync.Mutex

+}

-	// TODO(stevvooe): Use channel for closed instead, remove mutex. Using a

-	// channel will allow one to properly order the operations.

+type flusher interface {

+	Flush()

 }

-var errWriteFlusherClosed = errors.New("writeflusher: closed")

+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()

-

-	if wf.closed != nil {

-		return wf.closed

+	wf.closeLock.Lock()

+	defer wf.closeLock.Unlock()

+

+	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{})}

 }