Fix risk of unfinished span when panic occurred in http handler (#38)

nethttp/server.go

@@ -126,16 +126,18 @@ func MiddlewareFunc(tr opentracing.Tracer, h http.HandlerFunc, options ...MWOpti
 		}
 		ext.Component.Set(sp, componentName)
 
-		sct := &statusCodeTracker{w, 200}
+		sct := &statusCodeTracker{ResponseWriter: w}
 		r = r.WithContext(opentracing.ContextWithSpan(r.Context(), sp))
 
-		h(sct.wrappedResponseWriter(), r)
+		defer func() {
+			ext.HTTPStatusCode.Set(sp, uint16(sct.status))
+			if sct.status >= http.StatusInternalServerError || !sct.wroteheader {
+				ext.Error.Set(sp, true)
+			}
+			sp.Finish()
+		}()
 
-		ext.HTTPStatusCode.Set(sp, uint16(sct.status))
-		if sct.status >= http.StatusInternalServerError {
-			ext.Error.Set(sp, true)
-		}
-		sp.Finish()
+		h(sct.wrappedResponseWriter(), r)
 	}
 	return http.HandlerFunc(fn)
 }

nethttp/server_test.go

@@ -218,8 +218,8 @@ func TestSpanError(t *testing.T) {
 	wantTags := map[string]interface{}{string(ext.Error): true}
 
 	tests := []struct {
-		url     string
-		Tags    map[string]interface{}
+		url  string
+		Tags map[string]interface{}
 	}{
 		{"/root", make(map[string]interface{})},
 		{"/error", wantTags},
@@ -273,3 +273,70 @@ func BenchmarkStatusCodeTrackingOverhead(b *testing.B) {
 		}
 	})
 }
+
+func TestMiddlewareHandlerPanic(t *testing.T) {
+	tests := []struct {
+		handler func(w http.ResponseWriter, r *http.Request)
+		status  uint16
+		isError bool
+		tag     string
+	}{
+		{
+			func(w http.ResponseWriter, r *http.Request) {
+				w.Write([]byte("OK"))
+			},
+			200,
+			false,
+			"OK",
+		},
+		{
+			func(w http.ResponseWriter, r *http.Request) {
+				panic("panic test")
+			},
+			0,
+			true,
+			"Panic",
+		},
+		{
+			func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(http.StatusInternalServerError)
+				w.Write([]byte("InternalServerError"))
+			},
+			500,
+			true,
+			"InternalServerError",
+		},
+	}
+
+	for _, tt := range tests {
+		testCase := tt
+		t.Run(testCase.tag, func(t *testing.T) {
+			mux := http.NewServeMux()
+			mux.HandleFunc("/root", testCase.handler)
+			tr := &mocktracer.MockTracer{}
+			srv := httptest.NewServer(MiddlewareFunc(tr, mux.ServeHTTP))
+			defer srv.Close()
+
+			_, err := http.Get(srv.URL + "/root")
+			if err != nil {
+				t.Logf("server returned error: %v", err)
+			}
+
+			spans := tr.FinishedSpans()
+			if got, want := len(spans), 1; got != want {
+				t.Fatalf("got %d spans, expected %d", got, want)
+			}
+			actualStatus := spans[0].Tag(string(ext.HTTPStatusCode)).(uint16)
+			if testCase.status != actualStatus {
+				t.Fatalf("got status code %d, expected %d", actualStatus, testCase.status)
+			}
+			actualErr, ok := spans[0].Tag(string(ext.Error)).(bool)
+			if !ok {
+				actualErr = false
+			}
+			if testCase.isError != actualErr {
+				t.Fatalf("got span error %v, expected %v", actualErr, testCase.isError)
+			}
+		})
+	}
+}

nethttp/status-code-tracker-old.go

@@ -9,14 +9,24 @@ import (
 
 type statusCodeTracker struct {
 	http.ResponseWriter
-	status int
+	status      int
+	wroteheader bool
 }
 
 func (w *statusCodeTracker) WriteHeader(status int) {
 	w.status = status
+	w.wroteheader = true
 	w.ResponseWriter.WriteHeader(status)
 }
 
+func (w *statusCodeTracker) Write(b []byte) (int, error) {
+	if !w.wroteheader {
+		w.wroteheader = true
+		w.status = 200
+	}
+	return w.ResponseWriter.Write(b)
+}
+
 // wrappedResponseWriter returns a wrapped version of the original
 // ResponseWriter and only implements the same combination of additional
 // interfaces as the original.  This implementation is based on

nethttp/status-code-tracker.go

@@ -9,14 +9,24 @@ import (
 
 type statusCodeTracker struct {
 	http.ResponseWriter
-	status int
+	status      int
+	wroteheader bool
 }
 
 func (w *statusCodeTracker) WriteHeader(status int) {
 	w.status = status
+	w.wroteheader = true
 	w.ResponseWriter.WriteHeader(status)
 }
 
+func (w *statusCodeTracker) Write(b []byte) (int, error) {
+	if !w.wroteheader {
+		w.wroteheader = true
+		w.status = 200
+	}
+	return w.ResponseWriter.Write(b)
+}
+
 // wrappedResponseWriter returns a wrapped version of the original
 // ResponseWriter and only implements the same combination of additional
 // interfaces as the original.  This implementation is based on