(WIP) Experiment for applying all conditions in a rule to a single span in a trace

config/sampler_config.go

@@ -62,6 +62,7 @@ type RulesBasedSamplerRule struct {
 	SampleRate int
 	Sampler    *RulesBasedDownstreamSampler
 	Drop       bool
+	MatchSpan  bool
 	Condition  []*RulesBasedSamplerCondition
 }
 

sample/rules.go

@@ -65,93 +65,15 @@ func (s *RulesBasedSampler) GetSampleRate(trace *types.Trace) (rate uint, keep b
 	})
 
 	for _, rule := range s.Config.Rule {
-		var matched int
+		var matched bool
 
-		for _, condition := range rule.Condition {
-		span:
-			for _, span := range trace.GetSpans() {
-				var match bool
-				value, exists := span.Data[condition.Field]
-				if !exists && s.Config.CheckNestedFields {
-					jsonStr, err := json.Marshal(span.Data)
-					if err == nil {
-						result := gjson.Get(string(jsonStr), condition.Field)
-						if result.Exists() {
-							value = result.String()
-							exists = true
-						}
-					}
-				}
-
-				switch exists {
-				case true:
-					switch condition.Operator {
-					case "exists":
-						match = exists
-					case "!=":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison != equal
-						}
-					case "=":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison == equal
-						}
-					case ">":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison == more
-						}
-					case ">=":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison == more || comparison == equal
-						}
-					case "<":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison == less
-						}
-					case "<=":
-						if comparison, ok := compare(value, condition.Value); ok {
-							match = comparison == less || comparison == equal
-						}
-					case "starts-with":
-						switch a := value.(type) {
-						case string:
-							switch b := condition.Value.(type) {
-							case string:
-								match = strings.HasPrefix(a, b)
-							}
-						}
-					case "contains":
-						switch a := value.(type) {
-						case string:
-							switch b := condition.Value.(type) {
-							case string:
-								match = strings.Contains(a, b)
-							}
-						}
-					case "does-not-contain":
-						switch a := value.(type) {
-						case string:
-							switch b := condition.Value.(type) {
-							case string:
-								match = !strings.Contains(a, b)
-							}
-						}
-					}
-				case false:
-					switch condition.Operator {
-					case "not-exists":
-						match = !exists
-					}
-				}
-
-				if match {
-					matched++
-					break span
-				}
-			}
+		if rule.MatchSpan {
+			matched = ruleMatchesSpanInTrace(trace, rule, s.Config.CheckNestedFields)
+		} else {
+			matched = ruleMatchesTrace(trace, rule, s.Config.CheckNestedFields)
 		}
 
-		if rule.Condition == nil || matched == len(rule.Condition) {
+		if matched {
 			var rate uint
 			var keep bool
 
@@ -188,6 +110,141 @@ func (s *RulesBasedSampler) GetSampleRate(trace *types.Trace) (rate uint, keep b
 	return 1, true
 }
 
+func ruleMatchesTrace(t *types.Trace, rule *config.RulesBasedSamplerRule, checkNestedFields bool) bool {
+	// We treat a rule with no conditions as a match.
+	if rule.Condition == nil {
+		return true
+	}
+
+	var matched int
+
+	for _, condition := range rule.Condition {
+	span:
+		for _, span := range t.GetSpans() {
+			value, exists := extractValueFromSpan(span, condition, checkNestedFields)
+
+			if conditionMatchesValue(condition, value, exists) {
+				matched++
+				break span
+			}
+		}
+	}
+
+	return matched == len(rule.Condition)
+}
+
+func ruleMatchesSpanInTrace(trace *types.Trace, rule *config.RulesBasedSamplerRule, checkNestedFields bool) bool {
+	// We treat a rule with no conditions as a match.
+	if rule.Condition == nil {
+		return true
+	}
+
+	for _, span := range trace.GetSpans() {
+		// the number of conditions that match this span.
+		// incremented later on after we match a condition
+		// since we need to match *all* conditions on a single span, we reset in each iteration of the loop.
+		matchCount := 0
+		for _, condition := range rule.Condition {
+			// whether this condition is matched by this span.
+			value, exists := extractValueFromSpan(span, condition, checkNestedFields)
+
+			if conditionMatchesValue(condition, value, exists) {
+				matchCount++
+			}
+		}
+		// If this span was matched by every condition, then the rule as a whole
+		// matches (and we can return)
+		if matchCount == len(rule.Condition) {
+			return true
+		}
+	}
+
+	// if the rule didn't match above, then it doesn't match the trace.
+	return false
+}
+
+func extractValueFromSpan(span *types.Span, condition *config.RulesBasedSamplerCondition, checkNestedFields bool) (interface{}, bool) {
+	// whether this condition is matched by this span.
+	value, exists := span.Data[condition.Field]
+	if !exists && checkNestedFields {
+		jsonStr, err := json.Marshal(span.Data)
+		if err == nil {
+			result := gjson.Get(string(jsonStr), condition.Field)
+			if result.Exists() {
+				value = result.String()
+				exists = true
+			}
+		}
+	}
+
+	return value, exists
+}
+
+func conditionMatchesValue(condition *config.RulesBasedSamplerCondition, value interface{}, exists bool) bool {
+	var match bool
+	switch exists {
+	case true:
+		switch condition.Operator {
+		case "exists":
+			match = exists
+		case "!=":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison != equal
+			}
+		case "=":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison == equal
+			}
+		case ">":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison == more
+			}
+		case ">=":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison == more || comparison == equal
+			}
+		case "<":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison == less
+			}
+		case "<=":
+			if comparison, ok := compare(value, condition.Value); ok {
+				match = comparison == less || comparison == equal
+			}
+		case "starts-with":
+			switch a := value.(type) {
+			case string:
+				switch b := condition.Value.(type) {
+				case string:
+					match = strings.HasPrefix(a, b)
+				}
+			}
+		case "contains":
+			switch a := value.(type) {
+			case string:
+				switch b := condition.Value.(type) {
+				case string:
+					match = strings.Contains(a, b)
+				}
+			}
+		case "does-not-contain":
+			switch a := value.(type) {
+			case string:
+				switch b := condition.Value.(type) {
+				case string:
+					match = !strings.Contains(a, b)
+				}
+			}
+		}
+	case false:
+		switch condition.Operator {
+		case "not-exists":
+			match = !exists
+		}
+	}
+	return match
+}
+
 const (
 	less  = -1
 	equal = 0

sample/rules_test.go

@@ -835,3 +835,123 @@ func TestRulesWithEMADynamicSampler(t *testing.T) {
 		}
 	}
 }
+
+func TestRuleMatchesSpanMatchingSpan(t *testing.T) {
+	data := []TestRulesData{
+		{
+			Rules: &config.RulesBasedSamplerConfig{
+				Rule: []*config.RulesBasedSamplerRule{
+					{
+						Name:       "Rule to match span",
+						MatchSpan:  true,
+						SampleRate: 10,
+						Condition: []*config.RulesBasedSamplerCondition{
+							{
+								Field:    "rule_test",
+								Operator: "=",
+								Value:    int64(1),
+							},
+							{
+								Field:    "rule_test_2",
+								Operator: "=",
+								Value:    int64(2),
+							},
+						},
+					},
+				},
+			},
+			Spans: []*types.Span{
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"rule_test":        int64(1),
+							"http.status_code": "200",
+							"rule_test_2":      int64(2),
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"rule_test":        int64(1),
+							"http.status_code": "200",
+						},
+					},
+				},
+			},
+			ExpectedKeep: true,
+			ExpectedRate: 10,
+		},
+		{
+			Rules: &config.RulesBasedSamplerConfig{
+				Rule: []*config.RulesBasedSamplerRule{
+					{
+						Name:      "Rule to match span",
+						MatchSpan: true,
+						Condition: []*config.RulesBasedSamplerCondition{
+							{
+								Field:    "rule_test",
+								Operator: "=",
+								Value:    int64(1),
+							},
+							{
+								Field:    "rule_test_2",
+								Operator: "=",
+								Value:    int64(2),
+							},
+						},
+					},
+					{
+						Name:       "Default rule",
+						Drop:       true,
+						SampleRate: 1,
+					},
+				},
+			},
+			Spans: []*types.Span{
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"rule_test":        int64(1),
+							"http.status_code": "200",
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"rule_test_2":      int64(2),
+							"http.status_code": "200",
+						},
+					},
+				},
+			},
+			ExpectedKeep: false,
+			ExpectedRate: 1,
+		},
+	}
+
+	for _, d := range data {
+		sampler := &RulesBasedSampler{
+			Config:  d.Rules,
+			Logger:  &logger.NullLogger{},
+			Metrics: &metrics.NullMetrics{},
+		}
+
+		trace := &types.Trace{}
+
+		for _, span := range d.Spans {
+			trace.AddSpan(span)
+		}
+
+		sampler.Start()
+		rate, keep := sampler.GetSampleRate(trace)
+
+		assert.Equal(t, d.ExpectedRate, rate, d.Rules)
+
+		// we can only test when we don't expect to keep the trace
+		if !d.ExpectedKeep {
+			assert.Equal(t, d.ExpectedKeep, keep, d.Rules)
+		}
+	}
+}