feat: Add retry mechanism for storage initialization

cmd/dex/config.go

@@ -5,10 +5,12 @@ import (
 	"encoding/json"
 	"fmt"
 	"log/slog"
+	"math"
 	"net/http"
 	"net/netip"
 	"os"
 	"strings"
+	"time"
 
 	"golang.org/x/crypto/bcrypt"
 
@@ -78,6 +80,13 @@ func (c Config) Validate() error {
 		{c.GRPC.TLSMaxVersion != "" && c.GRPC.TLSMinVersion != "" && c.GRPC.TLSMinVersion > c.GRPC.TLSMaxVersion, "TLSMinVersion greater than TLSMaxVersion"},
 	}
 
+	if err := c.Storage.Retry.Validate(); err != nil {
+		checks = append(checks, struct {
+			bad    bool
+			errMsg string
+		}{true, err.Error()})
+	}
+
 	var checkErrors []string
 
 	for _, check := range checks {
@@ -256,6 +265,50 @@ type GRPC struct {
 type Storage struct {
 	Type   string        `json:"type"`
 	Config StorageConfig `json:"config"`
+	Retry  Retry         `json:"retry"`
+}
+
+// Retry holds retry mechanism configuration.
+type Retry struct {
+	MaxAttempts   int     `json:"maxAttempts"`   // Defaults to 5
+	InitialDelay  string  `json:"initialDelay"`  // Defaults to 1s
+	MaxDelay      string  `json:"maxDelay"`      // Defaults to 5s
+	BackoffFactor float64 `json:"backoffFactor"` // Defaults to 2
+}
+
+func (r *Retry) Validate() error {
+	// If retry is configured but empty, return an error
+	if r.MaxAttempts == 0 && r.InitialDelay == "" && r.MaxDelay == "" && r.BackoffFactor == 0 {
+		return fmt.Errorf("empty configuration is supplied for storage retry")
+	}
+
+	if r.MaxAttempts < 1 {
+		return fmt.Errorf("storage retry max attempts must be at least 1")
+	}
+
+	initialDelay, err := time.ParseDuration(r.InitialDelay)
+	if err != nil || initialDelay <= 0 {
+		return fmt.Errorf("storage retry initial delay must be a positive duration in go time format")
+	}
+
+	maxDelay, err := time.ParseDuration(r.MaxDelay)
+	if err != nil || maxDelay <= 0 {
+		return fmt.Errorf("storage retry max delay must be a positive duration in go time format")
+	}
+
+	if maxDelay < initialDelay {
+		return fmt.Errorf("storage retry max delay must be greater than or equal to initial delay")
+	}
+
+	if r.BackoffFactor <= 1 {
+		return fmt.Errorf("storage retry backoff factor must be greater than 1")
+	}
+	// exponential backoff algorithm-specific check
+	if float64(maxDelay) < float64(initialDelay)*math.Pow(r.BackoffFactor, float64(r.MaxAttempts-1)) {
+		return fmt.Errorf("storage retry max delay is too small for the given initial delay, backoff factor, and max attempts")
+	}
+
+	return nil
 }
 
 // StorageConfig is a configuration that can create a storage.
@@ -320,6 +373,7 @@ func (s *Storage) UnmarshalJSON(b []byte) error {
 	var store struct {
 		Type   string          `json:"type"`
 		Config json.RawMessage `json:"config"`
+		Retry  Retry           `json:"retry"`
 	}
 	if err := json.Unmarshal(b, &store); err != nil {
 		return fmt.Errorf("parse storage: %v", err)
@@ -355,9 +409,16 @@ func (s *Storage) UnmarshalJSON(b []byte) error {
 			return fmt.Errorf("parse storage config: %v", err)
 		}
 	}
+
 	*s = Storage{
 		Type:   store.Type,
 		Config: storageConfig,
+		Retry: Retry{
+			MaxAttempts:   value(store.Retry.MaxAttempts, 5),
+			InitialDelay:  value(store.Retry.InitialDelay, time.Second.String()),
+			MaxDelay:      value(store.Retry.MaxDelay, (5 * time.Second).String()),
+			BackoffFactor: value(store.Retry.BackoffFactor, 2),
+		},
 	}
 	return nil
 }
@@ -474,3 +535,11 @@ type RefreshToken struct {
 	AbsoluteLifetime  string `json:"absoluteLifetime"`
 	ValidIfNotUsedFor string `json:"validIfNotUsedFor"`
 }
+
+func value[T comparable](val, defaultValue T) T {
+	var zero T
+	if val == zero {
+		return defaultValue
+	}
+	return val
+}

cmd/dex/config_test.go

@@ -25,6 +25,12 @@ func TestValidConfiguration(t *testing.T) {
 			Config: &sql.SQLite3{
 				File: "examples/dex.db",
 			},
+			Retry: Retry{
+				MaxAttempts:   5,
+				InitialDelay:  "1s",
+				MaxDelay:      "5m",
+				BackoffFactor: 2,
+			},
 		},
 		Web: Web{
 			HTTP: "127.0.0.1:5556",
@@ -54,7 +60,8 @@ func TestInvalidConfiguration(t *testing.T) {
 	wanted := `invalid Config:
 	-	no issuer specified in config file
 	-	no storage supplied in config file
-	-	must supply a HTTP/HTTPS  address to listen on`
+	-	must supply a HTTP/HTTPS  address to listen on
+	-	empty configuration is supplied for storage retry`
 	if got != wanted {
 		t.Fatalf("Expected error message to be %q, got %q", wanted, got)
 	}
@@ -72,6 +79,11 @@ storage:
     maxIdleConns: 3
     connMaxLifetime: 30
     connectionTimeout: 3
+  retry:
+    maxAttempts: 10
+    initialDelay: "4s"
+    maxDelay: "5m"
+    backoffFactor: 2
 web:
   https: 127.0.0.1:5556
   tlsMinVersion: 1.3
@@ -152,6 +164,12 @@ additionalFeatures: [
 					ConnectionTimeout: 3,
 				},
 			},
+			Retry: Retry{
+				MaxAttempts:   10,
+				InitialDelay:  "4s",
+				MaxDelay:      "5m",
+				BackoffFactor: 2,
+			},
 		},
 		Web: Web{
 			HTTPS:         "127.0.0.1:5556",
@@ -295,6 +313,11 @@ storage:
     maxIdleConns: 3
     connMaxLifetime: 30
     connectionTimeout: 3
+  retry:
+    maxAttempts: 5
+    initialDelay: 2s
+    maxDelay: 10s
+    backoffFactor: 2
 web:
   http: 127.0.0.1:5556
 
@@ -375,6 +398,12 @@ logger:
 					ConnectionTimeout: 3,
 				},
 			},
+			Retry: Retry{
+				MaxAttempts:   5,
+				InitialDelay:  "2s",
+				MaxDelay:      "10s",
+				BackoffFactor: 2,
+			},
 		},
 		Web: Web{
 			HTTP: "127.0.0.1:5556",
@@ -452,3 +481,24 @@ 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
+//   retry:
+//     attempts: 10
+//     delay: 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.Retry.Attempts)
+// 	require.Equal(t, "1s", c.Storage.Retry.Delay)
+// }

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,40 @@ 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
+
+	maxAttempts := storageConfig.Retry.MaxAttempts
+	initialDelay, _ := time.ParseDuration(storageConfig.Retry.InitialDelay)
+	maxDelay, _ := time.ParseDuration(storageConfig.Retry.MaxDelay)
+	backoffFactor := storageConfig.Retry.BackoffFactor
+
+	delay := initialDelay
+
+	for attempt := 1; attempt <= maxAttempts; 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, maxAttempts),
+			"error", err)
+
+		if attempt < maxAttempts {
+			logger.Info("Retrying storage initialization",
+				"nextAttemptIn", delay.String())
+			time.Sleep(delay)
+
+			// Calculate next delay using exponential backoff
+			delay = time.Duration(float64(delay) * backoffFactor)
+			if delay > maxDelay {
+				delay = maxDelay
+			}
+		}
+	}
+	return nil, fmt.Errorf("failed to initialize storage after %d attempts: %v", maxAttempts, err)
+}

cmd/dex/serve_test.go

@@ -1,9 +1,12 @@
 package main
 
 import (
+	"fmt"
 	"log/slog"
 	"testing"
 
+	"github.com/dexidp/dex/storage"
+	"github.com/dexidp/dex/storage/memory"
 	"github.com/stretchr/testify/require"
 )
 
@@ -27,3 +30,51 @@ func TestNewLogger(t *testing.T) {
 		require.Equal(t, (*slog.Logger)(nil), logger)
 	})
 }
+func TestStorageInitializationRetry(t *testing.T) {
+	// 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,
+			Retry: Retry{
+				MaxAttempts:   5,
+				InitialDelay:  "1s",
+				MaxDelay:      "10s",
+				BackoffFactor: 2,
+			},
+		},
+		Web: Web{
+			HTTP: "127.0.0.1:5556",
+		},
+		Logger: Logger{
+			Level:  slog.LevelInfo,
+			Format: "json",
+		},
+	}
+
+	logger, err := newLogger(config.Logger.Level, config.Logger.Format)
+	require.NoError(t, err)
+
+	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, fmt.Errorf("mock storage failure")
+	}
+	return memory.New(logger), nil
+}

config.yaml.dist

@@ -47,6 +47,14 @@ storage:
   # config:
   #   kubeConfigFile: $HOME/.kube/config
 
+  # Configuration of the retry mechanism upon a failure to storage database
+  # If not defined, the defaults below are applied
+  # retry:
+  #   maxAttempts: 5
+  #   initialDelay: "1s"
+  #   maxDelay: "5s"
+  #   backoffFactor: 2
+
 # HTTP service configuration
 web:
   http: 127.0.0.1:5556

examples/config-dev.yaml

@@ -45,6 +45,14 @@ storage:
   # config:
   #   kubeConfigFile: $HOME/.kube/config
 
+  # Configuration of the retry mechanism upon a failure to storage database
+  # If not defined, the defaults below are applied
+  # retry:
+  #   maxAttempts: 5
+  #   initialDelay: "1s"
+  #   maxDelay: "5s"
+  #   backoffFactor: 2
+
 # Configuration for the HTTP endpoints.
 web:
   http: 0.0.0.0:5556