diff --git a/plugin/sampling/strategystore/static/fixtures/operation_strategies.json b/plugin/sampling/strategystore/static/fixtures/operation_strategies.json new file mode 100644 index 00000000000..f0c7525f357 --- /dev/null +++ b/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 + } + ] + } + ] +} diff --git a/plugin/sampling/strategystore/static/strategy.go b/plugin/sampling/strategystore/static/strategy.go index dcf683f9fe4..f16a41795ec 100644 --- a/plugin/sampling/strategystore/static/strategy.go +++ b/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 } diff --git a/plugin/sampling/strategystore/static/strategy_store.go b/plugin/sampling/strategystore/static/strategy_store.go index c04268c19bd..8ca49bc342e 100644 --- a/plugin/sampling/strategystore/static/strategy_store.go +++ b/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: diff --git a/plugin/sampling/strategystore/static/strategy_store_test.go b/plugin/sampling/strategystore/static/strategy_store_test.go index 4553d3ca7dc..16e05575b44 100644 --- a/plugin/sampling/strategystore/static/strategy_store_test.go +++ b/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 diff --git a/scripts/import-order-cleanup.py b/scripts/import-order-cleanup.py index 13293e68c0b..990a245bef1 100644 --- a/scripts/import-order-cleanup.py +++ b/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)