Keepalive

pkg/database/config.go

@@ -17,6 +17,7 @@ package database
 import (
 	"fmt"
 	"strings"
+	"time"
 
 	"github.com/google/exposure-notifications-server/pkg/keys"
 	"github.com/google/exposure-notifications-server/pkg/secrets"
@@ -25,16 +26,18 @@ import (
 
 // Config represents the env var based configuration for database connections.
 type Config struct {
-	Name              string `env:"DB_NAME" json:",omitempty"`
-	User              string `env:"DB_USER" json:",omitempty"`
-	Host              string `env:"DB_HOST, default=localhost" json:",omitempty"`
-	Port              string `env:"DB_PORT, default=5432" json:",omitempty"`
-	SSLMode           string `env:"DB_SSLMODE, default=require" json:",omitempty"`
-	ConnectionTimeout uint   `env:"DB_CONNECT_TIMEOUT" json:",omitempty"`
-	Password          string `env:"DB_PASSWORD" json:"-"` // ignored by zap's JSON formatter
-	SSLCertPath       string `env:"DB_SSLCERT" json:",omitempty"`
-	SSLKeyPath        string `env:"DB_SSLKEY" json:",omitempty"`
-	SSLRootCertPath   string `env:"DB_SSLROOTCERT" json:",omitempty"`
+	Name                   string        `env:"DB_NAME" json:",omitempty"`
+	User                   string        `env:"DB_USER" json:",omitempty"`
+	Host                   string        `env:"DB_HOST, default=localhost" json:",omitempty"`
+	Port                   string        `env:"DB_PORT, default=5432" json:",omitempty"`
+	SSLMode                string        `env:"DB_SSLMODE, default=require" json:",omitempty"`
+	ConnectionTimeout      uint          `env:"DB_CONNECT_TIMEOUT" json:",omitempty"`
+	Password               string        `env:"DB_PASSWORD" json:"-"` // ignored by zap's JSON formatter
+	SSLCertPath            string        `env:"DB_SSLCERT" json:",omitempty"`
+	SSLKeyPath             string        `env:"DB_SSLKEY" json:",omitempty"`
+	SSLRootCertPath        string        `env:"DB_SSLROOTCERT" json:",omitempty"`
+	MaxConnectionIdleTime  time.Duration `env:"DB_MAX_CONN_IDLE_TIME" json:",omitempty"`
+	ConnectionPingInterval time.Duration `env:"DB_CONNECTION_KEEPALIVE_TIME" json:",omitempty"`
 
 	// Debug is a boolean that indicates whether the database should log SQL
 	// commands.

pkg/database/database.go

@@ -20,6 +20,8 @@ import (
 	"encoding/base64"
 	"errors"
 	"fmt"
+	"sync/atomic"
+	"time"
 
 	"github.com/google/exposure-notifications-server/pkg/base64util"
 	"github.com/google/exposure-notifications-server/pkg/keys"
@@ -51,6 +53,9 @@ type Database struct {
 
 	// secretManager is used to resolve secrets.
 	secretManager secrets.SecretManager
+
+	keepAliveStarted int32
+	keepAliveStopCh  chan struct{}
 }
 
 // SupportsPerRealmSigning returns true if the configuration supports
@@ -112,6 +117,8 @@ func (c *Config) Load(ctx context.Context) (*Database, error) {
 		signingKeyManager: signingKeyManager,
 		logger:            logger,
 		secretManager:     secretManager,
+		keepAliveStarted:  0,
+		keepAliveStopCh:   make(chan struct{}),
 	}, nil
 }
 
@@ -130,6 +137,10 @@ func (db *Database) OpenWithCacher(ctx context.Context, cacher cache.Cacher) err
 		return fmt.Errorf("database gorm.Open: %w", err)
 	}
 
+	if c.MaxConnectionIdleTime > 0 {
+		rawDB.DB().SetConnMaxIdleTime(c.MaxConnectionIdleTime)
+	}
+
 	// Log SQL statements in debug mode.
 	if c.Debug {
 		rawDB = rawDB.LogMode(true)
@@ -173,15 +184,54 @@ func (db *Database) OpenWithCacher(ctx context.Context, cacher cache.Cacher) err
 
 	}
 
+	if c.ConnectionPingInterval > 0 {
+		db.KeepAlive(ctx, c.ConnectionPingInterval)
+	}
+
 	db.db = rawDB
 	return nil
 }
 
 // Close will close the database connection. Should be deferred right after Open.
 func (db *Database) Close() error {
+	if db.keepAliveStarted > 0 {
+		db.keepAliveStopCh <- struct{}{}
+	}
 	return db.db.Close()
 }
 
+// KeepAlive will ping database connections at the specified itnterval, reconnecting
+// if needed.
+func (db *Database) KeepAlive(ctx context.Context, interval time.Duration) {
+	if atomic.CompareAndSwapInt32(&db.keepAliveStarted, 0, 1) {
+		go func() {
+			logger := logging.FromContext(ctx).Named("dbkeepalive")
+
+			logger.Infow("started", "interval", interval.String())
+			defer logger.Infow("stopped")
+			var status string
+			ticker := time.NewTicker(interval)
+			for {
+				select {
+				case <-db.keepAliveStopCh:
+					logger.Info("stopping")
+					ticker.Stop()
+					return
+				case <-ticker.C:
+					ctx, cancel := context.WithTimeout(ctx, 1*time.Second)
+					defer cancel()
+					if err := db.db.DB().PingContext(ctx); err != nil {
+						status = "down"
+					} else {
+						status = "up"
+					}
+					logger.Debugw("finished ping", "status", status)
+				}
+			}
+		}()
+	}
+}
+
 // Ping attempts a connection and closes it to the database.
 func (db *Database) Ping(ctx context.Context) error {
 	return db.db.DB().PingContext(ctx)

terraform/services.tf

@@ -45,6 +45,7 @@ locals {
     DB_SSLROOTCERT                    = "secret://${google_secret_manager_secret_version.db-secret-version["sslrootcert"].id}?target=file"
     DB_USER                           = google_sql_user.user.name
     DB_VERIFICATION_CODE_DATABASE_KEY = "secret://${google_secret_manager_secret_version.db-verification-code-hmac.id}"
+    DB_CONNECTION_KEEPALIVE_TIME      = "1m"
   }
 
   firebase_config = {