Merge pull request #133 from openzipkin/128_request_sampler_for_client

feat(#128): adds support for request sampler in client.

Author: basvanbeek

middleware/http/transport.go

@@ -57,6 +57,7 @@ type transport struct {
 	errHandler        ErrHandler
 	errResponseReader *ErrResponseReader
 	logger            *log.Logger
+	requestSampler    func(r *http.Request) *bool
 }
 
 // TransportOption allows one to configure optional transport configuration.
@@ -106,6 +107,17 @@ func TransportLogger(l *log.Logger) TransportOption {
 	}
 }
 
+// TransportRequestSampler allows one to set the sampling decision based on
+// the details found in the http.Request. It has preference over the existing
+// 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 {
+	return func(t *transport) {
+		t.requestSampler = sampleFunc
+	}
+}
+
 // NewTransport returns a new Zipkin instrumented http RoundTripper which can be
 // used with a standard library http Client.
 func NewTransport(tracer *zipkin.Tracer, options ...TransportOption) (http.RoundTripper, error) {
@@ -167,7 +179,14 @@ func (t *transport) RoundTrip(req *http.Request) (res *http.Response, err error)
 	zipkin.TagHTTPMethod.Set(sp, req.Method)
 	zipkin.TagHTTPPath.Set(sp, req.URL.Path)
 
-	_ = b3.InjectHTTP(req)(sp.Context())
+	spCtx := sp.Context()
+	if t.requestSampler != nil {
+		if shouldSample := t.requestSampler(req); shouldSample != nil {
+			spCtx.Sampled = shouldSample
+		}
+	}
+
+	_ = b3.InjectHTTP(req)(spCtx)
 
 	res, err = t.rt.RoundTrip(req)
 	if err != nil {

middleware/http/transport_test.go

@@ -1,6 +1,7 @@
 package http
 
 import (
+	"context"
 	"errors"
 	"io"
 	"io/ioutil"
@@ -9,6 +10,7 @@ import (
 	"testing"
 
 	zipkin "github.com/openzipkin/zipkin-go"
+	"github.com/openzipkin/zipkin-go/reporter/recorder"
 )
 
 type errRoundTripper struct {
@@ -125,3 +127,90 @@ func TestRoundTripErrResponseReadingSuccess(t *testing.T) {
 		t.Errorf("unexpected body: want %s, have %s", want, have)
 	}
 }
+
+func boolToPtr(b bool) *bool {
+	return &b
+}
+
+func TestTransportRequestSamplerOverridesSamplingFromContext(t *testing.T) {
+	cases := []struct {
+		Sampler          func(uint64) bool
+		RequestSampler   func(*http.Request) *bool
+		ExpectedSampling string
+	}{
+		// Test proper handling when there is no RequestSampler
+		{
+			Sampler:          zipkin.AlwaysSample,
+			RequestSampler:   nil,
+			ExpectedSampling: "1",
+		},
+		// Test proper handling when there is no RequestSampler
+		{
+			Sampler:          zipkin.NeverSample,
+			RequestSampler:   nil,
+			ExpectedSampling: "0",
+		},
+		// Test RequestSampler override sample -> no sample
+		{
+			Sampler:          zipkin.AlwaysSample,
+			RequestSampler:   func(_ *http.Request) *bool { return boolToPtr(false) },
+			ExpectedSampling: "0",
+		},
+		// Test RequestSampler override no sample -> sample
+		{
+			Sampler:          zipkin.NeverSample,
+			RequestSampler:   func(_ *http.Request) *bool { return boolToPtr(true) },
+			ExpectedSampling: "1",
+		},
+		// Test RequestSampler pass through of sampled decision
+		{
+			Sampler: zipkin.AlwaysSample,
+			RequestSampler: func(r *http.Request) *bool {
+				return nil
+			},
+			ExpectedSampling: "1",
+		},
+		// Test RequestSampler pass through of not sampled decision
+		{
+			Sampler: zipkin.NeverSample,
+			RequestSampler: func(r *http.Request) *bool {
+				return nil
+			},
+			ExpectedSampling: "0",
+		},
+	}
+
+	for i, c := range cases {
+		srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
+			if want, have := c.ExpectedSampling, r.Header.Get("x-b3-sampled"); want != have {
+				t.Errorf("unexpected sampling decision in case #%d, want %q, have %q", i, want, have)
+			}
+		}))
+
+		// we need to use a valid reporter or Tracer will implement noop mode which makes this test invalid
+		rep := recorder.NewReporter()
+
+		tracer, err := zipkin.NewTracer(rep, zipkin.WithSampler(c.Sampler))
+		if err != nil {
+			t.Fatalf("unexpected error when creating tracer: %v", err)
+		}
+
+		sp := tracer.StartSpan("op1")
+		defer sp.Finish()
+		ctx := zipkin.NewContext(context.Background(), sp)
+
+		req, _ := http.NewRequest("GET", srv.URL, nil)
+		transport, _ := NewTransport(
+			tracer,
+			TransportRequestSampler(c.RequestSampler),
+		)
+
+		_, err = transport.RoundTrip(req.WithContext(ctx))
+		if err != nil {
+			t.Fatalf("unexpected error: %v", err)
+		}
+
+		rep.Close()
+		srv.Close()
+	}
+}