✨ feat: add pqsql configs for open and idle connections (#1583)

When Postgres is used as the backing database for headscale,
it does not set a limit on maximum open and idle connections
which leads to hundreds of open connections to the Postgres
server.

This commit introduces the configuration variables to set those
values and also sets default while opening a new postgres connection.

CHANGELOG.md

@@ -46,6 +46,7 @@ after improving the test harness as part of adopting [#1460](https://github.com/
 - Added the possibility to manually create a DERP-map entry which can be customized, instead of automatically creating it. [#1565](https://github.com/juanfont/headscale/pull/1565)
 - Change the structure of database configuration, see [config-example.yaml](./config-example.yaml) for the new structure. [#1700](https://github.com/juanfont/headscale/pull/1700)
   - Old structure is now considered deprecated and will be removed in the future.
+  - Adds additional configuration for PostgreSQL for setting max open, idle conection and idle connection lifetime.
 
 ## 0.22.3 (2023-05-12)
 

config-example.yaml

@@ -153,6 +153,9 @@ database:
   #   name: headscale
   #   user: foo
   #   pass: bar
+  #   max_open_conns: 10
+  #   max_idle_conns: 10
+  #   conn_max_idle_time_secs: 3600
 
   #   # If other 'sslmode' is required instead of 'require(true)' and 'disabled(false)', set the 'sslmode' you need
   #   # in the 'db_ssl' field. Refers to https://www.postgresql.org/docs/current/libpq-ssl.html Table 34.1.

hscontrol/db/db.go

@@ -12,13 +12,14 @@ import (
 
 	"github.com/glebarez/sqlite"
 	"github.com/go-gormigrate/gormigrate/v2"
-	"github.com/juanfont/headscale/hscontrol/notifier"
-	"github.com/juanfont/headscale/hscontrol/types"
-	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/rs/zerolog/log"
 	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 	"gorm.io/gorm/logger"
+
+	"github.com/juanfont/headscale/hscontrol/notifier"
+	"github.com/juanfont/headscale/hscontrol/types"
+	"github.com/juanfont/headscale/hscontrol/util"
 )
 
 var errDatabaseNotSupported = errors.New("database type not supported")
@@ -50,7 +51,10 @@ func NewHeadscaleDatabase(
 		return nil, err
 	}
 
-	migrations := gormigrate.New(dbConn, gormigrate.DefaultOptions, []*gormigrate.Migration{
+	migrations := gormigrate.New(
+		dbConn,
+		gormigrate.DefaultOptions,
+		[]*gormigrate.Migration{
 			// New migrations should be added as transactions at the end of this list.
 			// The initial commit here is quite messy, completely out of order and
 			// has no versioning and is the tech debt of not having versioned migrations
@@ -67,23 +71,28 @@ func NewHeadscaleDatabase(
 
 					// the big rename from Machine to Node
 					_ = tx.Migrator().RenameTable("machines", "nodes")
-				_ = tx.Migrator().RenameColumn(&types.Route{}, "machine_id", "node_id")
+					_ = tx.Migrator().
+						RenameColumn(&types.Route{}, "machine_id", "node_id")
 
 					err = tx.AutoMigrate(types.User{})
 					if err != nil {
 						return err
 					}
 
-				_ = tx.Migrator().RenameColumn(&types.Node{}, "namespace_id", "user_id")
+					_ = tx.Migrator().
-				_ = tx.Migrator().RenameColumn(&types.PreAuthKey{}, "namespace_id", "user_id")
+						RenameColumn(&types.Node{}, "namespace_id", "user_id")
+					_ = tx.Migrator().
+						RenameColumn(&types.PreAuthKey{}, "namespace_id", "user_id")
 
-				_ = tx.Migrator().RenameColumn(&types.Node{}, "ip_address", "ip_addresses")
+					_ = tx.Migrator().
+						RenameColumn(&types.Node{}, "ip_address", "ip_addresses")
 					_ = tx.Migrator().RenameColumn(&types.Node{}, "name", "hostname")
 
 					// GivenName is used as the primary source of DNS names, make sure
 					// the field is populated and normalized if it was not when the
 					// node was registered.
-				_ = tx.Migrator().RenameColumn(&types.Node{}, "nickname", "given_name")
+					_ = tx.Migrator().
+						RenameColumn(&types.Node{}, "nickname", "given_name")
 
 					// If the Node table has a column for registered,
 					// find all occourences of "false" and drop them. Then
@@ -180,7 +189,10 @@ func NewHeadscaleDatabase(
 						}
 
 						nodesAux := []NodeAux{}
-					err := tx.Table("nodes").Select("id, enabled_routes").Scan(&nodesAux).Error
+						err := tx.Table("nodes").
+							Select("id, enabled_routes").
+							Scan(&nodesAux).
+							Error
 						if err != nil {
 							log.Fatal().Err(err).Msg("Error accessing db")
 						}
@@ -226,7 +238,9 @@ func NewHeadscaleDatabase(
 
 						err = tx.Migrator().DropColumn(&types.Node{}, "enabled_routes")
 						if err != nil {
-						log.Error().Err(err).Msg("Error dropping enabled_routes column")
+							log.Error().
+								Err(err).
+								Msg("Error dropping enabled_routes column")
 						}
 					}
 
@@ -302,7 +316,8 @@ func NewHeadscaleDatabase(
 					return nil
 				},
 			},
-	})
+		},
+	)
 
 	if err = migrations.Migrate(); err != nil {
 		log.Fatal().Err(err).Msgf("Migration failed: %v", err)
@@ -319,7 +334,6 @@ func NewHeadscaleDatabase(
 }
 
 func openDB(cfg types.DatabaseConfig) (*gorm.DB, error) {
-
 	// TODO(kradalby): Integrate this with zerolog
 	var dbLogger logger.Interface
 	if cfg.Debug {
@@ -374,10 +388,22 @@ func openDB(cfg types.DatabaseConfig) (*gorm.DB, error) {
 			dbString += fmt.Sprintf(" password=%s", cfg.Postgres.Pass)
 		}
 
-		return gorm.Open(postgres.Open(dbString), &gorm.Config{
+		db, err := gorm.Open(postgres.Open(dbString), &gorm.Config{
 			DisableForeignKeyConstraintWhenMigrating: true,
 			Logger:                                   dbLogger,
 		})
+		if err != nil {
+			return nil, err
+		}
+
+		sqlDB, _ := db.DB()
+		sqlDB.SetMaxIdleConns(cfg.Postgres.MaxIdleConnections)
+		sqlDB.SetMaxOpenConns(cfg.Postgres.MaxOpenConnections)
+		sqlDB.SetConnMaxIdleTime(
+			time.Duration(cfg.Postgres.ConnMaxIdleTimeSecs) * time.Second,
+		)
+
+		return db, nil
 	}
 
 	return nil, fmt.Errorf(

hscontrol/types/config.go

@@ -11,7 +11,6 @@ import (
 	"time"
 
 	"github.com/coreos/go-oidc/v3/oidc"
-	"github.com/juanfont/headscale/hscontrol/util"
 	"github.com/prometheus/common/model"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
@@ -20,6 +19,8 @@ import (
 	"tailscale.com/net/tsaddr"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/dnstype"
+
+	"github.com/juanfont/headscale/hscontrol/util"
 )
 
 const (
@@ -81,6 +82,9 @@ type PostgresConfig struct {
 	User                string
 	Pass                string
 	Ssl                 string
+	MaxOpenConnections  int
+	MaxIdleConnections  int
+	ConnMaxIdleTimeSecs int
 }
 
 type DatabaseConfig struct {
@@ -213,6 +217,9 @@ func LoadConfig(path string, isFile bool) error {
 
 	viper.SetDefault("db_ssl", false)
 	viper.SetDefault("database.postgres.ssl", false)
+	viper.SetDefault("database.postgres.max_open_conns", 10)
+	viper.SetDefault("database.postgres.max_idle_conns", 10)
+	viper.SetDefault("database.postgres.conn_max_idle_time_secs", 3600)
 
 	viper.SetDefault("oidc.scope", []string{oidc.ScopeOpenID, "profile", "email"})
 	viper.SetDefault("oidc.strip_email_domain", true)
@@ -287,7 +294,7 @@ func LoadConfig(path string, isFile bool) error {
 	}
 
 	if errorText != "" {
-		//nolint
+		// nolint
 		return errors.New(strings.TrimSuffix(errorText, "\n"))
 	} else {
 		return nil
@@ -429,14 +436,17 @@ func GetDatabaseConfig() DatabaseConfig {
 	case "sqlite":
 		type_ = "sqlite3"
 	default:
-		log.Fatal().Msgf("invalid database type %q, must be sqlite, sqlite3 or postgres", type_)
+		log.Fatal().
+			Msgf("invalid database type %q, must be sqlite, sqlite3 or postgres", type_)
 	}
 
 	return DatabaseConfig{
 		Type:  type_,
 		Debug: debug,
 		Sqlite: SqliteConfig{
-			Path: util.AbsolutePathFromConfigPath(viper.GetString("database.sqlite.path")),
+			Path: util.AbsolutePathFromConfigPath(
+				viper.GetString("database.sqlite.path"),
+			),
 		},
 		Postgres: PostgresConfig{
 			Host:               viper.GetString("database.postgres.host"),
@@ -445,6 +455,11 @@ func GetDatabaseConfig() DatabaseConfig {
 			User:               viper.GetString("database.postgres.user"),
 			Pass:               viper.GetString("database.postgres.pass"),
 			Ssl:                viper.GetString("database.postgres.ssl"),
+			MaxOpenConnections: viper.GetInt("database.postgres.max_open_conns"),
+			MaxIdleConnections: viper.GetInt("database.postgres.max_idle_conns"),
+			ConnMaxIdleTimeSecs: viper.GetInt(
+				"database.postgres.conn_max_idle_time_secs",
+			),
 		},
 	}
 }