From eed4d38387cb9c74f7a29cb3eb9b06155a09d259 Mon Sep 17 00:00:00 2001
From: Gabriel Aszalos <gabriel.aszalos@gmail.com>
Date: Thu, 18 Apr 2019 03:16:48 -0400
Subject: [PATCH] ddtrace/tracer: add manual.keep and manual.drop tags (#430)

Add tag aliases for simplifying user-controlled sampling priorities.
---
 ddtrace/ext/tags.go         |  8 ++++++++
 ddtrace/tracer/span.go      | 41 +++++++++++++++++++++++++++----------
 ddtrace/tracer/span_test.go | 12 +++++++++++
 3 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/ddtrace/ext/tags.go b/ddtrace/ext/tags.go
index 3a5b19c304..c21f848259 100644
--- a/ddtrace/ext/tags.go
+++ b/ddtrace/ext/tags.go
@@ -67,4 +67,12 @@ const (
 	// AnalyticsEvent specifies whether the span should be recorded as a Trace
 	// Search & Analytics event.
 	AnalyticsEvent = "analytics.event"
+
+	// ManualKeep is a tag which specifies that the trace to which this span
+	// belongs to should be kept when set to true.
+	ManualKeep = "manual.keep"
+
+	// ManualDrop is a tag which specifies that the trace to which this span
+	// belongs to should be dropped when set to true.
+	ManualDrop = "manual.drop"
 )
diff --git a/ddtrace/tracer/span.go b/ddtrace/tracer/span.go
index 0e9f7d529c..e8d9432bc1 100644
--- a/ddtrace/tracer/span.go
+++ b/ddtrace/tracer/span.go
@@ -84,17 +84,10 @@ func (s *span) SetTag(key string, value interface{}) {
 	case ext.Error:
 		s.setTagError(value, true)
 		return
-	case ext.AnalyticsEvent:
-		// "analytics.event" is a boolean alias for setting the event sampling
-		// rate to 0.0 or 1.0
-		if set, ok := value.(bool); ok {
-			if set {
-				s.setTagNumeric(ext.EventSampleRate, 1.0)
-			} else {
-				s.setTagNumeric(ext.EventSampleRate, 0.0)
-			}
-			return
-		}
+	}
+	if v, ok := value.(bool); ok {
+		s.setTagBool(key, v)
+		return
 	}
 	if v, ok := value.(string); ok {
 		s.setTagString(key, v)
@@ -158,6 +151,32 @@ func (s *span) setTagString(key, v string) {
 	}
 }
 
+// setTagBool sets a boolean tag on the span.
+func (s *span) setTagBool(key string, v bool) {
+	switch key {
+	case ext.AnalyticsEvent:
+		if v {
+			s.setTagNumeric(ext.EventSampleRate, 1.0)
+		} else {
+			s.setTagNumeric(ext.EventSampleRate, 0.0)
+		}
+	case ext.ManualDrop:
+		if v {
+			s.setTagNumeric(ext.SamplingPriority, ext.PriorityUserReject)
+		}
+	case ext.ManualKeep:
+		if v {
+			s.setTagNumeric(ext.SamplingPriority, ext.PriorityUserKeep)
+		}
+	default:
+		if v {
+			s.setTagString(key, "true")
+		} else {
+			s.setTagString(key, "false")
+		}
+	}
+}
+
 // setTagNumeric sets a numeric tag, in our case called a metric. This method
 // is not safe for concurrent use.
 func (s *span) setTagNumeric(key string, v float64) {
diff --git a/ddtrace/tracer/span_test.go b/ddtrace/tracer/span_test.go
index caaa83591e..e03402232d 100644
--- a/ddtrace/tracer/span_test.go
+++ b/ddtrace/tracer/span_test.go
@@ -167,6 +167,18 @@ func TestSpanSetTag(t *testing.T) {
 
 	span.SetTag(ext.AnalyticsEvent, false)
 	assert.Equal(0.0, span.Metrics[ext.EventSampleRate])
+
+	span.SetTag(ext.ManualDrop, true)
+	assert.Equal(-1., span.Metrics[keySamplingPriority])
+
+	span.SetTag(ext.ManualKeep, true)
+	assert.Equal(2., span.Metrics[keySamplingPriority])
+
+	span.SetTag("some.bool", true)
+	assert.Equal("true", span.Meta["some.bool"])
+
+	span.SetTag("some.other.bool", false)
+	assert.Equal("false", span.Meta["some.other.bool"])
 }
 
 func TestSpanSetDatadogTags(t *testing.T) {