Merge branch 'ps-db-connect-hangs' into 'master'

praefect: Praefect should not hang if db is not reachable

Closes #3678

See merge request gitlab-org/gitaly!4034

7 files changed, 35 insertions, 20 deletions

diff --git a/cmd/praefect/main.go b/cmd/praefect/main.go
index 1976aa33f..e2210e659 100644
--- a/cmd/praefect/main.go
+++ b/cmd/praefect/main.go
@@ -208,20 +208,21 @@ func run(cfgs []starter.Config, conf config.Config, b bootstrap.Listener, promre
 		return err
 	}
 
-	var db *sql.DB
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
+	var db *sql.DB
 	if conf.NeedsSQL() {
-		dbConn, closedb, err := initDatabase(logger, conf)
+		logger.Infof("establishing database connection to %s:%d ...", conf.DB.Host, conf.DB.Port)
+		dbConn, closedb, err := initDatabase(ctx, logger, conf)
 		if err != nil {
 			return err
 		}
 		defer closedb()
 		db = dbConn
+		logger.Info("database connection established")
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
-
 	var queue datastore.ReplicationEventQueue
 	var rs datastore.RepositoryStore
 	var csg datastore.ConsistentStoragesGetter
@@ -430,6 +431,9 @@ func run(cfgs []starter.Config, conf config.Config, b bootstrap.Listener, promre
 	if err := b.Start(); err != nil {
 		return fmt.Errorf("unable to start the bootstrap: %v", err)
 	}
+	for _, cfg := range cfgs {
+		logger.WithFields(logrus.Fields{"schema": cfg.Name, "address": cfg.Addr}).Info("listening")
+	}
 
 	go repl.ProcessBacklog(ctx, praefect.ExpBackoffFactory{Start: time.Second, Max: 5 * time.Second})
 	logger.Info("background started: processing of the replication events")
@@ -510,8 +514,6 @@ func getStarterConfigs(conf config.Config) ([]starter.Config, error) {
 		unique[addrConf.Addr] = struct{}{}
 
 		cfgs = append(cfgs, addrConf)
-
-		logger.WithFields(logrus.Fields{"schema": schema, "address": addr}).Info("listening")
 	}
 
 	if len(cfgs) == 0 {
@@ -521,8 +523,10 @@ func getStarterConfigs(conf config.Config) ([]starter.Config, error) {
 	return cfgs, nil
 }
 
-func initDatabase(logger *logrus.Entry, conf config.Config) (*sql.DB, func(), error) {
-	db, err := glsql.OpenDB(conf.DB)
+func initDatabase(ctx context.Context, logger *logrus.Entry, conf config.Config) (*sql.DB, func(), error) {
+	openDBCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	db, err := glsql.OpenDB(openDBCtx, conf.DB)
 	if err != nil {
 		logger.WithError(err).Error("SQL connection open failed")
 		return nil, nil, err
diff --git a/cmd/praefect/subcmd.go b/cmd/praefect/subcmd.go
index 253c4d38f..05718b073 100644
--- a/cmd/praefect/subcmd.go
+++ b/cmd/praefect/subcmd.go
@@ -89,7 +89,11 @@ func getNodeAddress(cfg config.Config) (string, error) {
 }
 
 func openDB(conf config.DB) (*sql.DB, func(), error) {
-	db, err := glsql.OpenDB(conf)
+	ctx := context.Background()
+
+	openDBCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	db, err := glsql.OpenDB(openDBCtx, conf)
 	if err != nil {
 		return nil, nil, fmt.Errorf("sql open: %v", err)
 	}
diff --git a/cmd/praefect/subcmd_list_untracked_repositories.go b/cmd/praefect/subcmd_list_untracked_repositories.go
index 87e675078..1339cecc2 100644
--- a/cmd/praefect/subcmd_list_untracked_repositories.go
+++ b/cmd/praefect/subcmd_list_untracked_repositories.go
@@ -72,7 +72,9 @@ func (cmd listUntrackedRepositories) Exec(flags *flag.FlagSet, cfg config.Config
 	logger.Debug("connected to gitaly nodes")
 
 	logger.Debug("connecting to praefect database...")
-	db, err := glsql.OpenDB(cfg.DB)
+	openDBCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	db, err := glsql.OpenDB(openDBCtx, cfg.DB)
 	if err != nil {
 		return fmt.Errorf("connect to database: %w", err)
 	}
diff --git a/cmd/praefect/subcmd_remove_repository.go b/cmd/praefect/subcmd_remove_repository.go
index 1c0db75dd..6641ab7fe 100644
--- a/cmd/praefect/subcmd_remove_repository.go
+++ b/cmd/praefect/subcmd_remove_repository.go
@@ -62,14 +62,17 @@ func (cmd removeRepository) Exec(flags *flag.FlagSet, cfg config.Config) error {
 	case cmd.relativePath == "":
 		return requiredParameterError(paramRelativePath)
 	}
+	ctx := context.Background()
 
-	db, err := glsql.OpenDB(cfg.DB)
+	openDBCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	db, err := glsql.OpenDB(openDBCtx, cfg.DB)
 	if err != nil {
 		return fmt.Errorf("connect to database: %w", err)
 	}
 	defer func() { _ = db.Close() }()
 
-	ctx := correlation.ContextWithCorrelation(context.Background(), correlation.SafeRandomID())
+	ctx = correlation.ContextWithCorrelation(ctx, correlation.SafeRandomID())
 	logger := cmd.logger.WithField("correlation_id", correlation.ExtractFromContext(ctx))
 
 	return cmd.exec(ctx, logger, db, cfg)
diff --git a/cmd/praefect/subcmd_remove_repository_test.go b/cmd/praefect/subcmd_remove_repository_test.go
index 98ea5a8bd..5f9771234 100644
--- a/cmd/praefect/subcmd_remove_repository_test.go
+++ b/cmd/praefect/subcmd_remove_repository_test.go
@@ -61,7 +61,7 @@ func TestRemoveRepository_Exec_invalidArgs(t *testing.T) {
 		cfg := config.Config{DB: config.DB{Host: "stub", SSLMode: "disable"}}
 		err := cmd.Exec(flag.NewFlagSet("", flag.PanicOnError), cfg)
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "connect to database: dial tcp: lookup stub")
+		require.Contains(t, err.Error(), "connect to database: send ping: dial tcp: lookup stub")
 	})
 
 	t.Run("praefect address is not set in config ", func(t *testing.T) {
diff --git a/cmd/praefect/subcmd_track_repository.go b/cmd/praefect/subcmd_track_repository.go
index c1c50b359..1dda0f52f 100644
--- a/cmd/praefect/subcmd_track_repository.go
+++ b/cmd/praefect/subcmd_track_repository.go
@@ -65,15 +65,17 @@ func (cmd trackRepository) Exec(flags *flag.FlagSet, cfg config.Config) error {
 		}
 	}
 
-	db, err := glsql.OpenDB(cfg.DB)
+	ctx := correlation.ContextWithCorrelation(context.Background(), correlation.SafeRandomID())
+	logger := cmd.logger.WithField("correlation_id", correlation.ExtractFromContext(ctx))
+
+	openDBCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
+	defer cancel()
+	db, err := glsql.OpenDB(openDBCtx, cfg.DB)
 	if err != nil {
 		return fmt.Errorf("connect to database: %w", err)
 	}
 	defer func() { _ = db.Close() }()
 
-	ctx := correlation.ContextWithCorrelation(context.Background(), correlation.SafeRandomID())
-	logger := cmd.logger.WithField("correlation_id", correlation.ExtractFromContext(ctx))
-
 	return cmd.exec(ctx, logger, db, cfg)
 }
 
diff --git a/cmd/praefect/subcmd_track_repository_test.go b/cmd/praefect/subcmd_track_repository_test.go
index 609f798dd..c6195ad99 100644
--- a/cmd/praefect/subcmd_track_repository_test.go
+++ b/cmd/praefect/subcmd_track_repository_test.go
@@ -62,11 +62,11 @@ func TestAddRepository_Exec_invalidArgs(t *testing.T) {
 	})
 
 	t.Run("db connection error", func(t *testing.T) {
-		cmd := trackRepository{virtualStorage: "stub", relativePath: "stub", authoritativeStorage: "storage-0"}
+		cmd := trackRepository{virtualStorage: "stub", relativePath: "stub", authoritativeStorage: "storage-0", logger: testhelper.NewTestLogger(t)}
 		cfg := config.Config{DB: config.DB{Host: "stub", SSLMode: "disable"}}
 		err := cmd.Exec(flag.NewFlagSet("", flag.PanicOnError), cfg)
 		require.Error(t, err)
-		require.Contains(t, err.Error(), "connect to database: dial tcp: lookup stub")
+		require.Contains(t, err.Error(), "connect to database: send ping: dial tcp: lookup stub")
 	})
 }