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)
@@ -312,16 +316,17 @@ func (s *Storage) UnmarshalJSON(b []byte) error {
 	if len(store.Config) != 0 {
 		data := []byte(store.Config)
 		if featureflags.ExpandEnv.Enabled() {
-			// Caution, we're expanding in the raw JSON/YAML source. This may not be what the admin expects.
 			data = []byte(os.ExpandEnv(string(store.Config)))
 		}
 		if err := json.Unmarshal(data, storageConfig); err != nil {
 			return fmt.Errorf("parse storage config: %v", err)
 		}
 	}
 	*s = Storage{
-		Type:   store.Type,
-		Config: storageConfig,
+		Type:          store.Type,
+		Config:        storageConfig,
+		RetryAttempts: store.RetryAttempts,
+		RetryDelay:    store.RetryDelay,
 	}
 	return nil
 }

cmd/dex/serve.go

@@ -196,11 +196,16 @@ func runServe(options serveOptions) error {
 		grpcOptions = append(grpcOptions, grpc.Creds(credentials.NewTLS(tlsConfig)))
 	}
 
-	s, err := c.Storage.Config.Open(logger)
+	// s, err := c.Storage.Config.Open(logger)
+	// if err != nil {
+	// 	return fmt.Errorf("failed to initialize storage: %v", err)
+	// }
+	s, err := initializeStorageWithRetry(c, logger)
 	if err != nil {
 		return fmt.Errorf("failed to initialize storage: %v", err)
 	}
 	defer s.Close()
+	defer s.Close()
 
 	logger.Info("config storage", "storage_type", c.Storage.Type)
 
@@ -689,3 +694,32 @@ 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(c Config, logger *slog.Logger) (storage.Storage, error) {
+	var s storage.Storage
+	var err error
+	retryAttempts := c.Storage.RetryAttempts
+	if retryAttempts == 0 {
+		retryAttempts = 5 // Default to 5 attempts
+	}
+	retryDelay, err := time.ParseDuration(c.Storage.RetryDelay)
+	if err != nil {
+		retryDelay = 5 * time.Second // Default to 5 seconds
+	}
+	for attempt := 1; attempt <= retryAttempts; attempt++ {
+		s, err = c.Storage.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)
+			time.Sleep(retryDelay)
+		}
+	}
+	return nil, fmt.Errorf("failed to initialize storage after %d attempts: %v", retryAttempts, err)
+}