implement more fine-grained setting of "Connection: close" header

http/body.go

@@ -0,0 +1,19 @@
+package http
+
+import "io"
+
+// bodyWrapper wraps an io.Reader and calls onError whenever the Read function returns an error.
+// This was designed for wrapping the request body, so we can know whether it was closed.
+type bodyWrapper struct {
+	io.ReadCloser
+	onError func(err error)
+}
+
+func (bw bodyWrapper) Read(data []byte) (int, error) {
+	n, err := bw.ReadCloser.Read(data)
+	if err != nil {
+		bw.onError(err)
+	}
+
+	return n, err
+}

http/handler.go

@@ -109,6 +109,23 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// If we have a request body, make sure the preamble
+	// knows that it should close the body if it wants to
+	// write before completing reading.
+	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
+	// FIXME: https://github.com/golang/go/issues/15527
+	var bodyErrChan chan error
+	if r.Body != http.NoBody {
+		bodyErrChan = make(chan error)
+		bw := bodyWrapper{
+			ReadCloser: r.Body,
+			onError: func(err error) {
+				bodyErrChan <- err
+			},
+		}
+		r.Body = bw
+	}
+
 	req, err := parseRequest(ctx, r, h.root)
 	if err != nil {
 		if err == ErrNotFound {
@@ -157,14 +174,7 @@ func (h *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
-	// If we have a request body, close the connection when we're done.
-	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
-	// FIXME: https://github.com/golang/go/issues/15527
-	if r.Body != http.NoBody {
-		w.Header().Set("Connection", "close")
-	}
-
-	re := NewResponseEmitter(w, r.Method, req)
+	re := NewResponseEmitter(w, r.Method, req, WithRequestBodyErrorChan(bodyErrChan))
 	h.root.Call(req, re, h.env)
 }
 

http/responseemitter.go

@@ -28,11 +28,10 @@ var (
 )
 
 // NewResponeEmitter returns a new ResponseEmitter.
-func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request) ResponseEmitter {
+func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request, opts ...ResponseEmitterOption) ResponseEmitter {
 	encType := cmds.GetEncoding(req)
 
 	var enc cmds.Encoder
-
 	if _, ok := cmds.Encoders[encType]; ok {
 		enc = cmds.Encoders[encType](req)(w)
 	}
@@ -44,9 +43,24 @@ func NewResponseEmitter(w http.ResponseWriter, method string, req *cmds.Request)
 		method:  method,
 		req:     req,
 	}
+
+	for _, opt := range opts {
+		opt(re)
+	}
+
 	return re
 }
 
+type ResponseEmitterOption func(*responseEmitter)
+
+func WithRequestBodyErrorChan(ch <-chan error) ResponseEmitterOption {
+	return func(re *responseEmitter) {
+		if ch != nil {
+			re.bodyErrChan = ch
+		}
+	}
+}
+
 type ResponseEmitter interface {
 	cmds.ResponseEmitter
 	http.Flusher
@@ -63,6 +77,8 @@ type responseEmitter struct {
 	length uint64
 	err    *cmdkit.Error
 
+	bodyErrChan <-chan error
+
 	streaming bool
 	closed    bool
 	once      sync.Once
@@ -208,6 +224,22 @@ func (re *responseEmitter) doPreamble(value interface{}) {
 		mime   string
 	)
 
+	// If we have a request body, make sure we close the body
+	// if we want to write before completing reading.
+	// FIXME: https://github.com/ipfs/go-ipfs/issues/5168
+	// FIXME: https://github.com/golang/go/issues/15527
+	if re.bodyErrChan != nil {
+		select {
+		case <-re.bodyErrChan:
+			// all good, we received an error, so the body is read completely.
+			// we handle the error where it occurs, here we just want to know that we're done
+		default:
+			// set connection close header, because we want to write
+			// even though the body is not yet read completely.
+			h.Set("Connection", "close")
+		}
+	}
+
 	switch v := value.(type) {
 	case io.Reader:
 		// set streams output type to text to avoid issues with browsers rendering