Merge pull request #1 from SamuelTissot/concurrency_fetch

concurrent fetch secrets

Author: SamuelTissot

builder.go

@@ -4,12 +4,30 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
-	"github.com/BurntSushi/toml"
 	"io"
 	"reflect"
 	"strings"
+	"sync"
+	"sync/atomic"
+
+	"github.com/BurntSushi/toml"
+	"golang.org/x/sync/errgroup"
 )
 
+var builderWorkers atomic.Int32
+
+func init() {
+	builderWorkers.Store(10)
+}
+
+// SetBuilderWorkers set the amount of concurrent request to the secret manager
+func SetBuilderWorkers(n int32) {
+	if n < 1 {
+		n = 1
+	}
+	builderWorkers.Store(n)
+}
+
 type Builder struct {
 	sources  []*StaticSource
 	managers []SecretManager
@@ -59,12 +77,20 @@ func buildFromSources(config any, sources []*StaticSource) error {
 		}
 
 		if err != nil {
-			return fmt.Errorf("failed to process source: %s, %w", source.name, err)
+			return fmt.Errorf(
+				"failed to process source: %s, %w",
+				source.name,
+				err,
+			)
 		}
 
 		err = source.reader.Close()
 		if err != nil {
-			return fmt.Errorf("failed to close source: %s, %s", source.name, err)
+			return fmt.Errorf(
+				"failed to close source: %s, %s",
+				source.name,
+				err,
+			)
 		}
 	}
 
@@ -80,26 +106,66 @@ func parseJSON(config any, r io.Reader) error {
 	return json.NewDecoder(r).Decode(config)
 }
 
-func resolveSecrets(ctx context.Context, config any, managers []SecretManager) error {
+type replacement struct {
+	old string
+	new string
+}
+
+func resolveSecrets(
+	ctx context.Context,
+	config any,
+	managers []SecretManager,
+) error {
+	g, ctx := errgroup.WithContext(ctx)
+	g.SetLimit(int(builderWorkers.Load()))
+
+	var mu sync.Mutex
+
 	strConfig, err := configToJSON(config)
 	if err != nil {
 		return err
 	}
 
 	subs := findSubstitutions(strConfig)
+	replacements := make([]replacement, 0, len(subs))
 	for _, sub := range subs {
+		sub := sub
 		manager, err := getManagerForPrefix(sub.managerPrefix, managers)
 		if err != nil {
 			return err
 		}
 
-		secret, err := manager.Secret(ctx, sub.managerKey)
-		if err != nil {
-			return err
-		}
+		g.Go(func() error {
+			select {
+			case <-ctx.Done():
+				return nil
+			default:
+			}
+
+			secret, err := manager.Secret(ctx, sub.managerKey)
+			if err != nil {
+				return err
+			}
+
+			escapedSecret := escape(secret)
+
+			mu.Lock()
+			replacements = append(
+				replacements,
+				replacement{old: sub.value, new: escapedSecret},
+			)
+			mu.Unlock()
+			return nil
+		})
+	}
+
+	if err := g.Wait(); err != nil {
+		return err
+	}
 
-		escapedSecret := escape(secret)
-		strConfig = strings.Replace(strConfig, sub.value, escapedSecret, 1)
+	// actually do the replacement
+	for _, r := range replacements {
+		strConfig = strings.Replace(strConfig, r.old, r.new, 1)
 	}
 
 	return parseJSON(config, strings.NewReader(strConfig))
@@ -151,7 +217,10 @@ func findSubstitutions(str string) []substitutions {
 	return out
 }
 
-func getManagerForPrefix(prefix string, managers []SecretManager) (SecretManager, error) {
+func getManagerForPrefix(
+	prefix string,
+	managers []SecretManager,
+) (SecretManager, error) {
 	for _, manager := range managers {
 		if manager.Prefix() == prefix {
 			return manager, nil

builder_test.go

@@ -1,12 +1,13 @@
 package flowconf_test
 
 import (
-	"context"
+	"testing"
+	"time"
+
 	"github.com/SamuelTissot/flowconf"
 	"github.com/SamuelTissot/flowconf/test"
 	"github.com/stretchr/testify/assert"
-	"testing"
-	"time"
+	"github.com/stretchr/testify/mock"
 )
 
 func TestBuilder_Build_fromStaticSource(t *testing.T) {
@@ -20,7 +21,10 @@ func TestBuilder_Build_fromStaticSource(t *testing.T) {
 		{
 			name: "one source toml",
 			sources: func() []*flowconf.StaticSource {
-				sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(test.FileSystem, "data/config.toml")
+				sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(
+					test.FileSystem,
+					"data/config.toml",
+				)
 				if err != nil {
 					t.Fatalf("failed to load sources from embedded filesystem")
 				}
@@ -43,7 +47,10 @@ func TestBuilder_Build_fromStaticSource(t *testing.T) {
 		{
 			name: "one source json",
 			sources: func() []*flowconf.StaticSource {
-				sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(test.FileSystem, "data/config.json")
+				sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(
+					test.FileSystem,
+					"data/config.json",
+				)
 				if err != nil {
 					t.Fatalf("failed to load sources from embedded filesystem")
 				}
@@ -132,7 +139,6 @@ func TestBuilder_Build_fromStaticSource(t *testing.T) {
 func TestBuilder_Build_fetchesSecretsFromSecretsManagers(t *testing.T) {
 	// /////////////////////// GIVEN ///////////////////////
 	var (
-		ctx              = context.Background()
 		configSourceFile = "data/config.toml"
 		prefix           = "managerprefix"                      // this prefix needs to be the same as in the config file [[test/data/config.toml]]
 		secretKey        = "projects/id/secrets/name-of-secret" // needs to be the same as in [[test/data/config.toml]]
@@ -152,12 +158,15 @@ func TestBuilder_Build_fetchesSecretsFromSecretsManagers(t *testing.T) {
 	)
 
 	// setup sources
-	sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(test.FileSystem, configSourceFile)
+	sources, err := flowconf.NewSourcesFromEmbeddedFileSystem(
+		test.FileSystem,
+		configSourceFile,
+	)
 	assert.NoError(t, err)
 
 	// setup mock
 	managerMock.On("Prefix").Return(prefix).Once()
-	managerMock.On("Secret", ctx, secretKey).Return(secret, nil)
+	managerMock.On("Secret", mock.Anything, secretKey).Return(secret, nil)
 
 	// setup builder
 	builder := flowconf.NewBuilder(sources...)