Ver código fonte

Expose db connection max lifetime and idle time settings

Tulir Asokan 3 anos atrás
pai
commit
dbe25bbce6
5 arquivos alterados com 48 adições e 20 exclusões
  1. 13 8
      config/config.go
  2. 3 1
      config/upgrade.go
  3. 24 4
      database/database.go
  4. 4 0
      example-config.yaml
  5. 4 7
      main.go

+ 13 - 8
config/config.go

@@ -1,5 +1,5 @@
 // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
-// Copyright (C) 2019 Tulir Asokan
+// Copyright (C) 2022 Tulir Asokan
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -41,13 +41,7 @@ type Config struct {
 		Hostname string `yaml:"hostname"`
 		Port     uint16 `yaml:"port"`
 
-		Database struct {
-			Type string `yaml:"type"`
-			URI  string `yaml:"uri"`
-
-			MaxOpenConns int `yaml:"max_open_conns"`
-			MaxIdleConns int `yaml:"max_idle_conns"`
-		} `yaml:"database"`
+		Database DatabaseConfig `yaml:"database"`
 
 		Provisioning struct {
 			Prefix       string `yaml:"prefix"`
@@ -131,3 +125,14 @@ func (config *Config) MakeAppService() (*appservice.AppService, error) {
 	as.Registration, err = config.GetRegistration()
 	return as, err
 }
+
+type DatabaseConfig struct {
+	Type string `yaml:"type"`
+	URI  string `yaml:"uri"`
+
+	MaxOpenConns int `yaml:"max_open_conns"`
+	MaxIdleConns int `yaml:"max_idle_conns"`
+
+	ConnMaxIdleTime string `yaml:"conn_max_idle_time"`
+	ConnMaxLifetime string `yaml:"conn_max_lifetime"`
+}

+ 3 - 1
config/upgrade.go

@@ -1,5 +1,5 @@
 // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
-// Copyright (C) 2021 Tulir Asokan
+// Copyright (C) 2022 Tulir Asokan
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -40,6 +40,8 @@ func (helper *UpgradeHelper) doUpgrade() {
 	helper.Copy(Str, "appservice", "database", "uri")
 	helper.Copy(Int, "appservice", "database", "max_open_conns")
 	helper.Copy(Int, "appservice", "database", "max_idle_conns")
+	helper.Copy(Str|Null, "appservice", "database", "max_conn_idle_time")
+	helper.Copy(Str|Null, "appservice", "database", "max_conn_lifetime")
 	helper.Copy(Str, "appservice", "provisioning", "prefix")
 	if secret, ok := helper.Get(Str, "appservice", "provisioning", "shared_secret"); !ok || secret == "generate" {
 		sharedSecret := appservice.RandomString(64)

+ 24 - 4
database/database.go

@@ -1,5 +1,5 @@
 // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
-// Copyright (C) 2019 Tulir Asokan
+// Copyright (C) 2022 Tulir Asokan
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -18,6 +18,8 @@ package database
 
 import (
 	"database/sql"
+	"fmt"
+	"time"
 
 	"github.com/lib/pq"
 	_ "github.com/mattn/go-sqlite3"
@@ -25,6 +27,7 @@ import (
 	log "maunium.net/go/maulogger/v2"
 
 	"go.mau.fi/whatsmeow/store/sqlstore"
+	"maunium.net/go/mautrix-whatsapp/config"
 	"maunium.net/go/mautrix-whatsapp/database/upgrades"
 )
 
@@ -45,8 +48,8 @@ type Database struct {
 	DisappearingMessage *DisappearingMessageQuery
 }
 
-func New(dbType string, uri string, baseLog log.Logger) (*Database, error) {
-	conn, err := sql.Open(dbType, uri)
+func New(cfg config.DatabaseConfig, baseLog log.Logger) (*Database, error) {
+	conn, err := sql.Open(cfg.Type, cfg.URI)
 	if err != nil {
 		return nil, err
 	}
@@ -54,7 +57,7 @@ func New(dbType string, uri string, baseLog log.Logger) (*Database, error) {
 	db := &Database{
 		DB:      conn,
 		log:     baseLog.Sub("Database"),
-		dialect: dbType,
+		dialect: cfg.Type,
 	}
 	db.User = &UserQuery{
 		db:  db,
@@ -76,6 +79,23 @@ func New(dbType string, uri string, baseLog log.Logger) (*Database, error) {
 		db:  db,
 		log: db.log.Sub("DisappearingMessage"),
 	}
+
+	db.SetMaxOpenConns(cfg.MaxOpenConns)
+	db.SetMaxIdleConns(cfg.MaxIdleConns)
+	if len(cfg.ConnMaxIdleTime) > 0 {
+		maxIdleTimeDuration, err := time.ParseDuration(cfg.ConnMaxIdleTime)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse max_conn_idle_time: %w", err)
+		}
+		db.SetConnMaxIdleTime(maxIdleTimeDuration)
+	}
+	if len(cfg.ConnMaxLifetime) > 0 {
+		maxLifetimeDuration, err := time.ParseDuration(cfg.ConnMaxLifetime)
+		if err != nil {
+			return nil, fmt.Errorf("failed to parse max_conn_idle_time: %w", err)
+		}
+		db.SetConnMaxLifetime(maxLifetimeDuration)
+	}
 	return db, nil
 }
 

+ 4 - 0
example-config.yaml

@@ -36,6 +36,10 @@ appservice:
         # Maximum number of connections. Mostly relevant for Postgres.
         max_open_conns: 20
         max_idle_conns: 2
+        # Maximum connection idle time and lifetime before they're closed. Disabled if null.
+        # Parsed with https://pkg.go.dev/time#ParseDuration
+        max_conn_idle_time: null
+        max_conn_lifetime: null
 
     # Settings for provisioning API
     provisioning:

+ 4 - 7
main.go

@@ -1,5 +1,5 @@
 // mautrix-whatsapp - A Matrix-WhatsApp puppeting bridge.
-// Copyright (C) 2021 Tulir Asokan
+// Copyright (C) 2022 Tulir Asokan
 //
 // This program is free software: you can redistribute it and/or modify
 // it under the terms of the GNU Affero General Public License as published by
@@ -136,7 +136,7 @@ func (bridge *Bridge) GenerateRegistration() {
 }
 
 func (bridge *Bridge) MigrateDatabase() {
-	oldDB, err := database.New(flag.Arg(0), flag.Arg(1), log.DefaultLogger)
+	oldDB, err := database.New(config.DatabaseConfig{Type: flag.Arg(0), URI: flag.Arg(1)}, log.DefaultLogger)
 	if err != nil {
 		fmt.Println("Failed to open old database:", err)
 		os.Exit(30)
@@ -147,7 +147,7 @@ func (bridge *Bridge) MigrateDatabase() {
 		os.Exit(31)
 	}
 
-	newDB, err := database.New(bridge.Config.AppService.Database.Type, bridge.Config.AppService.Database.URI, log.DefaultLogger)
+	newDB, err := database.New(bridge.Config.AppService.Database, log.DefaultLogger)
 	if err != nil {
 		fmt.Println("Failed to open new database:", err)
 		os.Exit(32)
@@ -250,7 +250,7 @@ func (bridge *Bridge) Init() {
 	bridge.Log.Infoln("Initializing", VersionString)
 
 	bridge.Log.Debugln("Initializing database connection")
-	bridge.DB, err = database.New(bridge.Config.AppService.Database.Type, bridge.Config.AppService.Database.URI, bridge.Log)
+	bridge.DB, err = database.New(bridge.Config.AppService.Database, bridge.Log)
 	if err != nil {
 		bridge.Log.Fatalln("Failed to initialize database connection:", err)
 		os.Exit(14)
@@ -260,9 +260,6 @@ func (bridge *Bridge) Init() {
 	bridge.StateStore = database.NewSQLStateStore(bridge.DB)
 	bridge.AS.StateStore = bridge.StateStore
 
-	bridge.DB.SetMaxOpenConns(bridge.Config.AppService.Database.MaxOpenConns)
-	bridge.DB.SetMaxIdleConns(bridge.Config.AppService.Database.MaxIdleConns)
-
 	bridge.WAContainer = sqlstore.NewWithDB(bridge.DB.DB, bridge.Config.AppService.Database.Type, nil)
 
 	ss := bridge.Config.AppService.Provisioning.SharedSecret