Make operation_strategies part also be part of default_strategy

plugin/sampling/strategystore/static/constants.go

@@ -32,13 +32,13 @@ const (
 	defaultSamplingProbability = 0.001
 )
 
-var (
-	// defaultStrategy is the default sampling strategy the Strategy Store will return
-	// if none is provided.
-	defaultStrategy = sampling.SamplingStrategyResponse{
+// defaultStrategy is the default sampling strategy the Strategy Store will return
+// if none is provided.
+func defaultStrategyResponse() *sampling.SamplingStrategyResponse {
+	return &sampling.SamplingStrategyResponse{
 		StrategyType: sampling.SamplingStrategyType_PROBABILISTIC,
 		ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{
 			SamplingRate: defaultSamplingProbability,
 		},
 	}
-)
+}

plugin/sampling/strategystore/static/fixtures/operation_strategies.json

@@ -1,14 +1,41 @@
 {
   "default_strategy": {
     "type": "probabilistic",
-    "param": 0.5
+    "param": 0.5,
+    "operation_strategies": [
+      {
+        "operation": "op0",
+        "type": "probabilistic",
+        "param": 0.2
+      },
+      {
+        "operation": "op6",
+        "type": "probabilistic",
+        "param": 0
+      },
+      {
+        "operation": "spam",
+        "type": "ratelimiting",
+        "param": 1
+      },
+      {
+        "operation": "op7",
+        "type": "probabilistic",
+        "param": 1
+      }
+    ]
   },
   "service_strategies": [
     {
       "service": "foo",
       "type": "probabilistic",
       "param": 0.8,
       "operation_strategies": [
+        {
+          "operation": "op6",
+          "type": "probabilistic",
+          "param": 0.5
+        },
         {
           "operation": "op1",
           "type": "probabilistic",

plugin/sampling/strategystore/static/strategy.go

@@ -36,6 +36,6 @@ type serviceStrategy struct {
 
 // strategies holds a default sampling strategy and service specific sampling strategies.
 type strategies struct {
-	DefaultStrategy   *strategy          `json:"default_strategy"`
+	DefaultStrategy   *serviceStrategy   `json:"default_strategy"`
 	ServiceStrategies []*serviceStrategy `json:"service_strategies"`
 }

plugin/sampling/strategystore/static/strategy_store.go

@@ -74,17 +74,57 @@ func loadStrategies(strategiesFile string) (*strategies, error) {
 }
 
 func (h *strategyStore) parseStrategies(strategies *strategies) {
-	h.defaultStrategy = &defaultStrategy
+	h.defaultStrategy = defaultStrategyResponse()
 	if strategies == nil {
 		h.logger.Info("No sampling strategies provided, using defaults")
 		return
 	}
 	if strategies.DefaultStrategy != nil {
-		h.defaultStrategy = h.parseStrategy(strategies.DefaultStrategy)
+		h.defaultStrategy = h.parseServiceStrategies(strategies.DefaultStrategy)
 	}
+
+	merge := true
+	if h.defaultStrategy.OperationSampling == nil ||
+		h.defaultStrategy.OperationSampling.PerOperationStrategies == nil {
+		merge = false
+	}
+
 	for _, s := range strategies.ServiceStrategies {
 		h.serviceStrategies[s.Service] = h.parseServiceStrategies(s)
+
+		// Merge with the default operation strategies, because only merging with
+		// the default strategy has no effect on service strategies (the default strategy
+		// is not merged with and only used as a fallback).
+		opS := h.serviceStrategies[s.Service].OperationSampling
+		if opS == nil {
+			// Service has no per-operation strategies, so just reference the default settings.
+			h.serviceStrategies[s.Service].OperationSampling = h.defaultStrategy.OperationSampling
+			continue
+		}
+
+		if merge {
+			opS.PerOperationStrategies = mergePerOperationSamplingStrategies(
+				opS.PerOperationStrategies,
+				h.defaultStrategy.OperationSampling.PerOperationStrategies)
+		}
+	}
+}
+
+// mergePerOperationStrategies merges two operation strategies a and b, where a takes precedence over b.
+func mergePerOperationSamplingStrategies(
+	a, b []*sampling.OperationSamplingStrategy,
+) []*sampling.OperationSamplingStrategy {
+	m := make(map[string]bool)
+	for _, aOp := range a {
+		m[aOp.Operation] = true
+	}
+	for _, bOp := range b {
+		if m[bOp.Operation] {
+			continue
+		}
+		a = append(a, bOp)
 	}
+	return a
 }
 
 func (h *strategyStore) parseServiceStrategies(strategy *serviceStrategy) *sampling.SamplingStrategyResponse {
@@ -99,17 +139,11 @@ func (h *strategyStore) parseServiceStrategies(strategy *serviceStrategy) *sampl
 		opS.DefaultSamplingProbability = resp.ProbabilisticSampling.SamplingRate
 	}
 	for _, operationStrategy := range strategy.OperationStrategies {
-		s := h.parseStrategy(&operationStrategy.strategy)
-		if s.StrategyType == sampling.SamplingStrategyType_RATE_LIMITING {
-			// TODO OperationSamplingStrategy only supports probabilistic sampling
-			h.logger.Warn(
-				fmt.Sprintf(
-					"Operation strategies only supports probabilistic sampling at the moment,"+
-						"'%s' defaulting to probabilistic sampling with probability %f",
-					operationStrategy.Operation, opS.DefaultSamplingProbability),
-				zap.Any("strategy", operationStrategy))
+		s, ok := h.parseOperationStrategy(operationStrategy, opS)
+		if !ok {
 			continue
 		}
+
 		opS.PerOperationStrategies = append(opS.PerOperationStrategies,
 			&sampling.OperationSamplingStrategy{
 				Operation:             operationStrategy.Operation,
@@ -120,6 +154,24 @@ func (h *strategyStore) parseServiceStrategies(strategy *serviceStrategy) *sampl
 	return resp
 }
 
+func (h *strategyStore) parseOperationStrategy(
+	strategy *operationStrategy,
+	parent *sampling.PerOperationSamplingStrategies,
+) (s *sampling.SamplingStrategyResponse, ok bool) {
+	s = h.parseStrategy(&strategy.strategy)
+	if s.StrategyType == sampling.SamplingStrategyType_RATE_LIMITING {
+		// TODO OperationSamplingStrategy only supports probabilistic sampling
+		h.logger.Warn(
+			fmt.Sprintf(
+				"Operation strategies only supports probabilistic sampling at the moment,"+
+					"'%s' defaulting to probabilistic sampling with probability %f",
+				strategy.Operation, parent.DefaultSamplingProbability),
+			zap.Any("strategy", strategy))
+		return nil, false
+	}
+	return s, true
+}
+
 func (h *strategyStore) parseStrategy(strategy *strategy) *sampling.SamplingStrategyResponse {
 	switch strategy.Type {
 	case samplerTypeProbabilistic:
@@ -138,7 +190,7 @@ func (h *strategyStore) parseStrategy(strategy *strategy) *sampling.SamplingStra
 		}
 	default:
 		h.logger.Warn("Failed to parse sampling strategy", zap.Any("strategy", strategy))
-		return deepCopy(&defaultStrategy)
+		return defaultStrategyResponse()
 	}
 }
 

plugin/sampling/strategystore/static/strategy_store_test.go

@@ -15,6 +15,7 @@
 package static
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -77,9 +78,16 @@ func TestPerOperationSamplingStrategies(t *testing.T) {
 	require.NotNil(t, s.OperationSampling)
 	os := s.OperationSampling
 	assert.EqualValues(t, os.DefaultSamplingProbability, 0.8)
-	require.Len(t, os.PerOperationStrategies, 1)
-	assert.Equal(t, "op1", os.PerOperationStrategies[0].Operation)
-	assert.EqualValues(t, 0.2, os.PerOperationStrategies[0].ProbabilisticSampling.SamplingRate)
+	require.Len(t, os.PerOperationStrategies, 4)
+	fmt.Println(os)
+	assert.Equal(t, "op6", os.PerOperationStrategies[0].Operation)
+	assert.EqualValues(t, 0.5, os.PerOperationStrategies[0].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op1", os.PerOperationStrategies[1].Operation)
+	assert.EqualValues(t, 0.2, os.PerOperationStrategies[1].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op0", os.PerOperationStrategies[2].Operation)
+	assert.EqualValues(t, 0.2, os.PerOperationStrategies[2].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op7", os.PerOperationStrategies[3].Operation)
+	assert.EqualValues(t, 1, os.PerOperationStrategies[3].ProbabilisticSampling.SamplingRate)
 
 	expected = makeResponse(sampling.SamplingStrategyType_RATE_LIMITING, 5)
 
@@ -91,15 +99,45 @@ func TestPerOperationSamplingStrategies(t *testing.T) {
 	require.NotNil(t, s.OperationSampling)
 	os = s.OperationSampling
 	assert.EqualValues(t, os.DefaultSamplingProbability, 0.001)
-	require.Len(t, os.PerOperationStrategies, 2)
+	require.Len(t, os.PerOperationStrategies, 5)
 	assert.Equal(t, "op3", os.PerOperationStrategies[0].Operation)
 	assert.EqualValues(t, 0.3, os.PerOperationStrategies[0].ProbabilisticSampling.SamplingRate)
 	assert.Equal(t, "op5", os.PerOperationStrategies[1].Operation)
 	assert.EqualValues(t, 0.4, os.PerOperationStrategies[1].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op0", os.PerOperationStrategies[2].Operation)
+	assert.EqualValues(t, 0.2, os.PerOperationStrategies[2].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op6", os.PerOperationStrategies[3].Operation)
+	assert.EqualValues(t, 0, os.PerOperationStrategies[3].ProbabilisticSampling.SamplingRate)
+	assert.Equal(t, "op7", os.PerOperationStrategies[4].Operation)
+	assert.EqualValues(t, 1, os.PerOperationStrategies[4].ProbabilisticSampling.SamplingRate)
 
 	s, err = store.GetSamplingStrategy("default")
 	require.NoError(t, err)
-	assert.EqualValues(t, makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.5), *s)
+	expectedRsp := makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.5)
+	expectedRsp.OperationSampling = &sampling.PerOperationSamplingStrategies{
+		DefaultSamplingProbability: 0.5,
+		PerOperationStrategies: []*sampling.OperationSamplingStrategy{
+			{
+				Operation: "op0",
+				ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{
+					SamplingRate: 0.2,
+				},
+			},
+			{
+				Operation: "op6",
+				ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{
+					SamplingRate: 0,
+				},
+			},
+			{
+				Operation: "op7",
+				ProbabilisticSampling: &sampling.ProbabilisticSamplingStrategy{
+					SamplingRate: 1,
+				},
+			},
+		},
+	}
+	assert.EqualValues(t, expectedRsp, *s)
 }
 
 func TestMissingServiceSamplingStrategyTypes(t *testing.T) {