feat: extract key fields from rules config (#1327)

## Which problem is this PR solving?

Extract sampling fields from config so that we can use it to construct
key fields only spans

fixes: #1325 

## Short description of the changes

- add `GetSamplingFielder` interface to sampler config
- add `GetKeyFields` to all samplers

config/sampler_config.go

@@ -149,10 +149,22 @@ type V2SamplerConfig struct {
 	Samplers     map[string]*V2SamplerChoice `json:"samplers" yaml:"Samplers,omitempty" validate:"required"`
 }
 
+type GetSamplingFielder interface {
+	GetSamplingFields() []string
+}
+
+var _ GetSamplingFielder = (*DeterministicSamplerConfig)(nil)
+
 type DeterministicSamplerConfig struct {
 	SampleRate int `json:"samplerate" yaml:"SampleRate,omitempty" default:"1" validate:"required,gte=1"`
 }
 
+func (d *DeterministicSamplerConfig) GetSamplingFields() []string {
+	return nil
+}
+
+var _ GetSamplingFielder = (*DynamicSamplerConfig)(nil)
+
 type DynamicSamplerConfig struct {
 	SampleRate     int64    `json:"samplerate" yaml:"SampleRate,omitempty" validate:"required,gte=1"`
 	ClearFrequency Duration `json:"clearfrequency" yaml:"ClearFrequency,omitempty"`
@@ -161,6 +173,12 @@ type DynamicSamplerConfig struct {
 	UseTraceLength bool     `json:"usetracelength" yaml:"UseTraceLength,omitempty"`
 }
 
+func (d *DynamicSamplerConfig) GetSamplingFields() []string {
+	return d.FieldList
+}
+
+var _ GetSamplingFielder = (*EMADynamicSamplerConfig)(nil)
+
 type EMADynamicSamplerConfig struct {
 	GoalSampleRate      int      `json:"goalsamplerate" yaml:"GoalSampleRate,omitempty" validate:"gte=1"`
 	AdjustmentInterval  Duration `json:"adjustmentinterval" yaml:"AdjustmentInterval,omitempty"`
@@ -173,6 +191,12 @@ type EMADynamicSamplerConfig struct {
 	UseTraceLength      bool     `json:"usetracelength" yaml:"UseTraceLength,omitempty"`
 }
 
+func (d *EMADynamicSamplerConfig) GetSamplingFields() []string {
+	return d.FieldList
+}
+
+var _ GetSamplingFielder = (*EMAThroughputSamplerConfig)(nil)
+
 type EMAThroughputSamplerConfig struct {
 	GoalThroughputPerSec int      `json:"goalthroughputpersec" yaml:"GoalThroughputPerSec,omitempty"`
 	UseClusterSize       bool     `json:"useclustersize" yaml:"UseClusterSize,omitempty"`
@@ -187,6 +211,12 @@ type EMAThroughputSamplerConfig struct {
 	UseTraceLength       bool     `json:"usetracelength" yaml:"UseTraceLength,omitempty"`
 }
 
+func (d *EMAThroughputSamplerConfig) GetSamplingFields() []string {
+	return d.FieldList
+}
+
+var _ GetSamplingFielder = (*WindowedThroughputSamplerConfig)(nil)
+
 type WindowedThroughputSamplerConfig struct {
 	UpdateFrequency      Duration `json:"updatefrequency" yaml:"UpdateFrequency,omitempty"`
 	LookbackFrequency    Duration `json:"lookbackfrequency" yaml:"LookbackFrequency,omitempty"`
@@ -197,6 +227,12 @@ type WindowedThroughputSamplerConfig struct {
 	UseTraceLength       bool     `json:"usetracelength" yaml:"UseTraceLength,omitempty"`
 }
 
+func (d *WindowedThroughputSamplerConfig) GetSamplingFields() []string {
+	return d.FieldList
+}
+
+var _ GetSamplingFielder = (*TotalThroughputSamplerConfig)(nil)
+
 type TotalThroughputSamplerConfig struct {
 	GoalThroughputPerSec int      `json:"goalthroughputpersec" yaml:"GoalThroughputPerSec,omitempty" validate:"gte=1"`
 	UseClusterSize       bool     `json:"useclustersize" yaml:"UseClusterSize,omitempty"`
@@ -206,12 +242,50 @@ type TotalThroughputSamplerConfig struct {
 	UseTraceLength       bool     `json:"usetracelength" yaml:"UseTraceLength,omitempty"`
 }
 
+func (d *TotalThroughputSamplerConfig) GetSamplingFields() []string {
+	return d.FieldList
+}
+
+var _ GetSamplingFielder = (*RulesBasedSamplerConfig)(nil)
+
 type RulesBasedSamplerConfig struct {
 	// Rules has deliberately different names for json and yaml for conversion from old to new format
 	Rules             []*RulesBasedSamplerRule `json:"rule" yaml:"Rules,omitempty"`
 	CheckNestedFields bool                     `json:"checknestedfields" yaml:"CheckNestedFields,omitempty"`
 }
 
+func (r *RulesBasedSamplerConfig) GetSamplingFields() []string {
+	fields := make(generics.Set[string], 0)
+
+	for _, rule := range r.Rules {
+		if rule == nil {
+			continue
+		}
+
+		for _, condition := range rule.Conditions {
+			// Field and Fields are mutually exclusive, so we only need to check one.
+			if condition.Fields != nil {
+				fields.Add(condition.Fields...)
+				continue
+			}
+
+			if condition.Field != "" {
+				fields.Add(condition.Field)
+			}
+		}
+
+		if rule.Sampler != nil {
+			fields.Add(rule.Sampler.GetSamplingFields()...)
+		}
+	}
+
+	return fields.Members()
+}
+
+var _ GetSamplingFielder = (*RulesBasedDownstreamSampler)(nil)
+
+// RulesBasedDownstreamSampler is a sampler that can be used as a downstream sampler in a rules-based sampler.
+// Only one of the fields should be set.
 type RulesBasedDownstreamSampler struct {
 	DynamicSampler            *DynamicSamplerConfig            `json:"dynamicsampler" yaml:"DynamicSampler,omitempty"`
 	EMADynamicSampler         *EMADynamicSamplerConfig         `json:"emadynamicsampler" yaml:"EMADynamicSampler,omitempty"`
@@ -221,6 +295,35 @@ type RulesBasedDownstreamSampler struct {
 	DeterministicSampler      *DeterministicSamplerConfig      `json:"deterministicsampler" yaml:"DeterministicSampler,omitempty"`
 }
 
+func (r *RulesBasedDownstreamSampler) GetSamplingFields() []string {
+
+	if r.DeterministicSampler != nil {
+		return r.DeterministicSampler.GetSamplingFields()
+	}
+
+	if r.DynamicSampler != nil {
+		return r.DynamicSampler.GetSamplingFields()
+	}
+
+	if r.EMADynamicSampler != nil {
+		return r.EMADynamicSampler.GetSamplingFields()
+	}
+
+	if r.EMAThroughputSampler != nil {
+		return r.EMAThroughputSampler.GetSamplingFields()
+	}
+
+	if r.WindowedThroughputSampler != nil {
+		return r.WindowedThroughputSampler.GetSamplingFields()
+	}
+
+	if r.TotalThroughputSampler != nil {
+		return r.TotalThroughputSampler.GetSamplingFields()
+	}
+
+	return []string{}
+}
+
 type RulesBasedSamplerRule struct {
 	// Conditions has deliberately different names for json and yaml for conversion from old to new format
 	Name       string                        `json:"name" yaml:"Name,omitempty"`

sample/deterministic.go

@@ -57,3 +57,7 @@ func (d *DeterministicSampler) GetSampleRate(trace *types.Trace) (rate uint, kee
 
 	return uint(d.sampleRate), shouldKeep, "deterministic/chance", ""
 }
+
+func (d *DeterministicSampler) GetKeyFields() []string {
+	return d.Config.GetSamplingFields()
+}

sample/dynamic.go

@@ -23,7 +23,8 @@ type DynamicSampler struct {
 	prefix         string
 	lastMetrics    map[string]int64
 
-	key *traceKey
+	key       *traceKey
+	keyFields []string
 
 	dynsampler dynsampler.Sampler
 }
@@ -42,6 +43,7 @@ func (d *DynamicSampler) Start() error {
 		d.maxKeys = 500
 	}
 	d.prefix = "dynamic_"
+	d.keyFields = d.Config.GetSamplingFields()
 
 	// spin up the actual dynamic sampler
 	d.dynsampler = &dynsampler.AvgSampleRate{
@@ -96,3 +98,7 @@ func (d *DynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep bool
 	}
 	return rate, shouldKeep, "dynamic", key
 }
+
+func (d *DynamicSampler) GetKeyFields() []string {
+	return d.keyFields
+}

sample/dynamic_ema.go

@@ -27,7 +27,8 @@ type EMADynamicSampler struct {
 	prefix              string
 	lastMetrics         map[string]int64
 
-	key *traceKey
+	key       *traceKey
+	keyFields []string
 
 	dynsampler dynsampler.Sampler
 }
@@ -47,6 +48,7 @@ func (d *EMADynamicSampler) Start() error {
 		d.maxKeys = 500
 	}
 	d.prefix = "emadynamic_"
+	d.keyFields = d.Config.GetSamplingFields()
 
 	// spin up the actual dynamic sampler
 	d.dynsampler = &dynsampler.EMASampleRate{
@@ -104,3 +106,7 @@ func (d *EMADynamicSampler) GetSampleRate(trace *types.Trace) (rate uint, keep b
 	}
 	return rate, shouldKeep, "emadynamic", key
 }
+
+func (d *EMADynamicSampler) GetKeyFields() []string {
+	return d.keyFields
+}

sample/ema_throughput.go

@@ -30,7 +30,8 @@ type EMAThroughputSampler struct {
 	prefix               string
 	lastMetrics          map[string]int64
 
-	key *traceKey
+	key       *traceKey
+	keyFields []string
 
 	dynsampler *dynsampler.EMAThroughput
 }
@@ -56,6 +57,7 @@ func (d *EMAThroughputSampler) Start() error {
 	}
 	d.prefix = "emathroughput_"
 
+	d.keyFields = d.Config.GetSamplingFields()
 	// spin up the actual dynamic sampler
 	d.dynsampler = &dynsampler.EMAThroughput{
 		GoalThroughputPerSec: d.goalThroughputPerSec / d.clusterSize,
@@ -121,3 +123,7 @@ func (d *EMAThroughputSampler) GetSampleRate(trace *types.Trace) (rate uint, kee
 	}
 	return rate, shouldKeep, "emathroughput", key
 }
+
+func (d *EMAThroughputSampler) GetKeyFields() []string {
+	return d.keyFields
+}

sample/rules.go

@@ -15,11 +15,12 @@ import (
 var _ ClusterSizer = (*RulesBasedSampler)(nil)
 
 type RulesBasedSampler struct {
-	Config   *config.RulesBasedSamplerConfig
-	Logger   logger.Logger
-	Metrics  metrics.Metrics
-	samplers map[string]Sampler
-	prefix   string
+	Config    *config.RulesBasedSamplerConfig
+	Logger    logger.Logger
+	Metrics   metrics.Metrics
+	samplers  map[string]Sampler
+	prefix    string
+	keyFields []string
 }
 
 const RootPrefix = "root."
@@ -35,6 +36,7 @@ func (s *RulesBasedSampler) Start() error {
 	s.Metrics.Register(s.prefix+"sample_rate", "histogram")
 
 	s.samplers = make(map[string]Sampler)
+	s.keyFields = s.Config.GetSamplingFields()
 
 	for _, rule := range s.Config.Rules {
 		for _, cond := range rule.Conditions {
@@ -162,6 +164,10 @@ func (s *RulesBasedSampler) GetSampleRate(trace *types.Trace) (rate uint, keep b
 	return 1, true, "no rule matched", ""
 }
 
+func (s *RulesBasedSampler) GetKeyFields() []string {
+	return s.keyFields
+}
+
 func ruleMatchesTrace(t *types.Trace, rule *config.RulesBasedSamplerRule, checkNestedFields bool) bool {
 	// We treat a rule with no conditions as a match.
 	if rule.Conditions == nil {