Added SAMPLING_STORAGE_TYPE support in factory config

Signed-off-by: Joe Elliott <[email protected]>

plugin/sampling/strategystore/adaptive/factory.go

@@ -61,7 +61,7 @@ func (f *Factory) InitFromViper(v *viper.Viper, logger *zap.Logger) {
 // Initialize implements strategystore.Factory
 func (f *Factory) Initialize(metricsFactory metrics.Factory, ssFactory storage.SamplingStoreFactory, logger *zap.Logger) error {
 	if ssFactory == nil {
-		return errors.New("lock or SamplingStore nil. Please configure a backend that supports adaptive sampling")
+		return errors.New("sampling store factory is nil. Please configure a backend that supports adaptive sampling")
 	}
 
 	var err error

plugin/storage/factory_config.go

@@ -29,13 +29,17 @@ const (
 	// DependencyStorageTypeEnvVar is the name of the env var that defines the type of backend used for dependencies storage.
 	DependencyStorageTypeEnvVar = "DEPENDENCY_STORAGE_TYPE"
 
+	// SamplingStorageTypeEnvVar is the name of the env var that defines the type of backend used for sampling data storage when using adaptive sampling.
+	SamplingStorageTypeEnvVar = "SAMPLING_STORAGE_TYPE"
+
 	spanStorageFlag = "--span-storage.type"
 )
 
 // FactoryConfig tells the Factory which types of backends it needs to create for different storage types.
 type FactoryConfig struct {
 	SpanWriterTypes         []string
 	SpanReaderType          string
+	SamplingStorageType     string
 	DependenciesStorageType string
 	DownsamplingRatio       float64
 	DownsamplingHashSalt    string
@@ -73,11 +77,13 @@ func FactoryConfigFromEnvAndCLI(args []string, log io.Writer) FactoryConfig {
 	if depStorageType == "" {
 		depStorageType = spanWriterTypes[0]
 	}
+	samplingStorageType := os.Getenv(SamplingStorageTypeEnvVar)
 	// TODO support explicit configuration for readers
 	return FactoryConfig{
 		SpanWriterTypes:         spanWriterTypes,
 		SpanReaderType:          spanWriterTypes[0],
 		DependenciesStorageType: depStorageType,
+		SamplingStorageType:     samplingStorageType,
 	}
 }
 

plugin/storage/factory_config_test.go

@@ -26,6 +26,7 @@ import (
 func clearEnv() {
 	os.Setenv(SpanStorageTypeEnvVar, "")
 	os.Setenv(DependencyStorageTypeEnvVar, "")
+	os.Setenv(SamplingStorageTypeEnvVar, "")
 }
 
 func TestFactoryConfigFromEnv(t *testing.T) {
@@ -37,15 +38,18 @@ func TestFactoryConfigFromEnv(t *testing.T) {
 	assert.Equal(t, cassandraStorageType, f.SpanWriterTypes[0])
 	assert.Equal(t, cassandraStorageType, f.SpanReaderType)
 	assert.Equal(t, cassandraStorageType, f.DependenciesStorageType)
+	assert.Equal(t, "", f.SamplingStorageType)
 
 	os.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType)
 	os.Setenv(DependencyStorageTypeEnvVar, memoryStorageType)
+	os.Setenv(SamplingStorageTypeEnvVar, cassandraStorageType)
 
 	f = FactoryConfigFromEnvAndCLI(nil, &bytes.Buffer{})
 	assert.Equal(t, 1, len(f.SpanWriterTypes))
 	assert.Equal(t, elasticsearchStorageType, f.SpanWriterTypes[0])
 	assert.Equal(t, elasticsearchStorageType, f.SpanReaderType)
 	assert.Equal(t, memoryStorageType, f.DependenciesStorageType)
+	assert.Equal(t, cassandraStorageType, f.SamplingStorageType)
 
 	os.Setenv(SpanStorageTypeEnvVar, elasticsearchStorageType+","+kafkaStorageType)