feat: use a computed field for current descendant count in rules

config/sampler_config.go

@@ -30,6 +30,16 @@ const (
 	NotIn          = "not-in"
 )
 
+// ComputedField is a virtual field. It's value is calculated during rule evaluation.
+// We use the `?.` prefix to distinguish computed fields from regular fields.
+type ComputedField string
+
+const (
+	// ComputedFieldPrefix is the prefix for computed fields.
+	ComputedFieldPrefix               = "?."
+	NUM_DESCENDANTS     ComputedField = ComputedFieldPrefix + "NUM_DESCENDANTS"
+)
+
 // The json tags in this file are used for conversion from the old format (see tools/convert for details).
 // They are deliberately all lowercase.
 // The yaml tags are used for the new format and are PascalCase.
@@ -250,6 +260,14 @@ func (r *RulesBasedSamplerCondition) String() string {
 	return fmt.Sprintf("%+v", *r)
 }
 
+func (r *RulesBasedSamplerCondition) GetComputedField() (ComputedField, bool) {
+	if strings.HasPrefix(r.Field, ComputedFieldPrefix) {
+		return ComputedField(r.Field), true
+	}
+	return "", false
+
+}
+
 func (r *RulesBasedSamplerCondition) setMatchesFunction() error {
 	switch r.Operator {
 	case Exists:

sample/rules.go

@@ -167,7 +167,7 @@ func ruleMatchesTrace(t *types.Trace, rule *config.RulesBasedSamplerRule, checkN
 
 	span:
 		for _, span := range t.GetSpans() {
-			value, exists := extractValueFromSpan(span, condition, checkNestedFields)
+			value, exists := extractValueFromSpan(t, span, condition, checkNestedFields)
 			if condition.Matches == nil {
 				if conditionMatchesValue(condition, value, exists) {
 					matched++
@@ -194,7 +194,7 @@ func ruleMatchesSpanInTrace(trace *types.Trace, rule *config.RulesBasedSamplerRu
 		ruleMatched := true
 		for _, condition := range rule.Conditions {
 			// whether this condition is matched by this span.
-			value, exists := extractValueFromSpan(span, condition, checkNestedFields)
+			value, exists := extractValueFromSpan(trace, span, condition, checkNestedFields)
 			if condition.Matches == nil {
 				if !conditionMatchesValue(condition, value, exists) {
 					ruleMatched = false
@@ -220,7 +220,15 @@ func ruleMatchesSpanInTrace(trace *types.Trace, rule *config.RulesBasedSamplerRu
 	return false
 }
 
-func extractValueFromSpan(span *types.Span, condition *config.RulesBasedSamplerCondition, checkNestedFields bool) (interface{}, bool) {
+func extractValueFromSpan(trace *types.Trace, span *types.Span, condition *config.RulesBasedSamplerCondition, checkNestedFields bool) (interface{}, bool) {
+	// If the condition is a descendant count, we extract the count from trace and return it.
+	if f, ok := condition.GetComputedField(); ok {
+		switch f {
+		case config.NUM_DESCENDANTS:
+			return int64(trace.DescendantCount()), true
+		}
+	}
+
 	// whether this condition is matched by this span.
 	var value any
 	var exists bool

sample/rules_test.go

@@ -772,6 +772,126 @@ func TestRules(t *testing.T) {
 			ExpectedName: "no rule matched",
 			ExpectedRate: 1,
 		},
+		{
+			Rules: &config.RulesBasedSamplerConfig{
+				Rules: []*config.RulesBasedSamplerRule{
+					{
+						Name:       "Check that the number of descendants is greater than 3",
+						SampleRate: 1,
+						Conditions: []*config.RulesBasedSamplerCondition{
+							{
+								Field:    string(config.NUM_DESCENDANTS),
+								Operator: config.GT,
+								Value:    int(3),
+								Datatype: "int",
+							},
+						},
+						Drop: true,
+					},
+				},
+			},
+			Spans: []*types.Span{
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "54322",
+							"trace.parent_id": "54321",
+							"meta.span_count": int64(2),
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "654321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "754321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "754321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+			},
+			ExpectedName: "Check that the number of descendants is greater than 3",
+			ExpectedKeep: false,
+			ExpectedRate: 1,
+		},
+		{
+			Rules: &config.RulesBasedSamplerConfig{
+				Rules: []*config.RulesBasedSamplerRule{
+					{
+						Name:       "Check that the number of descendants is less than 3",
+						SampleRate: 1,
+						Conditions: []*config.RulesBasedSamplerCondition{
+							{
+								Field:    string(config.NUM_DESCENDANTS),
+								Operator: config.LT,
+								Value:    int(3),
+							},
+						},
+					},
+				},
+			},
+			Spans: []*types.Span{
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "54322",
+							"trace.parent_id": "54321",
+							"meta.span_count": int64(2),
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "654321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "754321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+				{
+					Event: types.Event{
+						Data: map[string]interface{}{
+							"trace.trace_id":  "12345",
+							"trace.span_id":   "754321",
+							"trace.parent_id": "54322",
+						},
+					},
+				},
+			},
+			ExpectedName: "no rule matched",
+			ExpectedKeep: true,
+			ExpectedRate: 1,
+		},
 	}
 
 	for _, d := range data {