test changes

Author: nhulston

internal/trace/listener.go

@@ -13,6 +13,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"os"
+	"regexp"
 	"strings"
 
 	"github.com/DataDog/datadog-lambda-go/internal/extension"
@@ -91,6 +92,58 @@ func (l *Listener) HandlerStarted(ctx context.Context, msg json.RawMessage) cont
 			tracer.WithGlobalTag("_dd.origin", "lambda"),
 			tracer.WithSendRetries(2),
 		}, l.tracerOptions...)
+
+		// Begin parsing DD_TRACE_SAMPLING_RULES
+		if samplingRulesEnv := os.Getenv("DD_TRACE_SAMPLING_RULES"); samplingRulesEnv != "" {
+			var rawRules []struct {
+				Service    string  `json:"service,omitempty"`
+				Name       string  `json:"name,omitempty"`
+				SampleRate float64 `json:"sample_rate"`
+			}
+			err := json.Unmarshal([]byte(samplingRulesEnv), &rawRules)
+			if err != nil {
+				logger.Error(fmt.Errorf("Error parsing sampling rules from DD_TRACE_SAMPLING_RULES: %v", err))
+			} else {
+				var samplingRules []tracer.SamplingRule
+				for _, rule := range rawRules {
+					// Validate the sample rate
+					if rule.SampleRate < 0 || rule.SampleRate > 1 {
+						logger.Error(fmt.Errorf("Invalid sample rate %v in sampling rules", rule.SampleRate))
+						continue
+					}
+
+					// Compile regex patterns for Service and Name if provided
+					var serviceRegex, nameRegex *regexp.Regexp
+					if rule.Service != "" {
+						serviceRegex, err = regexp.Compile(rule.Service)
+						if err != nil {
+							logger.Error(fmt.Errorf("Invalid service regex %v: %v", rule.Service, err))
+							continue
+						}
+					}
+					if rule.Name != "" {
+						nameRegex, err = regexp.Compile(rule.Name)
+						if err != nil {
+							logger.Error(fmt.Errorf("Invalid name regex %v: %v", rule.Name, err))
+							continue
+						}
+					}
+
+					// Create the sampling rule
+					samplingRule := tracer.SamplingRule{
+						Service: serviceRegex,
+						Name:    nameRegex,
+						Rate:    rule.SampleRate,
+					}
+					samplingRules = append(samplingRules, samplingRule)
+				}
+				if len(samplingRules) > 0 {
+					opts = append(opts, tracer.WithSamplingRules(samplingRules...))
+				}
+			}
+		}
+		// End parsing DD_TRACE_SAMPLING_RULES
+
 		if l.otelTracerEnabled {
 			provider := ddotel.NewTracerProvider(
 				opts...,