Retry AI service database connection on startup

internal/store/store.go

@@ -41,17 +41,48 @@ func Open(ctx context.Context, databaseURL string) (*Store, error) {
 	if err != nil {
 		return nil, fmt.Errorf("parse database url: %w", err)
 	}
-	pool, err := pgxpool.NewWithConfig(ctx, cfg)
+	pool, err := connectWithRetry(ctx, cfg, 2*time.Minute)
 	if err != nil {
-		return nil, fmt.Errorf("connect postgres: %w", err)
+		return nil, err
-	}
-	if err := pool.Ping(ctx); err != nil {
-		pool.Close()
-		return nil, fmt.Errorf("ping postgres: %w", err)
 	}
 	return &Store{pool: pool}, nil
 }
 
+func connectWithRetry(ctx context.Context, cfg *pgxpool.Config, maxWait time.Duration) (*pgxpool.Pool, error) {
+	deadline := time.Now().Add(maxWait)
+	var lastErr error
+
+	for attempt := 1; ; attempt++ {
+		pool, err := pgxpool.NewWithConfig(ctx, cfg)
+		if err == nil {
+			if pingErr := pool.Ping(ctx); pingErr == nil {
+				return pool, nil
+			} else {
+				err = fmt.Errorf("ping postgres: %w", pingErr)
+				pool.Close()
+			}
+		} else {
+			err = fmt.Errorf("connect postgres: %w", err)
+		}
+		lastErr = err
+
+		if time.Now().After(deadline) {
+			return nil, fmt.Errorf("connect postgres after retry: %w", lastErr)
+		}
+		sleep := time.Duration(attempt) * time.Second
+		if sleep > 5*time.Second {
+			sleep = 5 * time.Second
+		}
+		timer := time.NewTimer(sleep)
+		select {
+		case <-ctx.Done():
+			timer.Stop()
+			return nil, fmt.Errorf("connect postgres cancelled: %w", ctx.Err())
+		case <-timer.C:
+		}
+	}
+}
+
 func (s *Store) Close() {
 	s.pool.Close()
 }