Only use CompositeWriter for multiple writers

Signed-off-by: Davit Yeghshatyan <davo@uber.com>

plugin/storage/factory.go

@@ -117,6 +117,9 @@ func (f *Factory) CreateSpanWriter() (spanstore.Writer, error) {
 		}
 		writers = append(writers, writer)
 	}
+	if len(f.SpanWriterTypes) == 1 {
+		return writers[0], nil
+	}
 	return spanstore.NewCompositeWriter(writers...), nil
 }
 

plugin/storage/factory_config.go

@@ -66,6 +66,7 @@ func FactoryConfigFromEnvAndCLI(args []string, log io.Writer) FactoryConfig {
 	if depStorageType == "" {
 		depStorageType = spanWriterTypes[0]
 	}
+	// TODO support explicit configuration for readers
 	return FactoryConfig{
 		SpanWriterTypes:         spanWriterTypes,
 		SpanReaderType:          spanWriterTypes[0],

plugin/storage/factory_test.go

@@ -94,16 +94,21 @@ func TestInitialize(t *testing.T) {
 }
 
 func TestCreate(t *testing.T) {
-	f, err := NewFactory(defaultCfg())
+	cfg := defaultCfg()
+	cfg.SpanWriterTypes = append(cfg.SpanWriterTypes, elasticsearchStorageType)
+	f, err := NewFactory(cfg)
 	require.NoError(t, err)
 	assert.NotEmpty(t, f.factories)
 	assert.NotEmpty(t, f.factories[cassandraStorageType])
 
 	mock := new(mocks.Factory)
+	mock2 := new(mocks.Factory)
 	f.factories[cassandraStorageType] = mock
+	f.factories[elasticsearchStorageType] = mock2
 
 	spanReader := new(spanStoreMocks.Reader)
 	spanWriter := new(spanStoreMocks.Writer)
+	spanWriter2 := new(spanStoreMocks.Writer)
 	depReader := new(depStoreMocks.Reader)
 
 	mock.On("CreateSpanReader").Return(spanReader, errors.New("span-reader-error"))
@@ -129,9 +134,10 @@ func TestCreate(t *testing.T) {
 	assert.EqualError(t, err, "Archive storage not supported")
 
 	mock.On("CreateSpanWriter").Return(spanWriter, nil)
+	mock2.On("CreateSpanWriter").Return(spanWriter2, nil)
 	w, err = f.CreateSpanWriter()
 	assert.NoError(t, err)
-	assert.Equal(t, spanstore.NewCompositeWriter(spanWriter), w)
+	assert.Equal(t, spanstore.NewCompositeWriter(spanWriter, spanWriter2), w)
 }
 
 func TestCreateArchive(t *testing.T) {