Added factory support and tests

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

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 {

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 {