Merge pull request #17634 from stevvooe/avoid-panic-on-flush

Avoid panic on write after close in http

api/server/router/local/container.go

@@ -63,7 +63,9 @@ func (s *router) getContainersStats(ctx context.Context, w http.ResponseWriter,
 		w.Header().Set("Content-Type", "application/json")
 		out = w
 	} else {
-		out = ioutils.NewWriteFlusher(w)
+		wf := ioutils.NewWriteFlusher(w)
+		out = wf
+		defer wf.Close()
 	}
 
 	var closeNotifier <-chan bool
@@ -116,11 +118,16 @@ func (s *router) getContainersLogs(ctx context.Context, w http.ResponseWriter, r
 		return derr.ErrorCodeNoSuchContainer.WithArgs(containerName)
 	}
 
-	outStream := ioutils.NewWriteFlusher(w)
 	// write an empty chunk of data (this is to ensure that the
 	// HTTP Response is sent immediately, even if the container has
 	// not yet produced any data)
-	outStream.Write(nil)
+	w.WriteHeader(http.StatusOK)
+	if flusher, ok := w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+
+	output := ioutils.NewWriteFlusher(w)
+	defer output.Close()
 
 	logsConfig := &daemon.ContainerLogsConfig{
 		Follow:     httputils.BoolValue(r, "follow"),
@@ -129,7 +136,7 @@ func (s *router) getContainersLogs(ctx context.Context, w http.ResponseWriter, r
 		Tail:       r.Form.Get("tail"),
 		UseStdout:  stdout,
 		UseStderr:  stderr,
-		OutStream:  outStream,
+		OutStream:  output,
 		Stop:       closeNotifier,
 	}
 

api/server/router/local/image.go

@@ -105,6 +105,7 @@ func (s *router) postImagesCreate(ctx context.Context, w http.ResponseWriter, r
 		err    error
 		output = ioutils.NewWriteFlusher(w)
 	)
+	defer output.Close()
 
 	w.Header().Set("Content-Type", "application/json")
 
@@ -184,6 +185,7 @@ func (s *router) postImagesPush(ctx context.Context, w http.ResponseWriter, r *h
 
 	name := vars["name"]
 	output := ioutils.NewWriteFlusher(w)
+	defer output.Close()
 	imagePushConfig := &graph.ImagePushConfig{
 		MetaHeaders: metaHeaders,
 		AuthConfig:  authConfig,
@@ -211,6 +213,7 @@ func (s *router) getImagesGet(ctx context.Context, w http.ResponseWriter, r *htt
 	w.Header().Set("Content-Type", "application/x-tar")
 
 	output := ioutils.NewWriteFlusher(w)
+	defer output.Close()
 	var names []string
 	if name, ok := vars["name"]; ok {
 		names = []string{name}
@@ -283,6 +286,7 @@ func (s *router) postBuild(ctx context.Context, w http.ResponseWriter, r *http.R
 
 	version := httputils.VersionFromContext(ctx)
 	output := ioutils.NewWriteFlusher(w)
+	defer output.Close()
 	sf := streamformatter.NewJSONStreamFormatter()
 	errf := func(err error) error {
 		// Do not write the error in the http output if it's still empty.

api/server/router/local/info.go

@@ -52,16 +52,6 @@ func (s *router) getInfo(ctx context.Context, w http.ResponseWriter, r *http.Req
 	return httputils.WriteJSON(w, http.StatusOK, info)
 }
 
-func buildOutputEncoder(w http.ResponseWriter) *json.Encoder {
-	w.Header().Set("Content-Type", "application/json")
-	outStream := ioutils.NewWriteFlusher(w)
-	// Write an empty chunk of data.
-	// This is to ensure that the HTTP status code is sent immediately,
-	// so that it will not block the receiver.
-	outStream.Write(nil)
-	return json.NewEncoder(outStream)
-}
-
 func (s *router) getEvents(ctx context.Context, w http.ResponseWriter, r *http.Request, vars map[string]string) error {
 	if err := httputils.ParseForm(r); err != nil {
 		return err
@@ -87,7 +77,19 @@ func (s *router) getEvents(ctx context.Context, w http.ResponseWriter, r *http.R
 		return err
 	}
 
-	enc := buildOutputEncoder(w)
+	w.Header().Set("Content-Type", "application/json")
+
+	// This is to ensure that the HTTP status code is sent immediately,
+	// so that it will not block the receiver.
+	w.WriteHeader(http.StatusOK)
+	if flusher, ok := w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+
+	output := ioutils.NewWriteFlusher(w)
+	defer output.Close()
+
+	enc := json.NewEncoder(output)
 
 	current, l, cancel := s.daemon.SubscribeToEvents()
 	defer cancel()

pkg/ioutils/writeflusher.go

@@ -1,32 +1,54 @@
 package ioutils
 
 import (
+	"errors"
 	"io"
 	"net/http"
 	"sync"
 )
 
-// WriteFlusher wraps the Write and Flush operation.
+// WriteFlusher wraps the Write and Flush operation ensuring that every write
+// 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 {
-	sync.Mutex
+	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.
 }
 
+var errWriteFlusherClosed = errors.New("writeflusher: closed")
+
 func (wf *WriteFlusher) Write(b []byte) (n int, err error) {
-	wf.Lock()
-	defer wf.Unlock()
+	wf.mu.Lock()
+	defer wf.mu.Unlock()
+	if wf.closed != nil {
+		return 0, wf.closed
+	}
+
 	n, err = wf.w.Write(b)
-	wf.flushed = true
-	wf.flusher.Flush()
+	wf.flush() // every write is a flush.
 	return n, err
 }
 
 // Flush the stream immediately.
 func (wf *WriteFlusher) Flush() {
-	wf.Lock()
-	defer wf.Unlock()
+	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 {
+		return
+	}
+
 	wf.flushed = true
 	wf.flusher.Flush()
 }
@@ -34,11 +56,30 @@ func (wf *WriteFlusher) Flush() {
 // Flushed returns the state of flushed.
 // If it's flushed, return true, or else it return false.
 func (wf *WriteFlusher) Flushed() bool {
-	wf.Lock()
-	defer wf.Unlock()
+	// 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
 }
 
+// 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.closed = errWriteFlusherClosed
+	return nil
+}
+
 // NewWriteFlusher returns a new WriteFlusher.
 func NewWriteFlusher(w io.Writer) *WriteFlusher {
 	var flusher http.Flusher