adds convenience functions Sample() and Discard() for request sampler…

… decisions
kujenga · Jul 10, 2019 · 8774349 · 8774349
1 parent 0ef9e89
commit 8774349
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 18 deletions.
diff --git a/middleware/http/request_sampler.go b/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
+}
diff --git a/middleware/http/server.go b/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
 }
 
@@ -66,7 +66,7 @@ func SpanName(name string) ServerOption {
 // RequestSampler allows one to set the sampling decision based on the details
 // 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 func(r *http.Request) *bool) ServerOption {
+func RequestSampler(sampleFunc RequestSamplerFunc) ServerOption {
 	return func(h *handler) {
 		h.requestSampler = sampleFunc
 	}

diff --git a/middleware/http/server_test.go b/middleware/http/server_test.go
@@ -207,10 +207,6 @@ func TestHTTPDefaultSpanName(t *testing.T) {
 
 func TestHTTPRequestSampler(t *testing.T) {
 	var (
-		sample            = true
-		noSample          = false
-		passThrough *bool = nil
-
 		spanRecorder    = &recorder.ReporterRecorder{}
 		httpRecorder    = httptest.NewRecorder()
 		requestBuf      = bytes.NewBufferString("incoming data")
@@ -220,9 +216,9 @@ func TestHTTPRequestSampler(t *testing.T) {
 
 	samplers := [](func(r *http.Request) *bool){
 		nil,
-		func(r *http.Request) *bool { return &sample },
-		func(r *http.Request) *bool { return &noSample },
-		func(r *http.Request) *bool { return passThrough },
+		func(r *http.Request) *bool { return mw.Sample() },
+		func(r *http.Request) *bool { return mw.Discard() },
+		func(r *http.Request) *bool { return nil },
 	}
 
 	for idx, sampler := range samplers {

diff --git a/middleware/http/transport.go b/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
 	}

diff --git a/middleware/http/transport_test.go b/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