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/serve.go

@@ -196,11 +196,12 @@ 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)
 	}
 	defer s.Close()
+	defer s.Close()
 
 	logger.Info("config storage", "storage_type", c.Storage.Type)
 
@@ -689,3 +690,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)
+			time.Sleep(retryDelay)
+		}
+	}
+	return nil, fmt.Errorf("failed to initialize storage after %d attempts: %v", retryAttempts, err)
+}

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