From 765ff6ca1ad4b17e18127e6a4264d7ccfc3fb261 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Carlos=20Ch=C3=A1vez?= <jcchavezs@gmail.com>
Date: Fri, 14 Dec 2018 13:09:06 +0100
Subject: [PATCH] [#812] adds request sampler into zipkin tracing.

---
 tracing/zipkin/http.go      |  6 ++++++
 tracing/zipkin/http_test.go | 40 ++++++++++++++++++++++++++++++++++++-
 tracing/zipkin/options.go   | 23 ++++++++++++++++-----
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/tracing/zipkin/http.go b/tracing/zipkin/http.go
index 9dd515d5b..85f079bad 100644
--- a/tracing/zipkin/http.go
+++ b/tracing/zipkin/http.go
@@ -156,6 +156,12 @@ func HTTPServerTrace(tracer *zipkin.Tracer, options ...TracerOption) kithttp.Ser
 
 			if config.propagate {
 				spanContext = tracer.Extract(b3.ExtractHTTP(req))
+
+				if config.requestSampler != nil && spanContext.Sampled == nil {
+					sample := config.requestSampler(req)
+					spanContext.Sampled = &sample
+				}
+
 				if spanContext.Err != nil {
 					config.logger.Log("err", spanContext.Err)
 				}
diff --git a/tracing/zipkin/http_test.go b/tracing/zipkin/http_test.go
index 64c94439f..c1164a7fd 100644
--- a/tracing/zipkin/http_test.go
+++ b/tracing/zipkin/http_test.go
@@ -27,7 +27,7 @@ const (
 	testTagValue = "test_value"
 )
 
-func TestHttpClientTracePropagatesParentSpan(t *testing.T) {
+func TestHTTPClientTracePropagatesParentSpan(t *testing.T) {
 	rec := recorder.NewReporter()
 	defer rec.Close()
 
@@ -220,3 +220,41 @@ func TestHTTPServerTrace(t *testing.T) {
 		t.Fatalf("incorrect error tag, want %s, have %s", want, have)
 	}
 }
+
+func TestHTTPServerTraceIsRequestBasedSampled(t *testing.T) {
+	rec := recorder.NewReporter()
+	defer rec.Close()
+
+	const httpMethod = "DELETE"
+
+	tr, _ := zipkin.NewTracer(rec)
+
+	handler := kithttp.NewServer(
+		endpoint.Nop,
+		func(context.Context, *http.Request) (interface{}, error) { return nil, nil },
+		func(context.Context, http.ResponseWriter, interface{}) error { return nil },
+		zipkinkit.HTTPServerTrace(tr, zipkinkit.RequestSampler(func(r *http.Request) bool {
+			return r.Method == httpMethod
+		})),
+	)
+
+	server := httptest.NewServer(handler)
+	defer server.Close()
+
+	req, err := http.NewRequest(httpMethod, server.URL, nil)
+	if err != nil {
+		t.Fatalf("unable to create HTTP request: %s", err.Error())
+	}
+
+	client := http.Client{}
+	resp, err := client.Do(req)
+	if err != nil {
+		t.Fatalf("unable to send HTTP request: %s", err.Error())
+	}
+	resp.Body.Close()
+
+	spans := rec.Flush()
+	if want, have := 1, len(spans); want != have {
+		t.Fatalf("incorrect number of spans, want %d, have %d", want, have)
+	}
+}
diff --git a/tracing/zipkin/options.go b/tracing/zipkin/options.go
index bdb4f97ee..8d54801b7 100644
--- a/tracing/zipkin/options.go
+++ b/tracing/zipkin/options.go
@@ -1,6 +1,10 @@
 package zipkin
 
-import "github.com/go-kit/kit/log"
+import (
+	"net/http"
+
+	"github.com/go-kit/kit/log"
+)
 
 // TracerOption allows for functional options to our Zipkin tracing middleware.
 type TracerOption func(o *tracerOptions)
@@ -46,9 +50,18 @@ func AllowPropagation(propagate bool) TracerOption {
 	}
 }
 
+// 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) TracerOption {
+	return func(o *tracerOptions) {
+		o.requestSampler = sampleFunc
+	}
+}
+
 type tracerOptions struct {
-	tags      map[string]string
-	name      string
-	logger    log.Logger
-	propagate bool
+	tags           map[string]string
+	name           string
+	logger         log.Logger
+	propagate      bool
+	requestSampler func(r *http.Request) bool
 }