lee/pentagi
1
0 Fork 0
mirror of https://github.com/vxcontrol/pentagi.git synced 2026-11-06 14:04:49 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(database): enhance PostgreSQL connection pooling and configuration

Browse Source 
- Introduced shared connection pooling for PostgreSQL using `*sql.DB` for sqlc and GORM, optimizing resource usage.
- Added new environment variables: `DB_MAX_OPEN_CONNS`, `DB_MAX_IDLE_CONNS`, and `DB_VECTOR_MAX_CONNS` for configurable connection limits.
- Updated documentation to reflect new connection pooling strategy and provide operational commands for monitoring.
- Implemented shared `pgxpool` for pgvector stores to reduce connection overhead and improve performance.
- Adjusted various components to utilize the new connection pooling setup, ensuring efficient database interactions.
This commit is contained in:
Dmitry Ng committed 2026-05-18 11:13:43 +03:00
1 parent f970922ea5
commit 077ddce476
8 files changed
+221 -42
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+55
View File
@@ -2613,6 +2613,61 @@ Optional:
- `SERVER_PORT` - Port to run the server (default: `8443`)
- `SERVER_USE_SSL` - Enable SSL for the server (default: `false`)
##### PostgreSQL / pgvector connection pool sizing
PentAGI opens two independent connection pools to the same Postgres instance:
| Pool | Env var | Default | Used by |
|---|---|---|---|
| Shared `sql.DB` | `DB_MAX_OPEN_CONNS` | `25` | All sqlc queries and GORM handlers share a single `*sql.DB` |
| Shared `pgxpool` | `DB_VECTOR_MAX_CONNS` | `10` | All pgvector stores (agent memory + knowledge API) share a single pool |
Additional tuning knob:
- `DB_MAX_IDLE_CONNS` — maximum idle connections kept open in the `sql.DB` pool between requests (default: `5`).
**Budget for the stock `vxcontrol/pgvector` image** (`max_connections = 100`, `superuser_reserved_connections = 3`):
```
Available for client connections = 97
pentagi sql.DB (DB_MAX_OPEN_CONNS) = 25
pentagi pgxpool (DB_VECTOR_MAX_CONNS) = 10
pgexporter = 3
autovacuum workers = 3
─────────────────────────────────────────
Total consumed = 41
Free buffer = 56 (≈ 58 %)
```
The defaults are sized for **10 parallel flows** with concurrent API requests. If you run more flows or deploy multiple PentAGI instances against the same Postgres, raise `max_connections` via the `command` override in `docker-compose.yml` and increase the pool sizes proportionally:
```yaml
pgvector:
image: vxcontrol/pgvector:latest
command: postgres -c max_connections=200
```
To inspect the live connection budget on a running deployment:
```bash
# Postgres limits
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT name, setting FROM pg_settings
WHERE name IN ('"'"'max_connections'"'"', '"'"'superuser_reserved_connections'"'"');"'
# Current usage vs. available
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT max_conn, used, max_conn - used AS available
FROM (SELECT current_setting('"'"'max_connections'"'"')::int AS max_conn,
count(*) AS used FROM pg_stat_activity) t;"'
# Breakdown by client
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT application_name, client_addr, state, count(*)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()
GROUP BY 1, 2, 3 ORDER BY count DESC;"'
```
#### Frontend Configuration
Edit the configuration for `frontend` in `.vscode/launch.json` file:
backend/cmd/pentagi/main.go
+24 -3
View File
@@ -24,6 +24,7 @@ import (
router "pentagi/pkg/server"
"pentagi/pkg/version"
"github.com/jackc/pgx/v5/pgxpool"
_ "github.com/lib/pq"
"github.com/pressly/goose/v3"
"github.com/sirupsen/logrus"
@@ -80,17 +81,37 @@ func main() {
log.Fatalf("Unable to open database: %v\n", err)
}
db.SetMaxOpenConns(20)
db.SetMaxIdleConns(5)
db.SetMaxOpenConns(cfg.DBMaxOpenConns)
db.SetMaxIdleConns(cfg.DBMaxIdleConns)
db.SetConnMaxLifetime(time.Hour)
queries := database.New(db)
orm, err := database.NewGorm(cfg.DatabaseURL, "postgres")
// Pass the same *sql.DB to GORM so both sqlc and GORM share one connection
// pool. Previously each opened its own *sql.DB, consuming up to 40
// Postgres connections. Now together they consume at most DBMaxOpenConns.
orm, err := database.NewGorm(db, cfg.Debug)
if err != nil {
log.Fatalf("Unable to open database with gorm: %v\n", err)
}
// Create a shared pgxpool for all pgvector stores so that each executor
// reuses pooled connections instead of opening a dedicated pgx.Connect.
pgPoolConfig, err := pgxpool.ParseConfig(cfg.DatabaseURL)
if err != nil {
log.Fatalf("Failed to parse pgxpool config: %v\n", err)
}
pgPoolConfig.MaxConns = int32(cfg.DBVectorMaxConns)
pgPool, err := pgxpool.NewWithConfig(ctx, pgPoolConfig)
if err != nil {
log.Fatalf("Failed to create pgxpool: %v\n", err)
}
defer pgPool.Close()
// Attach the live pool to the config so router and tool executors can use
// pgvector.WithConn(cfg.PgxPool) without any interface changes.
cfg.PgxPool = pgPool
goose.SetBaseFS(migrations.EmbedMigrations)
if err := goose.SetDialect("postgres"); err != nil {
backend/docs/config.md
+19 -10
View File
@@ -130,6 +130,9 @@ These settings control basic application behavior and are foundational for the s
| Option | Environment Variable | Default Value | Description |
| -------------- | -------------------- | ---------------------------------------------------------------------------- | ------------------------------------------------------------------------ |
| DatabaseURL | `DATABASE_URL` | `postgres://pentagiuser:pentagipass@pgvector:5432/pentagidb?sslmode=disable` | Connection string for the PostgreSQL database with pgvector extension |
| DBMaxOpenConns | `DB_MAX_OPEN_CONNS` | `25` | Maximum open connections in the shared `sql.DB` pool (sqlc + GORM combined). See [database.md §Connection Pooling](database.md#connection-pooling). |
| DBMaxIdleConns | `DB_MAX_IDLE_CONNS` | `5` | Maximum idle connections kept open between requests |
| DBVectorMaxConns | `DB_VECTOR_MAX_CONNS` | `10` | Maximum connections in the shared `pgxpool` for all pgvector stores |
| Debug | `DEBUG` | `false` | Enables debug mode with additional logging |
| DataDir | `DATA_DIR` | `./data` | Directory for storing persistent data |
| AskUser | `ASK_USER` | `false` | When enabled, requires explicit user confirmation for certain operations |
@@ -138,21 +141,27 @@ These settings control basic application behavior and are foundational for the s
### Usage Details
- **DatabaseURL**: This is a critical setting used throughout the application for all database connections. It's used to:
- Initialize the primary SQL database connection in `main.go`
- Create GORM ORM instances for model operations
- Configure pgvector connectivity for embedding operations
- Set up connection pools in various tools and executors
- **DatabaseURL**: This is a critical setting used throughout the application for all database connections. It is used to:
- Initialize the single shared `sql.DB` connection pool in `main.go` (used by both sqlc `Queries` and GORM)
- Seed the shared `pgxpool.Pool` for all pgvector stores (agent memory + knowledge API)
For connection pool sizing and operational monitoring commands see [database.md §Connection Pooling](database.md#connection-pooling).
```go
// In main.go for SQL connection
// In main.go — one pool shared by sqlc Queries and GORM
db, err := sql.Open("postgres", cfg.DatabaseURL)
db.SetMaxOpenConns(cfg.DBMaxOpenConns)
db.SetMaxIdleConns(cfg.DBMaxIdleConns)
db.SetConnMaxLifetime(time.Hour)
// In main.go for GORM connection
orm, err := database.NewGorm(cfg.DatabaseURL, "postgres")
queries := database.New(db)
orm, err := database.NewGorm(db) // GORM wraps the same *sql.DB
// In tools for vector database operations
pgvector.WithConnectionURL(fte.cfg.DatabaseURL)
// Shared pgxpool for all pgvector stores
pgPoolConfig, _ := pgxpool.ParseConfig(cfg.DatabaseURL)
pgPoolConfig.MaxConns = int32(cfg.DBVectorMaxConns)
pgPool, _ := pgxpool.NewWithConfig(ctx, pgPoolConfig)
cfg.PgxPool = pgPool // passed to tools and router
```
- **Debug**: Controls debug mode throughout the application, enabling additional logging and development features:
backend/docs/database.md
+69 -11
View File
@@ -1199,23 +1199,81 @@ return tx.Commit()
### Connection Pooling
The package provides optimized connection pooling through GORM:
PentAGI opens two independent connection pools to the same Postgres instance:
| Pool | Env var | Default | Used by |
|---|---|---|---|
| Shared `sql.DB` | `DB_MAX_OPEN_CONNS` | `25` | sqlc `Queries` and GORM — both clients are backed by the **same** `*sql.DB` created in `main.go` |
| Shared `pgxpool` | `DB_VECTOR_MAX_CONNS` | `10` | All `pgvector.Store` instances: every flow/assistant tool executor + knowledge API |
Additional knob: `DB_MAX_IDLE_CONNS` (default `5`) — idle connections kept open between requests.
`NewGorm` accepts the already-configured `*sql.DB` so GORM never opens its own pool:
```go
func NewGorm(dsn, dbType string) (*gorm.DB, error) {
db, err := gorm.Open(dbType, dsn)
if err != nil {
return nil, err
}
// main.go — one pool, used by both sqlc and GORM
db.SetMaxOpenConns(cfg.DBMaxOpenConns)
db.SetMaxIdleConns(cfg.DBMaxIdleConns)
db.SetConnMaxLifetime(time.Hour)
// Optimized connection settings
db.DB().SetMaxIdleConns(5)
db.DB().SetMaxOpenConns(20)
db.DB().SetConnMaxLifetime(time.Hour)
orm, err := database.NewGorm(db)
return db, nil
// pkg/database/database.go
func NewGorm(db *sql.DB) (*gorm.DB, error) {
orm, err := gorm.Open("postgres", db)
...
}
```
The shared `pgxpool` is created in `main.go` and stored in `cfg.PgxPool`. All pgvector stores
(tool executors via `SetEmbedder`, knowledge API in `router.go`) use `pgvector.WithConn(cfg.PgxPool)`
instead of opening individual `pgx.Connect` calls per executor.
**Connection budget** for the stock `vxcontrol/pgvector` image (`max_connections = 100`,
`superuser_reserved_connections = 3`):
```
Available for client connections = 97
pentagi sql.DB (DB_MAX_OPEN_CONNS) = 25
pentagi pgxpool (DB_VECTOR_MAX_CONNS) = 10
pgexporter = 3
autovacuum workers = 3
─────────────────────────────────────────────
Total consumed = 41
Free buffer = 56 (≈ 58%)
```
Defaults are sized for 10 parallel flows with concurrent API requests. To inspect the live
budget on a running deployment:
```bash
# Postgres limits
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT name, setting FROM pg_settings
WHERE name IN ('"'"'max_connections'"'"', '"'"'superuser_reserved_connections'"'"');"'
# Current usage vs. available
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT max_conn, used, max_conn - used AS available
FROM (SELECT current_setting('"'"'max_connections'"'"')::int AS max_conn,
count(*) AS used FROM pg_stat_activity) t;"'
# Breakdown by client
docker exec pgvector sh -c 'psql -U "$POSTGRES_USER" -d "$POSTGRES_DB" -c \
"SELECT application_name, client_addr, state, count(*)
FROM pg_stat_activity
WHERE pid <> pg_backend_pid()
GROUP BY 1, 2, 3 ORDER BY count DESC;"'
```
To raise the Postgres limit, add a `command` override in `docker-compose.yml`:
```yaml
pgvector:
image: vxcontrol/pgvector:latest
command: postgres -c max_connections=200
```
### Vector Operations
For pgvector operations, consider:
backend/pkg/config/config.go
+14
View File
@@ -10,6 +10,7 @@ import (
"github.com/caarlos0/env/v10"
"github.com/google/uuid"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/joho/godotenv"
"github.com/vxcontrol/cloud/anonymizer/patterns"
"github.com/vxcontrol/cloud/sdk"
@@ -223,6 +224,19 @@ type Config struct {
// === Agent Planning Phase Configuration ===
AgentPlanningStepEnabled bool `env:"AGENT_PLANNING_STEP_ENABLED" envDefault:"false"`
// === Database Connection Pool Sizing ===
// See backend/docs/database.md §Connection Pooling for budget calculation and
// operational commands. Both sqlc (Queries) and GORM share a single *sql.DB,
// so DB_MAX_OPEN_CONNS is the total sql.DB budget for the process.
// Ensure DB_MAX_OPEN_CONNS + DB_VECTOR_MAX_CONNS < Postgres max_connections.
DBMaxOpenConns int `env:"DB_MAX_OPEN_CONNS" envDefault:"25"`
DBMaxIdleConns int `env:"DB_MAX_IDLE_CONNS" envDefault:"5"`
DBVectorMaxConns int `env:"DB_VECTOR_MAX_CONNS" envDefault:"10"`
// PgxPool is the shared pgxpool.Pool for all pgvector stores. Populated by
// main after pool creation; NOT sourced from environment variables.
PgxPool *pgxpool.Pool `env:"-"`
}
func NewConfig() (*Config, error) {
backend/pkg/database/database.go
+9 -8
View File
@@ -123,15 +123,16 @@ func (*GormLogger) Print(v ...interface{}) {
}
}
func NewGorm(dsn, dbType string) (*gorm.DB, error) {
db, err := gorm.Open(dbType, dsn)
// NewGorm creates a GORM instance backed by the provided *sql.DB.
// Passing an existing db ensures GORM shares the same connection pool as
// the sqlc Queries client, so the two clients together consume at most
// DBMaxOpenConns connections instead of doubling the Postgres load.
func NewGorm(db *sql.DB, debug bool) (*gorm.DB, error) {
orm, err := gorm.Open("postgres", db)
if err != nil {
return nil, err
}
db.DB().SetMaxIdleConns(5)
db.DB().SetMaxOpenConns(20)
db.DB().SetConnMaxLifetime(time.Hour)
db.SetLogger(&GormLogger{})
db.LogMode(true)
return db, nil
orm.SetLogger(&GormLogger{})
orm.LogMode(debug)
return orm, nil
}
backend/pkg/server/router.go
+8 -4
View File
@@ -132,12 +132,16 @@ func NewRouter(
embedder := providers.Embedder()
var pgStore *pgvector.Store
if embedder.IsAvailable() {
if s, err := pgvector.New(
context.Background(),
pgvector.WithConnectionURL(cfg.DatabaseURL),
opts := []pgvector.Option{
pgvector.WithEmbedder(embedder),
pgvector.WithCollectionName("langchain"),
); err == nil {
}
if cfg.PgxPool != nil {
opts = append(opts, pgvector.WithConn(cfg.PgxPool))
} else {
opts = append(opts, pgvector.WithConnectionURL(cfg.DatabaseURL))
}
if s, err := pgvector.New(context.Background(), opts...); err == nil {
pgStore = &s
} else {
logrus.WithError(err).Warn("failed to initialise pgvector store for knowledge API; embedding operations will be unavailable")
backend/pkg/tools/tools.go
+23 -6
View File
@@ -382,15 +382,26 @@ func (fte *flowToolsExecutor) SetEmbedder(embedder embeddings.Embedder) {
}
if fte.store != nil {
fte.store.Close()
// Do NOT close the store when it is backed by the shared pgxpool — calling
// Close() on that store would terminate the pool for every other executor.
// With a per-connection URL the store owns its connection and must close it.
if fte.cfg.PgxPool == nil {
fte.store.Close()
}
fte.store = nil
}
store, err := pgvector.New(
context.Background(),
pgvector.WithConnectionURL(fte.cfg.DatabaseURL),
opts := []pgvector.Option{
pgvector.WithEmbedder(embedder),
pgvector.WithCollectionName("langchain"),
)
}
if fte.cfg.PgxPool != nil {
opts = append(opts, pgvector.WithConn(fte.cfg.PgxPool))
} else {
opts = append(opts, pgvector.WithConnectionURL(fte.cfg.DatabaseURL))
}
store, err := pgvector.New(context.Background(), opts...)
if err == nil {
fte.store = &store
}
@@ -682,7 +693,13 @@ func (fte *flowToolsExecutor) cachedResourcesDir() (string, error) {
func (fte *flowToolsExecutor) Release(ctx context.Context) error {
if fte.store != nil {
fte.store.Close()
// Do NOT close the store when it is backed by the shared pgxpool — the pool
// outlives individual flows and is shared by all executors. Only close when
// the store owns its own connection (no shared pool configured).
if fte.cfg.PgxPool == nil {
fte.store.Close()
}
fte.store = nil
}
// TODO: here better to get flow containers list and purge all of them