Merge pull request #74 from vincentfree/middleware-add-response-data

add a custom ResponseWriter to capture the statuscode of a response in otlpmiddleware

otelmiddleware/responseWriter.go

@@ -0,0 +1,239 @@
+package otelmiddleware
+
+// The original work was derived from Goji's middleware, source:
+// https://github.com/zenazn/goji/tree/master/web/middleware
+// In this library the derivitive comes from the Chi routers middleware:
+// https://github.com/go-chi/chi/blob/master/middleware/wrap_writer.go
+
+import (
+	"bufio"
+	"io"
+	"net"
+	"net/http"
+)
+
+// NewWrapResponseWriter wraps an http.ResponseWriter, returning a proxy that allows you to
+// hook into various parts of the response process.
+func NewWrapResponseWriter(w http.ResponseWriter, protoMajor int) WrapResponseWriter {
+	_, fl := w.(http.Flusher)
+
+	bw := basicWriter{ResponseWriter: w}
+
+	if protoMajor == 2 {
+		_, ps := w.(http.Pusher)
+		if fl && ps {
+			return &http2FancyWriter{bw}
+		}
+	} else {
+		_, hj := w.(http.Hijacker)
+		_, rf := w.(io.ReaderFrom)
+		if fl && hj && rf {
+			return &httpFancyWriter{bw}
+		}
+		if fl && hj {
+			return &flushHijackWriter{bw}
+		}
+		if hj {
+			return &hijackWriter{bw}
+		}
+	}
+
+	if fl {
+		return &flushWriter{bw}
+	}
+
+	return &bw
+}
+
+// WrapResponseWriter is a proxy around an http.ResponseWriter that allows you to hook
+// into various parts of the response process.
+type WrapResponseWriter interface {
+	http.ResponseWriter
+	// Status returns the HTTP status of the request, or 0 if one has not
+	// yet been sent.
+	Status() int
+	// BytesWritten returns the total number of bytes sent to the client.
+	BytesWritten() int
+	// Tee causes the response body to be written to the given io.Writer in
+	// addition to proxying the writes through. Only one io.Writer can be
+	// tee'd to at once: setting a second one will overwrite the first.
+	// Writes will be sent to the proxy before being written to this
+	// io.Writer. It is illegal for the tee'd writer to be modified
+	// concurrently with writes.
+	Tee(io.Writer)
+	// Unwrap returns the original proxied target.
+	Unwrap() http.ResponseWriter
+	// Discard causes all writes to the original ResponseWriter be discarded,
+	// instead writing only to the tee'd writer if it's set.
+	// The caller is responsible for calling WriteHeader and Write on the
+	// original ResponseWriter once the processing is done.
+	Discard()
+}
+
+// basicWriter wraps a http.ResponseWriter that implements the minimal
+// http.ResponseWriter interface.
+type basicWriter struct {
+	http.ResponseWriter
+	wroteHeader bool
+	code        int
+	bytes       int
+	tee         io.Writer
+	discard     bool
+}
+
+func (b *basicWriter) WriteHeader(code int) {
+	if !b.wroteHeader {
+		b.code = code
+		b.wroteHeader = true
+		if !b.discard {
+			b.ResponseWriter.WriteHeader(code)
+		}
+	}
+}
+
+func (b *basicWriter) Write(buf []byte) (n int, err error) {
+	b.maybeWriteHeader()
+	if !b.discard {
+		n, err = b.ResponseWriter.Write(buf)
+		if b.tee != nil {
+			_, err2 := b.tee.Write(buf[:n])
+			// Prefer errors generated by the proxied writer.
+			if err == nil {
+				err = err2
+			}
+		}
+	} else if b.tee != nil {
+		n, err = b.tee.Write(buf)
+	} else {
+		n, err = io.Discard.Write(buf)
+	}
+	b.bytes += n
+	return n, err
+}
+
+func (b *basicWriter) maybeWriteHeader() {
+	if !b.wroteHeader {
+		b.WriteHeader(http.StatusOK)
+	}
+}
+
+func (b *basicWriter) Status() int {
+	return b.code
+}
+
+func (b *basicWriter) BytesWritten() int {
+	return b.bytes
+}
+
+func (b *basicWriter) Tee(w io.Writer) {
+	b.tee = w
+}
+
+func (b *basicWriter) Unwrap() http.ResponseWriter {
+	return b.ResponseWriter
+}
+
+func (b *basicWriter) Discard() {
+	b.discard = true
+}
+
+// flushWriter ...
+type flushWriter struct {
+	basicWriter
+}
+
+func (f *flushWriter) Flush() {
+	f.wroteHeader = true
+	fl := f.basicWriter.ResponseWriter.(http.Flusher)
+	fl.Flush()
+}
+
+var _ http.Flusher = &flushWriter{}
+
+// hijackWriter ...
+type hijackWriter struct {
+	basicWriter
+}
+
+func (f *hijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
+	return hj.Hijack()
+}
+
+var _ http.Hijacker = &hijackWriter{}
+
+// flushHijackWriter ...
+type flushHijackWriter struct {
+	basicWriter
+}
+
+func (f *flushHijackWriter) Flush() {
+	f.wroteHeader = true
+	fl := f.basicWriter.ResponseWriter.(http.Flusher)
+	fl.Flush()
+}
+
+func (f *flushHijackWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
+	return hj.Hijack()
+}
+
+var _ http.Flusher = &flushHijackWriter{}
+var _ http.Hijacker = &flushHijackWriter{}
+
+// httpFancyWriter is a HTTP writer that additionally satisfies
+// http.Flusher, http.Hijacker, and io.ReaderFrom. It exists for the common case
+// of wrapping the http.ResponseWriter that package http gives you, in order to
+// make the proxied object support the full method set of the proxied object.
+type httpFancyWriter struct {
+	basicWriter
+}
+
+func (f *httpFancyWriter) Flush() {
+	f.wroteHeader = true
+	fl := f.basicWriter.ResponseWriter.(http.Flusher)
+	fl.Flush()
+}
+
+func (f *httpFancyWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
+	hj := f.basicWriter.ResponseWriter.(http.Hijacker)
+	return hj.Hijack()
+}
+
+func (f *http2FancyWriter) Push(target string, opts *http.PushOptions) error {
+	return f.basicWriter.ResponseWriter.(http.Pusher).Push(target, opts)
+}
+
+func (f *httpFancyWriter) ReadFrom(r io.Reader) (int64, error) {
+	if f.basicWriter.tee != nil {
+		n, err := io.Copy(&f.basicWriter, r)
+		f.basicWriter.bytes += int(n)
+		return n, err
+	}
+	rf := f.basicWriter.ResponseWriter.(io.ReaderFrom)
+	f.basicWriter.maybeWriteHeader()
+	n, err := rf.ReadFrom(r)
+	f.basicWriter.bytes += int(n)
+	return n, err
+}
+
+var _ http.Flusher = &httpFancyWriter{}
+var _ http.Hijacker = &httpFancyWriter{}
+var _ http.Pusher = &http2FancyWriter{}
+var _ io.ReaderFrom = &httpFancyWriter{}
+
+// http2FancyWriter is a HTTP2 writer that additionally satisfies
+// http.Flusher, and io.ReaderFrom. It exists for the common case
+// of wrapping the http.ResponseWriter that package http gives you, in order to
+// make the proxied object support the full method set of the proxied object.
+type http2FancyWriter struct {
+	basicWriter
+}
+
+func (f *http2FancyWriter) Flush() {
+	f.wroteHeader = true
+	fl := f.basicWriter.ResponseWriter.(http.Flusher)
+	fl.Flush()
+}
+
+var _ http.Flusher = &http2FancyWriter{}

otelmiddleware/tracing.go

@@ -15,6 +15,7 @@
 package otelmiddleware
 
 import (
+	"go.opentelemetry.io/otel/codes"
 	"net/http"
 
 	"go.opentelemetry.io/otel"
@@ -25,7 +26,7 @@ import (
 )
 
 // version is used as the instrumentation version.
-const version = "0.0.11"
+const version = "0.1.0"
 
 // TraceOption takes a traceConfig struct and applies changes.
 // It can be passed to the TraceWithOptions function to configure a traceConfig struct.
@@ -50,7 +51,7 @@ func TraceWithOptions(opt ...TraceOption) func(next http.Handler) http.Handler {
 	}
 	// check for the traceConfig.tracer if absent use a default value.
 	if config.tracer == nil {
-		config.tracer = otel.Tracer("otel-tracer", trace.WithInstrumentationVersion(version))
+		config.tracer = otel.Tracer("github.com/vincentfree/opentelemetry/otelmiddleware", trace.WithInstrumentationVersion(version))
 	}
 	// check for the traceConfig.propagator if absent use a default value.
 	if config.propagator == nil {
@@ -86,7 +87,7 @@ func TraceWithOptions(opt ...TraceOption) func(next http.Handler) http.Handler {
 			spanName := extractRoute(r.RequestURI)
 			if spanName == "" {
 				// no path available
-				spanName = "HTTP " + r.Method + " /"
+				spanName = r.Proto + " " + r.Method + " /"
 			}
 
 			// create a good name to recognize where the span originated.
@@ -102,8 +103,20 @@ func TraceWithOptions(opt ...TraceOption) func(next http.Handler) http.Handler {
 			carrier := propagation.HeaderCarrier(r.Header)
 			otel.GetTextMapPropagator().Inject(ctx, carrier)
 
+			// use a wrapper for the http.responseWriter to capture the response status code;
+			// this information is added to the spans generated by the middleware
+			wrapperRes := NewWrapResponseWriter(w, r.ProtoMajor)
+
 			// serve the request to the next middleware.
-			next.ServeHTTP(w, r)
+			next.ServeHTTP(wrapperRes, r)
+			// add the response status code to the span
+			if span.IsRecording() {
+				statusCode := wrapperRes.Status()
+				span.SetAttributes(semconv.HTTPAttributesFromHTTPStatusCode(statusCode)...)
+				if statusCode >= 500 {
+					span.SetStatus(codes.Error, http.StatusText(statusCode))
+				}
+			}
 		}
 
 		return http.HandlerFunc(fn)