Add ability to adjust static sampling probabilities per operation (#827)

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

@@ -0,0 +1,42 @@
+{
+  "default_strategy": {
+    "type": "probabilistic",
+    "param": 0.5
+  },
+  "service_strategies": [
+    {
+      "service": "foo",
+      "type": "probabilistic",
+      "param": 0.8,
+      "operation_strategies": [
+        {
+          "operation": "op1",
+          "type": "probabilistic",
+          "param": 0.2
+        },
+        {
+          "operation": "op2",
+          "type": "ratelimiting",
+          "param": 10
+        }
+      ]
+    },
+    {
+      "service": "bar",
+      "type": "ratelimiting",
+      "param": 5,
+      "operation_strategies": [
+        {
+          "operation": "op3",
+          "type": "probabilistic",
+          "param": 0.3
+        },
+        {
+          "operation": "op4",
+          "type": "ratelimiting",
+          "param": 100
+        }
+      ]
+    }
+  ]
+}

plugin/sampling/strategystore/static/strategy.go

@@ -21,9 +21,16 @@ type strategy struct {
 	Param float64 `json:"param"`
 }
 
+// operationStrategy defines an operation specific sampling strategy.
+type operationStrategy struct {
+	Operation string `json:"operation"`
+	strategy
+}
+
 // serviceStrategy defines a service specific sampling strategy.
 type serviceStrategy struct {
-	Service string `json:"service"`
+	Service             string               `json:"service"`
+	OperationStrategies []*operationStrategy `json:"operation_strategies"`
 	strategy
 }
 

plugin/sampling/strategystore/static/strategy_store.go

@@ -16,6 +16,7 @@ package static
 
 import (
 	"encoding/json"
+	"fmt"
 	"io/ioutil"
 
 	"github.com/pkg/errors"
@@ -80,10 +81,43 @@ func (h *strategyStore) parseStrategies(strategies *strategies) {
 		h.defaultStrategy = h.parseStrategy(strategies.DefaultStrategy)
 	}
 	for _, s := range strategies.ServiceStrategies {
-		h.serviceStrategies[s.Service] = h.parseStrategy(&s.strategy)
+		h.serviceStrategies[s.Service] = h.parseServiceStrategies(s)
 	}
 }
 
+func (h *strategyStore) parseServiceStrategies(strategy *serviceStrategy) *sampling.SamplingStrategyResponse {
+	resp := h.parseStrategy(&strategy.strategy)
+	if len(strategy.OperationStrategies) == 0 {
+		return resp
+	}
+	opS := &sampling.PerOperationSamplingStrategies{
+		DefaultSamplingProbability: defaultSamplingProbability,
+	}
+	if resp.StrategyType == sampling.SamplingStrategyType_PROBABILISTIC {
+		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))
+			continue
+		}
+		opS.PerOperationStrategies = append(opS.PerOperationStrategies,
+			&sampling.OperationSamplingStrategy{
+				Operation:             operationStrategy.Operation,
+				ProbabilisticSampling: s.ProbabilisticSampling,
+			})
+	}
+	resp.OperationSampling = opS
+	return resp
+}
+
 func (h *strategyStore) parseStrategy(strategy *strategy) *sampling.SamplingStrategyResponse {
 	switch strategy.Type {
 	case samplerTypeProbabilistic:

plugin/sampling/strategystore/static/strategy_store_test.go

@@ -58,6 +58,48 @@ func TestStrategyStore(t *testing.T) {
 	assert.EqualValues(t, makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.5), *s)
 }
 
+func TestPerOperationSamplingStrategies(t *testing.T) {
+	logger, buf := testutils.NewLogger()
+	store, err := NewStrategyStore(Options{StrategiesFile: "fixtures/operation_strategies.json"}, logger)
+	assert.Contains(t, buf.String(), "Operation strategies only supports probabilistic sampling at the moment,"+
+		"'op2' defaulting to probabilistic sampling with probability 0.8")
+	assert.Contains(t, buf.String(), "Operation strategies only supports probabilistic sampling at the moment,"+
+		"'op4' defaulting to probabilistic sampling with probability 0.001")
+	require.NoError(t, err)
+
+	expected := makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.8)
+
+	s, err := store.GetSamplingStrategy("foo")
+	require.NoError(t, err)
+	assert.Equal(t, sampling.SamplingStrategyType_PROBABILISTIC, s.StrategyType)
+	assert.Equal(t, *expected.ProbabilisticSampling, *s.ProbabilisticSampling)
+
+	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)
+
+	expected = makeResponse(sampling.SamplingStrategyType_RATE_LIMITING, 5)
+
+	s, err = store.GetSamplingStrategy("bar")
+	require.NoError(t, err)
+	assert.Equal(t, sampling.SamplingStrategyType_RATE_LIMITING, s.StrategyType)
+	assert.Equal(t, *expected.RateLimitingSampling, *s.RateLimitingSampling)
+
+	require.NotNil(t, s.OperationSampling)
+	os = s.OperationSampling
+	assert.EqualValues(t, os.DefaultSamplingProbability, 0.001)
+	require.Len(t, os.PerOperationStrategies, 1)
+	assert.Equal(t, "op3", os.PerOperationStrategies[0].Operation)
+	assert.EqualValues(t, 0.3, os.PerOperationStrategies[0].ProbabilisticSampling.SamplingRate)
+
+	s, err = store.GetSamplingStrategy("default")
+	require.NoError(t, err)
+	assert.EqualValues(t, makeResponse(sampling.SamplingStrategyType_PROBABILISTIC, 0.5), *s)
+}
+
 func TestParseStrategy(t *testing.T) {
 	tests := []struct {
 		strategy serviceStrategy

scripts/import-order-cleanup.py

@@ -81,7 +81,7 @@ def main():
     for f in go_files:
         parsed, imports_reordered = parse_go_file(f)
         if output == "stdout" and imports_reordered:
-            print f + " imports out of order"
+            print(f + " imports out of order")
         else:
             with open(f, 'w') as ofile:
                 ofile.write(parsed)