Support Datadog tracer priority sampling

vendor/gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer/sampler.go

@@ -1,10 +1,13 @@
 package tracer
 
 import (
+	"encoding/json"
+	"io"
 	"math"
 	"sync"
 
 	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace"
+	"gopkg.in/DataDog/dd-trace-go.v1/ddtrace/ext"
 )
 
 // Sampler is the generic interface of any sampler. It must be safe for concurrent use.
@@ -59,14 +62,83 @@ const knuthFactor = uint64(1111111111111111111)
 
 // Sample returns true if the given span should be sampled.
 func (r *rateSampler) Sample(spn ddtrace.Span) bool {
+	if r.rate == 1 {
+		// fast path
+		return true
+	}
 	s, ok := spn.(*span)
 	if !ok {
 		return false
 	}
 	r.RLock()
 	defer r.RUnlock()
-	if r.rate < 1 {
-		return s.TraceID*knuthFactor < uint64(r.rate*math.MaxUint64)
+	return sampledByRate(s.TraceID, r.rate)
+}
+
+// sampledByRate verifies if the number n should be sampled at the specified
+// rate.
+func sampledByRate(n uint64, rate float64) bool {
+	if rate < 1 {
+		return n*knuthFactor < uint64(rate*math.MaxUint64)
 	}
 	return true
 }
+
+// prioritySampler holds a set of per-service sampling rates and applies
+// them to spans.
+type prioritySampler struct {
+	mu          sync.RWMutex
+	rates       map[string]float64
+	defaultRate float64
+}
+
+func newPrioritySampler() *prioritySampler {
+	return &prioritySampler{
+		rates:       make(map[string]float64),
+		defaultRate: 1.,
+	}
+}
+
+// readRatesJSON will try to read the rates as JSON from the given io.ReadCloser.
+func (ps *prioritySampler) readRatesJSON(rc io.ReadCloser) error {
+	var payload struct {
+		Rates map[string]float64 `json:"rate_by_service"`
+	}
+	if err := json.NewDecoder(rc).Decode(&payload); err != nil {
+		return err
+	}
+	rc.Close()
+	const defaultRateKey = "service:,env:"
+	ps.mu.Lock()
+	defer ps.mu.Unlock()
+	ps.rates = payload.Rates
+	if v, ok := ps.rates[defaultRateKey]; ok {
+		ps.defaultRate = v
+		delete(ps.rates, defaultRateKey)
+	}
+	return nil
+}
+
+// getRate returns the sampling rate to be used for the given span. Callers must
+// guard the span.
+func (ps *prioritySampler) getRate(spn *span) float64 {
+	key := "service:" + spn.Service + ",env:" + spn.Meta[ext.Environment]
+	ps.mu.RLock()
+	defer ps.mu.RUnlock()
+	if rate, ok := ps.rates[key]; ok {
+		return rate
+	}
+	return ps.defaultRate
+}
+
+// apply applies sampling priority to the given span. Caller must ensure it is safe
+// to modify the span.
+func (ps *prioritySampler) apply(spn *span) {
+	rate := ps.getRate(spn)
+	if sampledByRate(spn.TraceID, rate) {
+		spn.SetTag(ext.SamplingPriority, ext.PriorityAutoKeep)
+	} else {
+		spn.SetTag(ext.SamplingPriority, ext.PriorityAutoReject)
+	}
+	spn.SetTag(samplingPriorityRateKey, rate)
+}