Allow overriding health endpoint check in handler (#1177)

* Allow overriding health endpoint check in handler

* Clarify fallback

* Clarify godoc for IsHealthEndpoint

* Combine TestIgnoreHealthz and TestHandlerIsHealthEndpoint into a single test TestIgnoreHealthEndpoints, and add cases

* Even more explicit godoc

* IsHealthEndpointFunc now accepts the http.Request

Author: tjcelaya

plugin/ochttp/server.go

@@ -70,6 +70,12 @@ type Handler struct {
 	// from the information found in the incoming HTTP Request. By default the
 	// name equals the URL Path.
 	FormatSpanName func(*http.Request) string
+
+	// IsHealthEndpoint holds the function to use for determining if the
+	// incoming HTTP request should be considered a health check. This is in
+	// addition to the private isHealthEndpoint func which may also indicate
+	// tracing should be skipped.
+	IsHealthEndpoint func(*http.Request) bool
 }
 
 func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
@@ -87,7 +93,7 @@ func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 }
 
 func (h *Handler) startTrace(w http.ResponseWriter, r *http.Request) (*http.Request, func()) {
-	if isHealthEndpoint(r.URL.Path) {
+	if h.IsHealthEndpoint != nil && h.IsHealthEndpoint(r) || isHealthEndpoint(r.URL.Path) {
 		return r, func() {}
 	}
 	var name string

plugin/ochttp/server_test.go

@@ -555,40 +555,61 @@ func TestHandlerImplementsHTTPCloseNotify(t *testing.T) {
 	}
 }
 
-func TestIgnoreHealthz(t *testing.T) {
-	var spans int
+func testHealthEndpointSkipArray(r *http.Request) bool {
+	for _, toSkip := range []string{"/health", "/metrics"} {
+		if r.URL.Path == toSkip {
+			return true
+		}
+	}
+	return false
+}
 
-	ts := httptest.NewServer(&Handler{
-		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			span := trace.FromContext(r.Context())
-			if span != nil {
-				spans++
-			}
-			fmt.Fprint(w, "ok")
-		}),
-		StartOptions: trace.StartOptions{
-			Sampler: trace.AlwaysSample(),
-		},
-	})
-	defer ts.Close()
+func TestIgnoreHealthEndpoints(t *testing.T) {
+	var spans int
 
 	client := &http.Client{}
+	tests := []struct {
+		path               string
+		healthEndpointFunc func(*http.Request) bool
+	}{
+		{"/healthz", nil},
+		{"/_ah/health", nil},
+		{"/healthz", testHealthEndpointSkipArray},
+		{"/_ah/health", testHealthEndpointSkipArray},
+		{"/health", testHealthEndpointSkipArray},
+		{"/metrics", testHealthEndpointSkipArray},
+	}
+	for _, tt := range tests {
+		t.Run(tt.path, func(t *testing.T) {
+			ts := httptest.NewServer(&Handler{
+				Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+					span := trace.FromContext(r.Context())
+					if span != nil {
+						spans++
+					}
+					fmt.Fprint(w, "ok")
+				}),
+				StartOptions: trace.StartOptions{
+					Sampler: trace.AlwaysSample(),
+				},
+				IsHealthEndpoint: tt.healthEndpointFunc,
+			})
+			defer ts.Close()
 
-	for _, path := range []string{"/healthz", "/_ah/health"} {
-		resp, err := client.Get(ts.URL + path)
-		if err != nil {
-			t.Fatalf("Cannot GET %q: %v", path, err)
-		}
-
-		b, err := ioutil.ReadAll(resp.Body)
-		if err != nil {
-			t.Fatalf("Cannot read body for %q: %v", path, err)
-		}
+			resp, err := client.Get(ts.URL + tt.path)
+			if err != nil {
+				t.Fatalf("Cannot GET %q: %v", tt.path, err)
+			}
+			b, err := ioutil.ReadAll(resp.Body)
+			if err != nil {
+				t.Fatalf("Cannot read body for %q: %v", tt.path, err)
+			}
 
-		if got, want := string(b), "ok"; got != want {
-			t.Fatalf("Body for %q = %q; want %q", path, got, want)
-		}
-		resp.Body.Close()
+			if got, want := string(b), "ok"; got != want {
+				t.Fatalf("Body for %q = %q; want %q", tt.path, got, want)
+			}
+			resp.Body.Close()
+		})
 	}
 
 	if spans > 0 {