internal/sampler: add defaults for each service

cmd/trace-agent/agent.go

@@ -56,7 +56,7 @@ type Agent struct {
 // NewAgent returns a new Agent object, ready to be started. It takes a context
 // which may be cancelled in order to gracefully stop the agent.
 func NewAgent(ctx context.Context, conf *config.AgentConfig) *Agent {
-	dynConf := sampler.NewDynamicConfig()
+	dynConf := sampler.NewDynamicConfig(conf.DefaultEnv)
 
 	// inter-component channels
 	rawTraceChan := make(chan agent.Trace, 5000) // about 1000 traces/sec for 5 sec, TODO: move to *model.Trace

internal/api/api_test.go

@@ -34,7 +34,7 @@ var headerFields = map[string]string{
 }
 
 func newTestReceiverFromConfig(conf *config.AgentConfig) *HTTPReceiver {
-	dynConf := sampler.NewDynamicConfig()
+	dynConf := sampler.NewDynamicConfig("none")
 
 	rawTraceChan := make(chan agent.Trace, 5000)
 	serviceChan := make(chan agent.ServicesMetadata, 50)

internal/sampler/dynamic_config.go

@@ -12,25 +12,27 @@ type DynamicConfig struct {
 	RateByService RateByService
 }
 
-// NewDynamicConfig creates a new dynamic config object.
-func NewDynamicConfig() *DynamicConfig {
-	// Not much logic here now, as RateByService is fine with
-	// being used unintialized, but external packages should use this.
-	return &DynamicConfig{}
+// NewDynamicConfig creates a new dynamic config object which maps service signatures
+// to their corresponding sampling rates. Each service will have a default assigned
+// matching the service rate of the specified env.
+func NewDynamicConfig(env string) *DynamicConfig {
+	return &DynamicConfig{RateByService: RateByService{defaultEnv: env}}
 }
 
 // RateByService stores the sampling rate per service. It is thread-safe, so
 // one can read/write on it concurrently, using getters and setters.
 type RateByService struct {
+	defaultEnv string // env. to use for service defaults
+
+	mu    sync.RWMutex // guards rates
 	rates map[ServiceSignature]float64
-	mutex sync.RWMutex
 }
 
 // SetAll the sampling rate for all services. If a service/env is not
 // in the map, then the entry is removed.
 func (rbs *RateByService) SetAll(rates map[ServiceSignature]float64) {
-	rbs.mutex.Lock()
-	defer rbs.mutex.Unlock()
+	rbs.mu.Lock()
+	defer rbs.mu.Unlock()
 
 	if rbs.rates == nil {
 		rbs.rates = make(map[ServiceSignature]float64, len(rates))
@@ -49,13 +51,18 @@ func (rbs *RateByService) SetAll(rates map[ServiceSignature]float64) {
 			v = 1
 		}
 		rbs.rates[k] = v
+		if k.Env == rbs.defaultEnv {
+			// if this is the default env, then this is also the
+			// service's default rate unbound to any env.
+			rbs.rates[ServiceSignature{Name: k.Name}] = v
+		}
 	}
 }
 
 // GetAll returns all sampling rates for all services.
 func (rbs *RateByService) GetAll() map[string]float64 {
-	rbs.mutex.RLock()
-	defer rbs.mutex.RUnlock()
+	rbs.mu.RLock()
+	defer rbs.mu.RUnlock()
 
 	ret := make(map[string]float64, len(rbs.rates))
 	for k, v := range rbs.rates {

internal/sampler/dynamic_config_test.go

@@ -11,7 +11,7 @@ import (
 func TestNewDynamicConfig(t *testing.T) {
 	assert := assert.New(t)
 
-	dc := NewDynamicConfig()
+	dc := NewDynamicConfig("none")
 	assert.NotNil(dc)
 
 	rates := map[ServiceSignature]float64{
@@ -87,6 +87,19 @@ func TestRateByServiceLimits(t *testing.T) {
 	assert.Equal(map[string]float64{"service:high,env:": 1, "service:low,env:": 0}, rbc.GetAll())
 }
 
+func TestRateByServiceDefaults(t *testing.T) {
+	rbc := RateByService{defaultEnv: "test"}
+	rbc.SetAll(map[ServiceSignature]float64{
+		ServiceSignature{"one", "prod"}: 0.5,
+		ServiceSignature{"two", "test"}: 0.4,
+	})
+	assert.Equal(t, map[string]float64{
+		"service:one,env:prod": 0.5,
+		"service:two,env:test": 0.4,
+		"service:two,env:":     0.4,
+	}, rbc.GetAll())
+}
+
 func TestRateByServiceConcurrency(t *testing.T) {
 	assert := assert.New(t)