ddtrace/tracer: adding support for single span matching

ddtrace/ext/tags.go

@@ -79,6 +79,19 @@ const (
 	// Search & Analytics event.
 	AnalyticsEvent = "analytics.event"
 
+	// SpanSamplingMechanism is the metric key holding a span sampling rule that a span was kept on
+	SpanSamplingMechanism = "_dd.span_sampling.mechanism"
+
+	// SingleSpanSamplingMechanism is a constant reserved to indicate that a span was kept on account of a span sampling rule.
+	SingleSpanSamplingMechanism = 8
+
+	// SingleSpanSamplingRuleRate is the metric key containing the configured sampling probability for the span sampling rule.
+	SingleSpanSamplingRuleRate = "_dd.span_sampling.rule_rate"
+
+	// SingleSpanSamplingMPS is the metric key contains the configured limit for the span sampling rule that the span matched.
+	//	// If there is no configured limit, then this tag is omitted.
+	SingleSpanSamplingMPS = "_dd.span_sampling.max_per_second"
+
 	// 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"

ddtrace/tracer/log_test.go

@@ -126,8 +126,8 @@ func TestLogSamplingRules(t *testing.T) {
 
 	lines := removeAppSec(tp.Lines())
 	assert.Len(lines, 2)
-	assert.Contains(lines[0], "WARN: at index 4: ignoring rule {Service: Name: Rate:9.10}: rate is out of [0.0, 1.0] range")
-	assert.Regexp(`Datadog Tracer v[0-9]+\.[0-9]+\.[0-9]+ WARN: DIAGNOSTICS Error\(s\) parsing DD_TRACE_SAMPLING_RULES: found errors:\n\tat index 1: rate not provided\n\tat index 3: rate not provided$`, lines[1])
+	assert.Contains(lines[0], "WARN: at index 4: ignoring rule {Service: Name: Rate:9.10 MaxPerSecond:0}: rate is out of [0.0, 1.0] range")
+	assert.Regexp(`Datadog Tracer v[0-9]+\.[0-9]+\.[0-9]+ WARN: DIAGNOSTICS Error\(s\) parsing sampling rules: found errors:\n\tat index 1: rate not provided\n\tat index 3: rate not provided$`, lines[1])
 }
 
 func TestLogAgentReachable(t *testing.T) {

ddtrace/tracer/sampler.go

@@ -188,19 +188,62 @@ func newRulesSampler(rules []SamplingRule) *rulesSampler {
 	}
 }
 
-// samplingRulesFromEnv parses sampling rules from the DD_TRACE_SAMPLING_RULES
-// environment variable.
+// samplingRulesFromEnv parses sampling rules from the DD_TRACE_SAMPLING_RULES environment variable
+// and combines them with spanSamplingRules
 func samplingRulesFromEnv() ([]SamplingRule, error) {
-	rulesFromEnv := os.Getenv("DD_TRACE_SAMPLING_RULES")
-	if rulesFromEnv == "" {
+	errs := []string{}
+	spanRules, err := spanSamplingRulesFromEnv()
+	if err != nil {
+		errs = append(errs, err.Error())
+	}
+	traceRules, err := processSamplingRulesFromEnv([]byte(os.Getenv("DD_TRACE_SAMPLING_RULES")), TraceSamplingRuleType)
+	if err != nil {
+		errs = append(errs, err.Error())
+	}
+
+	rules := append(spanRules, traceRules...)
+	if len(errs) != 0 {
+		return rules, fmt.Errorf("%s", strings.Join(errs, "\n\t"))
+	}
+	return rules, nil
+}
+
+// spanSamplingRulesFromEnv parses sampling rules from the DD_SPAN_SAMPLING_RULES and
+// DD_SPAN_SAMPLING_RULES_FILE environment variables.
+func spanSamplingRulesFromEnv() ([]SamplingRule, error) {
+	errs := []string{}
+	rules, err := processSamplingRulesFromEnv([]byte(os.Getenv("DD_SPAN_SAMPLING_RULES")), SingleSpanSamplingType)
+	if err != nil {
+		errs = append(errs, err.Error())
+	}
+	if len(rules) != 0 {
+		if os.Getenv("DD_SPAN_SAMPLING_RULES_FILE") != "" {
+			log.Warn("DIAGNOSTICS Error(s): DD_SPAN_SAMPLING_RULES is available and will take precedence over DD_SPAN_SAMPLING_RULES_FILE")
+		}
+		return rules, err
+	}
+	rulesFile := os.Getenv("DD_SPAN_SAMPLING_RULES_FILE")
+	if rulesFile == "" {
+		return rules, err
+	}
+	rulesFromEnvFile, err := os.ReadFile(rulesFile)
+	if err != nil {
+		log.Warn("Couldn't read file from DD_SPAN_SAMPLING_RULES_FILE")
+	}
+	return processSamplingRulesFromEnv(rulesFromEnvFile, SingleSpanSamplingType)
+}
+
+func processSamplingRulesFromEnv(b []byte, ruleType SamplingRuleType) ([]SamplingRule, error) {
+	if len(b) == 0 {
 		return nil, nil
 	}
-	jsonRules := []struct {
-		Service string      `json:"service"`
-		Name    string      `json:"name"`
-		Rate    json.Number `json:"sample_rate"`
-	}{}
-	err := json.Unmarshal([]byte(rulesFromEnv), &jsonRules)
+	var jsonRules []struct {
+		Service      string      `json:"service"`
+		Name         string      `json:"name"`
+		Rate         json.Number `json:"sample_rate"`
+		MaxPerSecond float64     `json:"max_per_second"`
+	}
+	err := json.Unmarshal(b, &jsonRules)
 	if err != nil {
 		return nil, fmt.Errorf("error unmarshalling JSON: %v", err)
 	}
@@ -220,6 +263,32 @@ func samplingRulesFromEnv() ([]SamplingRule, error) {
 			log.Warn("at index %d: ignoring rule %+v: rate is out of [0.0, 1.0] range", i, v)
 			continue
 		}
+		if ruleType == SingleSpanSamplingType {
+			if v.Service == "" {
+				v.Service = "*"
+			}
+			srvGlob, err := globMatch(v.Service)
+			if err != nil {
+				log.Warn("at index %d: ignoring rule %+v: service name regex pattern can't be compiled", i, v)
+				continue
+			}
+			if v.Name == "" {
+				v.Name = "*"
+			}
+			opGlob, err := globMatch(v.Name)
+			if err != nil {
+				log.Warn("at index %d: ignoring rule %+v: operation name regex pattern can't be compiled", i, v)
+				continue
+			}
+			rules = append(rules, SamplingRule{
+				Service:      srvGlob,
+				Name:         opGlob,
+				Rate:         rate,
+				MaxPerSecond: v.MaxPerSecond,
+				Type:         SingleSpanSamplingType,
+			})
+			continue
+		}
 		switch {
 		case v.Service != "" && v.Name != "":
 			rules = append(rules, NameServiceRule(v.Name, v.Service, rate))
@@ -230,7 +299,7 @@ func samplingRulesFromEnv() ([]SamplingRule, error) {
 		}
 	}
 	if len(errs) != 0 {
-		return rules, fmt.Errorf("found errors:\n\t%s", strings.Join(errs, "\n\t"))
+		return rules, fmt.Errorf("\n\t%s", strings.Join(errs, "\n\t"))
 	}
 	return rules, nil
 }
@@ -300,6 +369,13 @@ func (rs *rulesSampler) apply(span *span) bool {
 		if rule.match(span) {
 			matched = true
 			rate = rule.Rate
+			if rule.Type == SingleSpanSamplingType {
+				span.setMetric(ext.SpanSamplingMechanism, ext.SingleSpanSamplingMechanism)
+				span.setMetric(ext.SingleSpanSamplingRuleRate, rule.Rate)
+				if rule.MaxPerSecond != 0 {
+					span.setMetric(ext.SingleSpanSamplingMPS, rule.MaxPerSecond)
+				}
+			}
 			break
 		}
 	}
@@ -338,13 +414,22 @@ func (rs *rulesSampler) limit() (float64, bool) {
 	return math.NaN(), false
 }
 
+type SamplingRuleType int
+
+const (
+	TraceSamplingRuleType = iota
+	SingleSpanSamplingType
+)
+
 // SamplingRule is used for applying sampling rates to spans that match
 // the service name, operation name or both.
 // For basic usage, consider using the helper functions ServiceRule, NameRule, etc.
 type SamplingRule struct {
-	Service *regexp.Regexp
-	Name    *regexp.Regexp
-	Rate    float64
+	Service      *regexp.Regexp
+	Name         *regexp.Regexp
+	Rate         float64
+	MaxPerSecond float64
+	Type         SamplingRuleType
 
 	exactService string
 	exactName    string