Don't panic when a http.ResponseWriter does not implement CloseNotifier

Instead, provide a variant of instrumentedResponseWriter that does not
implement CloseNotifier, and use that when necessary. In
copyFullPayload, log instead of panicing when we encounter something
that doesn't implement CloseNotifier.

This is more complicated than I'd like, but it's necessary because
instrumentedResponseWriter must not embed CloseNotifier unless there's
really a CloseNotifier to embed.

Signed-off-by: Aaron Lehmann <[email protected]>

context/http.go

@@ -103,17 +103,21 @@ func GetRequestID(ctx Context) string {
 // WithResponseWriter returns a new context and response writer that makes
 // interesting response statistics available within the context.
 func WithResponseWriter(ctx Context, w http.ResponseWriter) (Context, http.ResponseWriter) {
-	closeNotifier, ok := w.(http.CloseNotifier)
-	if !ok {
-		panic("the ResponseWriter does not implement CloseNotifier")
-	}
-	irw := &instrumentedResponseWriter{
+	irw := instrumentedResponseWriter{
 		ResponseWriter: w,
-		CloseNotifier:  closeNotifier,
 		Context:        ctx,
 	}
 
-	return irw, irw
+	if closeNotifier, ok := w.(http.CloseNotifier); ok {
+		irwCN := &instrumentedResponseWriterCN{
+			instrumentedResponseWriter: irw,
+			CloseNotifier:              closeNotifier,
+		}
+
+		return irwCN, irwCN
+	}
+
+	return &irw, &irw
 }
 
 // GetResponseWriter returns the http.ResponseWriter from the provided
@@ -263,11 +267,19 @@ func (ctx *muxVarsContext) Value(key interface{}) interface{} {
 	return ctx.Context.Value(key)
 }
 
+// instrumentedResponseWriterCN provides response writer information in a
+// context. It implements http.CloseNotifier so that users can detect
+// early disconnects.
+type instrumentedResponseWriterCN struct {
+	instrumentedResponseWriter
+	http.CloseNotifier
+}
+
 // instrumentedResponseWriter provides response writer information in a
-// context.
+// context. This variant is only used in the case where CloseNotifier is not
+// implemented by the parent ResponseWriter.
 type instrumentedResponseWriter struct {
 	http.ResponseWriter
-	http.CloseNotifier
 	Context
 
 	mu      sync.Mutex

context/http_test.go

@@ -110,13 +110,6 @@ func (trw *testResponseWriter) Header() http.Header {
 	return trw.header
 }
 
-// CloseNotify is only here to make the testResponseWriter implement the
-// http.CloseNotifier interface, which WithResponseWriter expects to be
-// implemented.
-func (trw *testResponseWriter) CloseNotify() <-chan bool {
-	return make(chan bool)
-}
-
 func (trw *testResponseWriter) Write(p []byte) (n int, err error) {
 	if trw.status == 0 {
 		trw.status = http.StatusOK

registry/handlers/helpers.go

@@ -29,7 +29,7 @@ func copyFullPayload(responseWriter http.ResponseWriter, r *http.Request, destWr
 	if notifier, ok := responseWriter.(http.CloseNotifier); ok {
 		clientClosed = notifier.CloseNotify()
 	} else {
-		panic("the ResponseWriter does not implement CloseNotifier")
+		ctxu.GetLogger(context).Warn("the ResponseWriter does not implement CloseNotifier")
 	}
 
 	// Read in the data, if any.