diff --git a/context/http.go b/context/http.go index f61e3bc2..97b1067f 100644 --- a/context/http.go +++ b/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 diff --git a/context/http_test.go b/context/http_test.go index ae88a314..3d4b3c8e 100644 --- a/context/http_test.go +++ b/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 diff --git a/registry/handlers/helpers.go b/registry/handlers/helpers.go index 1f9a8ee1..a4f3abcc 100644 --- a/registry/handlers/helpers.go +++ b/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.