feat: Add retry mechanism for storage initialization

cmd/dex/config.go

@@ -254,8 +254,10 @@ type GRPC struct {
 
 // Storage holds app's storage configuration.
 type Storage struct {
-	Type   string        `json:"type"`
-	Config StorageConfig `json:"config"`
+	Type          string        `json:"type"`
+	Config        StorageConfig `json:"config"`
+	RetryAttempts int           `json:"retryAttempts"`
+	RetryDelay    string        `json:"retryDelay"`
 }
 
 // StorageConfig is a configuration that can create a storage.
@@ -297,8 +299,10 @@ var storages = map[string]func() StorageConfig{
 // dynamically determine the type of the storage config.
 func (s *Storage) UnmarshalJSON(b []byte) error {
 	var store struct {
-		Type   string          `json:"type"`
-		Config json.RawMessage `json:"config"`
+		Type          string          `json:"type"`
+		Config        json.RawMessage `json:"config"`
+		RetryAttempts int             `json:"retryAttempts"`
+		RetryDelay    string          `json:"retryDelay"`
 	}
 	if err := json.Unmarshal(b, &store); err != nil {
 		return fmt.Errorf("parse storage: %v", err)
@@ -320,8 +324,10 @@ func (s *Storage) UnmarshalJSON(b []byte) error {
 		}
 	}
 	*s = Storage{
-		Type:   store.Type,
-		Config: storageConfig,
+		Type:          store.Type,
+		Config:        storageConfig,
+		RetryAttempts: store.RetryAttempts,
+		RetryDelay:    store.RetryDelay,
 	}
 	return nil
 }

cmd/dex/config_test.go

@@ -7,6 +7,7 @@ import (
 
 	"github.com/ghodss/yaml"
 	"github.com/kylelemons/godebug/pretty"
+	"github.com/stretchr/testify/require"
 
 	"github.com/dexidp/dex/connector/mock"
 	"github.com/dexidp/dex/connector/oidc"
@@ -447,3 +448,23 @@ logger:
 		t.Errorf("got!=want: %s", diff)
 	}
 }
+
+func TestUnmarshalConfigWithRetry(t *testing.T) {
+	rawConfig := []byte(`
+storage:
+  type: postgres
+  config:
+    host: 10.0.0.1
+    port: 65432
+  retryAttempts: 10
+  retryDelay: "1s"
+`)
+
+	var c Config
+	err := yaml.Unmarshal(rawConfig, &c)
+	require.NoError(t, err)
+
+	require.Equal(t, "postgres", c.Storage.Type)
+	require.Equal(t, 10, c.Storage.RetryAttempts)
+	require.Equal(t, "1s", c.Storage.RetryDelay)
+}

cmd/dex/serve.go

@@ -196,7 +196,7 @@ func runServe(options serveOptions) error {
 		grpcOptions = append(grpcOptions, grpc.Creds(credentials.NewTLS(tlsConfig)))
 	}
 
-	s, err := c.Storage.Config.Open(logger)
+	s, err := initializeStorageWithRetry(c.Storage, logger)
 	if err != nil {
 		return fmt.Errorf("failed to initialize storage: %v", err)
 	}
@@ -689,3 +689,37 @@ func loadTLSConfig(certFile, keyFile, caFile string, baseConfig *tls.Config) (*t
 func recordBuildInfo() {
 	buildInfo.WithLabelValues(version, runtime.Version(), fmt.Sprintf("%s/%s", runtime.GOOS, runtime.GOARCH)).Set(1)
 }
+
+// initializeStorageWithRetry opens a connection to the storage backend with a retry mechanism.
+func initializeStorageWithRetry(storageConfig Storage, logger *slog.Logger) (storage.Storage, error) {
+	var s storage.Storage
+	var err error
+
+	retryAttempts := storageConfig.RetryAttempts
+	if retryAttempts == 0 {
+		retryAttempts = 5 // Default to 5 attempts
+	}
+
+	retryDelay, err := time.ParseDuration(storageConfig.RetryDelay)
+	if err != nil {
+		retryDelay = 5 * time.Second // Default to 5 seconds
+	}
+
+	for attempt := 1; attempt <= retryAttempts; attempt++ {
+		s, err = storageConfig.Config.Open(logger)
+		if err == nil {
+			return s, nil
+		}
+
+		logger.Error("Failed to initialize storage",
+			"attempt", fmt.Sprintf("%d/%d", attempt, retryAttempts),
+			"error", err)
+
+		if attempt < retryAttempts {
+			logger.Info("Retrying storage initialization",
+				"nextAttemptIn", retryDelay.String())
+			time.Sleep(retryDelay)
+		}
+	}
+	return nil, fmt.Errorf("failed to initialize storage after %d attempts: %v", retryAttempts, err)
+}

cmd/dex/serve_test.go

@@ -1,10 +1,15 @@
 package main
 
 import (
+	"context"
+	"errors"
 	"log/slog"
 	"testing"
 
 	"github.com/stretchr/testify/require"
+
+	"github.com/dexidp/dex/storage"
+	"github.com/dexidp/dex/storage/memory"
 )
 
 func TestNewLogger(t *testing.T) {
@@ -27,3 +32,50 @@ func TestNewLogger(t *testing.T) {
 		require.Equal(t, (*slog.Logger)(nil), logger)
 	})
 }
+
+func TestStorageInitializationRetry(t *testing.T) {
+	_, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	// Create a mock storage that fails a certain number of times before succeeding
+	mockStorage := &mockRetryStorage{
+		failuresLeft: 3,
+	}
+
+	config := Config{
+		Issuer: "http://127.0.0.1:5556/dex",
+		Storage: Storage{
+			Type:          "mock",
+			Config:        mockStorage,
+			RetryAttempts: 5,
+			RetryDelay:    "1s",
+		},
+		Web: Web{
+			HTTP: "127.0.0.1:5556",
+		},
+		Logger: Logger{
+			Level:  slog.LevelInfo,
+			Format: "json",
+		},
+	}
+
+	logger, _ := newLogger(config.Logger.Level, config.Logger.Format)
+
+	s, err := initializeStorageWithRetry(config.Storage, logger)
+	require.NoError(t, err)
+	require.NotNil(t, s)
+
+	require.Equal(t, 0, mockStorage.failuresLeft)
+}
+
+type mockRetryStorage struct {
+	failuresLeft int
+}
+
+func (m *mockRetryStorage) Open(logger *slog.Logger) (storage.Storage, error) {
+	if m.failuresLeft > 0 {
+		m.failuresLeft--
+		return nil, errors.New("mock storage failure")
+	}
+	return memory.New(logger), nil
+}

config.yaml.dist

@@ -47,6 +47,11 @@ storage:
   # config:
   #   kubeConfigFile: $HOME/.kube/config
 
+  # Configuration of the retry mechanism upon a failure to storage database
+  # If not defined, the default will apply 5 attempts with 5s timeout
+  # retryAttempts: 5
+  # retryDelay: "5s"
+
 # HTTP service configuration
 web:
   http: 127.0.0.1:5556

examples/config-dev.yaml

@@ -45,6 +45,11 @@ storage:
   # config:
   #   kubeConfigFile: $HOME/.kube/config
 
+  # Configuration of the retry mechanism upon a failure to storage database
+  # If not defined, the default will apply 5 attempts with 5s timeout
+  # retryAttempts: 5
+  # retryDelay: "5s"
+
 # Configuration for the HTTP endpoints.
 web:
   http: 0.0.0.0:5556