Merge pull request #134 from basvanbeek/http-server-request-sampler

Improved http server request sampler logic

Author: basvanbeek

middleware/http/request_sampler.go

@@ -0,0 +1,21 @@
+package http
+
+import "net/http"
+
+// RequestSamplerFunc can be implemented for client and/or server side sampling decisions that can override the existing
+// upstream sampling decision. If the implementation returns nil, the existing sampling decision stays as is.
+type RequestSamplerFunc func(r *http.Request) *bool
+
+// Sample is a convenience function that returns a pointer to a boolean true. Use this for RequestSamplerFuncs when
+// wanting the RequestSampler to override the sampling decision to yes.
+func Sample() *bool {
+	sample := true
+	return &sample
+}
+
+// Discard is a convenience function that returns a pointer to a boolean false. Use this for RequestSamplerFuncs when
+// wanting the RequestSampler to override the sampling decision to no.
+func Discard() *bool {
+	sample := false
+	return &sample
+}

middleware/http/server.go

@@ -31,7 +31,7 @@ type handler struct {
 	next            http.Handler
 	tagResponseSize bool
 	defaultTags     map[string]string
-	requestSampler  func(r *http.Request) bool
+	requestSampler  RequestSamplerFunc
 	errHandler      ErrHandler
 }
 
@@ -64,8 +64,9 @@ func SpanName(name string) ServerOption {
 }
 
 // RequestSampler allows one to set the sampling decision based on the details
-// found in the http.Request.
-func RequestSampler(sampleFunc func(r *http.Request) bool) ServerOption {
+// found in the http.Request. If wanting to keep the existing sampling decision
+// from upstream as is, this function should return nil.
+func RequestSampler(sampleFunc RequestSamplerFunc) ServerOption {
 	return func(h *handler) {
 		h.requestSampler = sampleFunc
 	}
@@ -100,9 +101,10 @@ func (h handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	// try to extract B3 Headers from upstream
 	sc := h.tracer.Extract(b3.ExtractHTTP(r))
 
-	if h.requestSampler != nil && sc.Sampled == nil {
-		sample := h.requestSampler(r)
-		sc.Sampled = &sample
+	if h.requestSampler != nil {
+		if sample := h.requestSampler(r); sample != nil {
+			sc.Sampled = sample
+		}
 	}
 
 	remoteEndpoint, _ := zipkin.NewEndpoint("", r.RemoteAddr)

middleware/http/server_test.go

@@ -214,13 +214,14 @@ func TestHTTPRequestSampler(t *testing.T) {
 		httpHandlerFunc = http.HandlerFunc(httpHandler(200, nil, bytes.NewBufferString("")))
 	)
 
-	samplers := [](func(r *http.Request) bool){
+	samplers := [](func(r *http.Request) *bool){
 		nil,
-		func(r *http.Request) bool { return true },
-		func(r *http.Request) bool { return false },
+		func(r *http.Request) *bool { return mw.Sample() },
+		func(r *http.Request) *bool { return mw.Discard() },
+		func(r *http.Request) *bool { return nil },
 	}
 
-	for _, sampler := range samplers {
+	for idx, sampler := range samplers {
 		tr, _ := zipkin.NewTracer(spanRecorder, zipkin.WithLocalEndpoint(lep), zipkin.WithSampler(zipkin.AlwaysSample))
 
 		request, err := http.NewRequest(methodType, "/test", requestBuf)
@@ -235,16 +236,16 @@ func TestHTTPRequestSampler(t *testing.T) {
 		spans := spanRecorder.Flush()
 
 		sampledSpans := 0
-		if sampler == nil || sampler(request) {
+		if sampler == nil || sampler(request) == nil || *(sampler(request)) {
 			sampledSpans = 1
 		}
 
 		if want, have := sampledSpans, len(spans); want != have {
-			t.Errorf("Expected %d spans, got %d", want, have)
+			t.Errorf("[%d] Expected %d spans, got %d", idx, want, have)
 		}
 	}
 
-	for _, sampler := range samplers {
+	for idx, sampler := range samplers {
 		tr, _ := zipkin.NewTracer(spanRecorder, zipkin.WithLocalEndpoint(lep), zipkin.WithSampler(zipkin.NeverSample))
 
 		request, err := http.NewRequest(methodType, "/test", requestBuf)
@@ -259,12 +260,12 @@ func TestHTTPRequestSampler(t *testing.T) {
 		spans := spanRecorder.Flush()
 
 		sampledSpans := 0
-		if sampler != nil && sampler(request) {
+		if sampler != nil && sampler(request) != nil && *(sampler(request)) {
 			sampledSpans = 1
 		}
 
 		if want, have := sampledSpans, len(spans); want != have {
-			t.Errorf("Expected %d spans, got %d", want, have)
+			t.Errorf("[%d] Expected %d spans, got %d", idx, want, have)
 		}
 	}
 

middleware/http/transport.go

@@ -57,7 +57,7 @@ type transport struct {
 	errHandler        ErrHandler
 	errResponseReader *ErrResponseReader
 	logger            *log.Logger
-	requestSampler    func(r *http.Request) *bool
+	requestSampler    RequestSamplerFunc
 }
 
 // TransportOption allows one to configure optional transport configuration.
@@ -112,7 +112,7 @@ func TransportLogger(l *log.Logger) TransportOption {
 // sampling decision contained in the context. The function returns a *bool,
 // if returning nil, sampling decision is not being changed whereas returning
 // something else than nil is being used as sampling decision.
-func TransportRequestSampler(sampleFunc func(r *http.Request) *bool) TransportOption {
+func TransportRequestSampler(sampleFunc RequestSamplerFunc) TransportOption {
 	return func(t *transport) {
 		t.requestSampler = sampleFunc
 	}

middleware/http/transport_test.go

@@ -128,14 +128,10 @@ func TestRoundTripErrResponseReadingSuccess(t *testing.T) {
 	}
 }
 
-func boolToPtr(b bool) *bool {
-	return &b
-}
-
 func TestTransportRequestSamplerOverridesSamplingFromContext(t *testing.T) {
 	cases := []struct {
 		Sampler          func(uint64) bool
-		RequestSampler   func(*http.Request) *bool
+		RequestSampler   RequestSamplerFunc
 		ExpectedSampling string
 	}{
 		// Test proper handling when there is no RequestSampler
@@ -153,13 +149,13 @@ func TestTransportRequestSamplerOverridesSamplingFromContext(t *testing.T) {
 		// Test RequestSampler override sample -> no sample
 		{
 			Sampler:          zipkin.AlwaysSample,
-			RequestSampler:   func(_ *http.Request) *bool { return boolToPtr(false) },
+			RequestSampler:   func(_ *http.Request) *bool { return Discard() },
 			ExpectedSampling: "0",
 		},
 		// Test RequestSampler override no sample -> sample
 		{
 			Sampler:          zipkin.NeverSample,
-			RequestSampler:   func(_ *http.Request) *bool { return boolToPtr(true) },
+			RequestSampler:   func(_ *http.Request) *bool { return Sample() },
 			ExpectedSampling: "1",
 		},
 		// Test RequestSampler pass through of sampled decision