Add sampling handler and sampling strategy store (#674)

Author: black-adder

cmd/collector/app/sampling/handler.go

@@ -0,0 +1,44 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sampling
+
+import (
+	"github.com/uber/tchannel-go/thrift"
+
+	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
+	"github.com/jaegertracing/jaeger/thrift-gen/sampling"
+)
+
+// Handler returns sampling strategies for specified services
+type Handler interface {
+	// GetSamplingStrategy returns recommended sampling strategy for a given service name.
+	GetSamplingStrategy(ctx thrift.Context, serviceName string) (*sampling.SamplingStrategyResponse, error)
+}
+
+type handler struct {
+	store strategystore.StrategyStore
+}
+
+// NewHandler creates a handler that controls sampling strategies for services.
+func NewHandler(store strategystore.StrategyStore) Handler {
+	return &handler{
+		store: store,
+	}
+}
+
+// GetSamplingStrategy returns allowed sampling strategy for a given service name.
+func (h *handler) GetSamplingStrategy(ctx thrift.Context, serviceName string) (*sampling.SamplingStrategyResponse, error) {
+	return h.store.GetSamplingStrategy(serviceName)
+}

cmd/collector/app/sampling/handler_test.go

@@ -0,0 +1,35 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sampling
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+
+	"github.com/jaegertracing/jaeger/thrift-gen/sampling"
+)
+
+func TestHandler(t *testing.T) {
+	handler := NewHandler(mockStore{})
+	_, err := handler.GetSamplingStrategy(nil, "")
+	assert.NoError(t, err)
+}
+
+type mockStore struct{}
+
+func (s mockStore) GetSamplingStrategy(serviceName string) (*sampling.SamplingStrategyResponse, error) {
+	return nil, nil
+}

cmd/collector/app/sampling/strategystore/factory.go

@@ -0,0 +1,34 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package strategystore
+
+import (
+	"github.com/uber/jaeger-lib/metrics"
+	"go.uber.org/zap"
+)
+
+// Factory defines an interface for a factory that can create implementations of different strategy storage components.
+// Implementations are also encouraged to implement plugin.Configurable interface.
+//
+// See also
+//
+// plugin.Configurable
+type Factory interface {
+	// Initialize performs internal initialization of the factory.
+	Initialize(metricsFactory metrics.Factory, logger *zap.Logger) error
+
+	// CreateStrategyStore initializes the StrategyStore and returns it.
+	CreateStrategyStore() (StrategyStore, error)
+}

cmd/collector/app/sampling/strategystore/interface.go

@@ -0,0 +1,25 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package strategystore
+
+import (
+	"github.com/jaegertracing/jaeger/thrift-gen/sampling"
+)
+
+// StrategyStore keeps track of service specific sampling strategies.
+type StrategyStore interface {
+	// GetSamplingStrategy retrieves the sampling strategy for the specified service.
+	GetSamplingStrategy(serviceName string) (*sampling.SamplingStrategyResponse, error)
+}

plugin/sampling/strategystore/factory.go

@@ -0,0 +1,105 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package strategystore
+
+import (
+	"flag"
+	"fmt"
+
+	"github.com/spf13/viper"
+	"github.com/uber/jaeger-lib/metrics"
+	"go.uber.org/zap"
+
+	"github.com/jaegertracing/jaeger/cmd/collector/app/sampling/strategystore"
+	"github.com/jaegertracing/jaeger/plugin"
+	"github.com/jaegertracing/jaeger/plugin/sampling/strategystore/static"
+)
+
+const (
+	staticStrategyStoreType   = "static"
+	adaptiveStrategyStoreType = "adaptive"
+)
+
+var allSamplingTypes = []string{staticStrategyStoreType} // TODO support adaptive
+
+// Factory implements strategystore.Factory interface as a meta-factory for strategy storage components.
+type Factory struct {
+	FactoryConfig
+
+	factories map[string]strategystore.Factory
+}
+
+// NewFactory creates the meta-factory.
+func NewFactory(config FactoryConfig) (*Factory, error) {
+	f := &Factory{FactoryConfig: config}
+	uniqueTypes := map[string]struct{}{
+		f.StrategyStoreType: {},
+	}
+	f.factories = make(map[string]strategystore.Factory)
+	for t := range uniqueTypes {
+		ff, err := f.getFactoryOfType(t)
+		if err != nil {
+			return nil, err
+		}
+		f.factories[t] = ff
+	}
+	return f, nil
+}
+
+func (f *Factory) getFactoryOfType(factoryType string) (strategystore.Factory, error) {
+	switch factoryType {
+	case staticStrategyStoreType:
+		return static.NewFactory(), nil
+	default:
+		return nil, fmt.Errorf("Unknown sampling strategy store type %s. Valid types are %v", factoryType, allSamplingTypes)
+	}
+}
+
+// AddFlags implements plugin.Configurable
+func (f *Factory) AddFlags(flagSet *flag.FlagSet) {
+	for _, factory := range f.factories {
+		if conf, ok := factory.(plugin.Configurable); ok {
+			conf.AddFlags(flagSet)
+		}
+	}
+}
+
+// InitFromViper implements plugin.Configurable
+func (f *Factory) InitFromViper(v *viper.Viper) {
+	for _, factory := range f.factories {
+		if conf, ok := factory.(plugin.Configurable); ok {
+			conf.InitFromViper(v)
+		}
+	}
+}
+
+// Initialize implements strategystore.Factory
+func (f *Factory) Initialize(metricsFactory metrics.Factory, logger *zap.Logger) error {
+	for _, factory := range f.factories {
+		if err := factory.Initialize(metricsFactory, logger); err != nil {
+			return err
+		}
+	}
+	return nil
+}
+
+// CreateStrategyStore implements strategystore.Factory
+func (f *Factory) CreateStrategyStore() (strategystore.StrategyStore, error) {
+	factory, ok := f.factories[f.StrategyStoreType]
+	if !ok {
+		return nil, fmt.Errorf("No %s strategy store registered", f.StrategyStoreType)
+	}
+	return factory.CreateStrategyStore()
+}

plugin/sampling/strategystore/factory_config.go

@@ -0,0 +1,42 @@
+// Copyright (c) 2018 The Jaeger Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package strategystore
+
+import (
+	"os"
+)
+
+const (
+	// SamplingTypeEnvVar is the name of the env var that defines the type of sampling strategy store used.
+	SamplingTypeEnvVar = "SAMPLING_TYPE"
+)
+
+// FactoryConfig tells the Factory what sampling type it needs to create.
+type FactoryConfig struct {
+	StrategyStoreType string
+}
+
+// FactoryConfigFromEnv reads the desired sampling type from the SAMPLING_TYPE environment variable. Allowed values:
+//   * `static` - built-in
+//   * `adaptive` - built-in // TODO
+func FactoryConfigFromEnv() FactoryConfig {
+	strategyStoreType := os.Getenv(SamplingTypeEnvVar)
+	if strategyStoreType == "" {
+		strategyStoreType = staticStrategyStoreType
+	}
+	return FactoryConfig{
+		StrategyStoreType: strategyStoreType,
+	}
+}

plugin/sampling/strategystore/factory_config_test.go

@@ -0,0 +1,34 @@
+// Copyright (c) 2018 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package strategystore
+
+import (
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func clearEnv() {
+	os.Setenv(SamplingTypeEnvVar, "")
+}
+
+func TestFactoryConfigFromEnv(t *testing.T) {
+	clearEnv()
+	defer clearEnv()
+
+	f := FactoryConfigFromEnv()
+	assert.Equal(t, staticStrategyStoreType, f.StrategyStoreType)
+}