diff --git a/plugin/storage/factory.go b/plugin/storage/factory.go index 5a6f2261638..7a1a7f61526 100644 --- a/plugin/storage/factory.go +++ b/plugin/storage/factory.go @@ -77,6 +77,10 @@ func NewFactory(config FactoryConfig) (*Factory, error) { for _, storageType := range f.SpanWriterTypes { uniqueTypes[storageType] = struct{}{} } + // skip SamplingStorageType if it is empty. See CreateSamplingStoreFactory for details + if f.SamplingStorageType != "" { + uniqueTypes[f.SamplingStorageType] = struct{}{} + } f.factories = make(map[string]storage.Factory) for t := range uniqueTypes { ff, err := f.getFactoryOfType(t) @@ -162,6 +166,20 @@ func (f *Factory) CreateSpanWriter() (spanstore.Writer, error) { // CreateSamplingStoreFactory creates a distributedlock.Lock and samplingstore.Store for use with adaptive sampling func (f *Factory) CreateSamplingStoreFactory() (storage.SamplingStoreFactory, error) { + // if a sampling storage type was specified then use it, otherwise search all factories + // for compatibility + if f.SamplingStorageType != "" { + factory, ok := f.factories[f.SamplingStorageType] + if !ok { + return nil, fmt.Errorf("no %s backend registered for sampling store", f.SamplingStorageType) + } + ss, ok := factory.(storage.SamplingStoreFactory) + if !ok { + return nil, fmt.Errorf("storage factory of type %s does not support sampling store", f.SamplingStorageType) + } + return ss, nil + } + for _, factory := range f.factories { ss, ok := factory.(storage.SamplingStoreFactory) if ok { diff --git a/plugin/storage/factory_test.go b/plugin/storage/factory_test.go index b2762d7c1a0..4ea9f8b8ef1 100644 --- a/plugin/storage/factory_test.go +++ b/plugin/storage/factory_test.go @@ -296,20 +296,39 @@ func TestCreateError(t *testing.T) { } } -func CreateSamplingStoreFactory(t *testing.T) { +func TestCreateSamplingStoreFactory(t *testing.T) { f, err := NewFactory(defaultCfg()) require.NoError(t, err) assert.NotEmpty(t, f.factories) assert.NotEmpty(t, f.factories[cassandraStorageType]) + // if not specified sampling store is chosen from available factories ssFactory, err := f.CreateSamplingStoreFactory() assert.Equal(t, f.factories[cassandraStorageType], ssFactory) assert.NoError(t, err) + // if not specified and there's no compatible factories then return nil delete(f.factories, cassandraStorageType) ssFactory, err = f.CreateSamplingStoreFactory() assert.Nil(t, ssFactory) assert.NoError(t, err) + + // if an incompatible factory is specified return err + cfg := defaultCfg() + cfg.SamplingStorageType = "elasticsearch" + f, err = NewFactory(cfg) + require.NoError(t, err) + ssFactory, err = f.CreateSamplingStoreFactory() + assert.Nil(t, ssFactory) + assert.EqualError(t, err, "storage factory of type elasticsearch does not support sampling store") + + // if a compatible factory is specified then return it + cfg.SamplingStorageType = "cassandra" + f, err = NewFactory(cfg) + require.NoError(t, err) + ssFactory, err = f.CreateSamplingStoreFactory() + assert.Equal(t, ssFactory, f.factories["cassandra"]) + assert.NoError(t, err) } type configurable struct {