From a5014aecfd83a7e996eb8559467228b55f89d011 Mon Sep 17 00:00:00 2001 From: pmahindrakar-oss Date: Wed, 2 Mar 2022 23:42:15 +0530 Subject: [PATCH] Added connection pooling to DB config and added defaults for the same (#358) * Added connection pooling to DB config and added defaults for the same Signed-off-by: Prafulla Mahindrakar * go mod changes for sql lite used for unit tests Signed-off-by: Prafulla Mahindrakar * linter fixes Signed-off-by: Prafulla Mahindrakar * test fix Signed-off-by: Prafulla Mahindrakar * feedback Signed-off-by: Prafulla Mahindrakar * updated the description for the flags Signed-off-by: Prafulla Mahindrakar --- flyteadmin/go.mod | 2 + flyteadmin/pkg/repositories/database.go | 21 ++++++- flyteadmin/pkg/repositories/database_test.go | 55 ++++++++++++++++++- .../runtime/application_config_provider.go | 5 ++ .../interfaces/application_configuration.go | 8 ++- 5 files changed, 86 insertions(+), 5 deletions(-) diff --git a/flyteadmin/go.mod b/flyteadmin/go.mod index a66f2709b7..5b2208b271 100644 --- a/flyteadmin/go.mod +++ b/flyteadmin/go.mod @@ -51,6 +51,7 @@ require ( google.golang.org/grpc v1.43.0 google.golang.org/protobuf v1.27.1 gorm.io/driver/postgres v1.2.3 + gorm.io/driver/sqlite v1.1.1 gorm.io/gorm v1.22.4 k8s.io/api v0.20.4 k8s.io/apimachinery v0.20.4 @@ -117,6 +118,7 @@ require ( github.com/lib/pq v1.10.4 // indirect github.com/mattn/go-colorable v0.1.12 // indirect github.com/mattn/go-isatty v0.0.14 // indirect + github.com/mattn/go-sqlite3 v2.0.3+incompatible // indirect github.com/mattn/goveralls v0.0.6 // indirect github.com/matttproud/golang_protobuf_extensions v1.0.2-0.20181231171920-c182affec369 // indirect github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect diff --git a/flyteadmin/pkg/repositories/database.go b/flyteadmin/pkg/repositories/database.go index 9ac1c9ef46..e4dd13cc3f 100644 --- a/flyteadmin/pkg/repositories/database.go +++ b/flyteadmin/pkg/repositories/database.go @@ -98,9 +98,26 @@ func GetDB(ctx context.Context, dbConfig *runtimeInterfaces.DbConfig, logConfig default: panic(fmt.Sprintf("Unrecognized database config %v", dbConfig)) } - - return gorm.Open(dialector, &gorm.Config{ + gormDb, err := gorm.Open(dialector, &gorm.Config{ Logger: gormLogger.Default.LogMode(logLevel), DisableForeignKeyConstraintWhenMigrating: !dbConfig.EnableForeignKeyConstraintWhenMigrating, }) + + if err != nil { + return nil, err + } + + // Setup connection pool settings + return gormDb, setupDbConnectionPool(gormDb, dbConfig) +} + +func setupDbConnectionPool(gormDb *gorm.DB, dbConfig *runtimeInterfaces.DbConfig) error { + genericDb, err := gormDb.DB() + if err != nil { + return err + } + genericDb.SetConnMaxLifetime(dbConfig.ConnMaxLifeTime.Duration) + genericDb.SetMaxIdleConns(dbConfig.MaxIdleConnections) + genericDb.SetMaxOpenConns(dbConfig.MaxOpenConnections) + return nil } diff --git a/flyteadmin/pkg/repositories/database_test.go b/flyteadmin/pkg/repositories/database_test.go index 5dae0033e2..5b1de474e7 100644 --- a/flyteadmin/pkg/repositories/database_test.go +++ b/flyteadmin/pkg/repositories/database_test.go @@ -3,13 +3,19 @@ package repositories import ( "context" "io/ioutil" + "os" + "path/filepath" "testing" + "time" runtimeInterfaces "github.com/flyteorg/flyteadmin/pkg/runtime/interfaces" + "github.com/flyteorg/flytestdlib/config" "github.com/flyteorg/flytestdlib/logger" - gormLogger "gorm.io/gorm/logger" "github.com/stretchr/testify/assert" + "gorm.io/driver/sqlite" + "gorm.io/gorm" + gormLogger "gorm.io/gorm/logger" ) func TestGetGormLogLevel(t *testing.T) { @@ -90,3 +96,50 @@ func TestGetPostgresDsn(t *testing.T) { assert.Equal(t, "host=localhost port=5432 dbname=postgres user=postgres password=123abc ", dsn) }) } + +func TestSetupDbConnectionPool(t *testing.T) { + t.Run("successful", func(t *testing.T) { + gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{}) + assert.Nil(t, err) + dbConfig := &runtimeInterfaces.DbConfig{ + DeprecatedPort: 5432, + MaxIdleConnections: 10, + MaxOpenConnections: 1000, + ConnMaxLifeTime: config.Duration{Duration: time.Hour}, + } + err = setupDbConnectionPool(gormDb, dbConfig) + assert.Nil(t, err) + genericDb, err := gormDb.DB() + assert.Nil(t, err) + assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000) + }) + t.Run("unset duration", func(t *testing.T) { + gormDb, err := gorm.Open(sqlite.Open(filepath.Join(os.TempDir(), "gorm.db")), &gorm.Config{}) + assert.Nil(t, err) + dbConfig := &runtimeInterfaces.DbConfig{ + DeprecatedPort: 5432, + MaxIdleConnections: 10, + MaxOpenConnections: 1000, + } + err = setupDbConnectionPool(gormDb, dbConfig) + assert.Nil(t, err) + genericDb, err := gormDb.DB() + assert.Nil(t, err) + assert.Equal(t, genericDb.Stats().MaxOpenConnections, 1000) + }) + t.Run("failed to get DB", func(t *testing.T) { + gormDb := &gorm.DB{ + Config: &gorm.Config{ + ConnPool: &gorm.PreparedStmtDB{}, + }, + } + dbConfig := &runtimeInterfaces.DbConfig{ + DeprecatedPort: 5432, + MaxIdleConnections: 10, + MaxOpenConnections: 1000, + ConnMaxLifeTime: config.Duration{Duration: time.Hour}, + } + err := setupDbConnectionPool(gormDb, dbConfig) + assert.NotNil(t, err) + }) +} diff --git a/flyteadmin/pkg/runtime/application_config_provider.go b/flyteadmin/pkg/runtime/application_config_provider.go index ff503ee13f..3731b7c726 100644 --- a/flyteadmin/pkg/runtime/application_config_provider.go +++ b/flyteadmin/pkg/runtime/application_config_provider.go @@ -1,6 +1,8 @@ package runtime import ( + "time" + "github.com/flyteorg/flyteadmin/pkg/common" "github.com/flyteorg/flyteadmin/pkg/runtime/interfaces" "github.com/flyteorg/flytestdlib/config" @@ -25,6 +27,9 @@ var databaseConfig = config.MustRegisterSection(database, &interfaces.DbConfig{ DeprecatedHost: postgres, DeprecatedDbName: postgres, DeprecatedExtraOptions: "sslmode=disable", + MaxIdleConnections: 10, + MaxOpenConnections: 1000, + ConnMaxLifeTime: config.Duration{Duration: time.Hour}, }) var flyteAdminConfig = config.MustRegisterSection(flyteAdmin, &interfaces.ApplicationConfig{ ProfilerPort: metricPort, diff --git a/flyteadmin/pkg/runtime/interfaces/application_configuration.go b/flyteadmin/pkg/runtime/interfaces/application_configuration.go index 52b7a116a4..2eca76c813 100644 --- a/flyteadmin/pkg/runtime/interfaces/application_configuration.go +++ b/flyteadmin/pkg/runtime/interfaces/application_configuration.go @@ -2,6 +2,7 @@ package interfaces import ( "github.com/flyteorg/flytestdlib/config" + "golang.org/x/time/rate" ) @@ -17,8 +18,11 @@ type DbConfig struct { DeprecatedExtraOptions string `json:"options" pflag:",deprecated"` DeprecatedDebug bool `json:"debug" pflag:",deprecated"` - EnableForeignKeyConstraintWhenMigrating bool `json:"enableForeignKeyConstraintWhenMigrating" pflag:",Whether to enable gorm foreign keys when migrating the db"` - PostgresConfig PostgresConfig `json:"postgres"` + EnableForeignKeyConstraintWhenMigrating bool `json:"enableForeignKeyConstraintWhenMigrating" pflag:",Whether to enable gorm foreign keys when migrating the db"` + MaxIdleConnections int `json:"maxIdleConnections" pflag:",maxIdleConnections sets the maximum number of connections in the idle connection pool."` + MaxOpenConnections int `json:"maxOpenConnections" pflag:",maxOpenConnections sets the maximum number of open connections to the database."` + ConnMaxLifeTime config.Duration `json:"connMaxLifeTime" pflag:",sets the maximum amount of time a connection may be reused"` + PostgresConfig PostgresConfig `json:"postgres"` } // PostgresConfig includes specific config options for opening a connection to a postgres database.