diff --git a/README.md b/README.md index 2515183..41054a9 100644 --- a/README.md +++ b/README.md @@ -205,6 +205,9 @@ LocalRecall uses environment variables to configure its behavior. These variable | `CHUNK_OVERLAP` | Overlap in characters between consecutive chunks (word-aligned). Default: 0. Use to improve context across chunk boundaries. | | `HYBRID_SEARCH_BM25_WEIGHT` | Weight for BM25 keyword search in hybrid search (default: 0.5, PostgreSQL only). | | `HYBRID_SEARCH_VECTOR_WEIGHT` | Weight for vector similarity search in hybrid search (default: 0.5, PostgreSQL only). | +| `POSTGRES_LOCK_TIMEOUT` | Per-connection `lock_timeout` for the PostgreSQL engine (default: `30s`). Bounds how long a statement waits to acquire a lock so a single stuck operation cannot make every other statement on the table queue indefinitely. Set to `0`/`off` to disable. | +| `POSTGRES_IDLE_IN_TRANSACTION_TIMEOUT` | Per-connection `idle_in_transaction_session_timeout` for the PostgreSQL engine (default: `300s`). Reaps abandoned transactions that would otherwise pin locks. Set to `0`/`off` to disable. | +| `POSTGRES_STATEMENT_TIMEOUT` | Per-connection `statement_timeout` for the PostgreSQL engine (default: unset). Bounds total statement runtime; useful to auto-abort a wedged query. Index builds are exempted, so it is safe to enable. Set to `0`/`off`/empty to disable. | | `API_KEYS` | Comma-separated list of API keys for securing access to the REST API (optional). | | `GIT_PRIVATE_KEY` | Base64-encoded SSH private key for accessing private Git repositories (optional). | diff --git a/rag/engine/postgres.go b/rag/engine/postgres.go index 43d6901..f81cefa 100644 --- a/rag/engine/postgres.go +++ b/rag/engine/postgres.go @@ -38,6 +38,9 @@ func NewPostgresDBCollection(collectionName, databaseURL string, openaiClient *o return nil, fmt.Errorf("failed to parse database URL: %w", err) } + // Apply per-connection safety timeouts before opening the pool. + applyConnTimeouts(config, os.Getenv) + // Create connection pool pool, err := pgxpool.NewWithConfig(context.Background(), config) if err != nil { @@ -124,6 +127,67 @@ func getTestEmbedding(client *openai.Client, model string) ([]float32, error) { return resp.Data[0].Embedding, nil } +// applyConnTimeouts sets per-connection safety timeouts on the pool config so +// that a single wedged or corrupt index can never hang a backend forever and +// head-of-line block every other operation on the table. +// +// A corrupt custom-index access method (e.g. a BM25 index left inconsistent by +// a past pg_resetwal) can make an INSERT spin indefinitely on a buffer-content +// lock. Such a backend holds its relation lock the whole time, so every later +// statement on that table queues behind it - one stuck insert silently stalls +// the entire vector store and saturates the connection pool. +// +// - lock_timeout (default 30s): bounds how long a statement waits to ACQUIRE +// a lock. This is the cascade-killer: queued statements fail fast instead +// of piling up for days. Always safe - it never interrupts a statement that +// is doing real work. +// - idle_in_transaction_session_timeout (default 300s): reaps abandoned +// transactions that would otherwise pin locks and the xmin horizon. +// - statement_timeout (default unset): bounds total statement runtime, which +// also kills the wedged insert itself. Left OFF by default because a +// legitimate large DiskANN/HNSW index build can exceed any fixed limit; +// index builds are exempted (see execNoStatementTimeout) so operators can +// safely opt in via POSTGRES_STATEMENT_TIMEOUT. +// +// Any value of "0" or "off" is treated as an explicit opt-out. +func applyConnTimeouts(config *pgxpool.Config, getenv func(string) string) { + if config.ConnConfig.RuntimeParams == nil { + config.ConnConfig.RuntimeParams = map[string]string{} + } + set := func(param, env, def string) { + v := def + if e := getenv(env); e != "" { + v = e + } + if v == "" || v == "0" || strings.EqualFold(v, "off") { + return + } + config.ConnConfig.RuntimeParams[param] = v + } + set("lock_timeout", "POSTGRES_LOCK_TIMEOUT", "30s") + set("idle_in_transaction_session_timeout", "POSTGRES_IDLE_IN_TRANSACTION_TIMEOUT", "300s") + set("statement_timeout", "POSTGRES_STATEMENT_TIMEOUT", "") +} + +// execNoStatementTimeout runs a DDL statement (typically a CREATE INDEX) with +// statement_timeout disabled for that statement only, so a configured +// POSTGRES_STATEMENT_TIMEOUT cannot abort a legitimately long index build. +// lock_timeout still applies, so the build never waits forever for a lock. +func (p *PostgresDB) execNoStatementTimeout(ctx context.Context, sql string) error { + tx, err := p.pool.Begin(ctx) + if err != nil { + return err + } + defer tx.Rollback(ctx) + if _, err := tx.Exec(ctx, "SET LOCAL statement_timeout = 0"); err != nil { + return err + } + if _, err := tx.Exec(ctx, sql); err != nil { + return err + } + return tx.Commit(ctx) +} + func (p *PostgresDB) setupDatabase() error { ctx := context.Background() @@ -195,7 +259,7 @@ func (p *PostgresDB) setupDatabase() error { // Create indexes // GIN index for native search - _, err = p.pool.Exec(ctx, fmt.Sprintf(` + err = p.execNoStatementTimeout(ctx, fmt.Sprintf(` CREATE INDEX IF NOT EXISTS idx_%s_search ON %s USING GIN(search_vector) `, p.tableName, p.tableName)) if err != nil { @@ -204,8 +268,8 @@ func (p *PostgresDB) setupDatabase() error { // BM25 index - required for hybrid search indexName := fmt.Sprintf("idx_%s_bm25", p.tableName) - _, err = p.pool.Exec(ctx, fmt.Sprintf(` - CREATE INDEX IF NOT EXISTS %s ON %s + err = p.execNoStatementTimeout(ctx, fmt.Sprintf(` + CREATE INDEX IF NOT EXISTS %s ON %s USING bm25(full_text) WITH (text_config='english') `, indexName, p.tableName)) if err != nil { @@ -227,7 +291,7 @@ func (p *PostgresDB) createVectorIndex(ctx context.Context, vectorscaleInstalled indexName := fmt.Sprintf("idx_%s_embedding", p.tableName) if vectorscaleInstalled { - _, err := p.pool.Exec(ctx, fmt.Sprintf(` + err := p.execNoStatementTimeout(ctx, fmt.Sprintf(` CREATE INDEX IF NOT EXISTS %s ON %s USING diskann(embedding) `, indexName, p.tableName)) @@ -238,7 +302,7 @@ func (p *PostgresDB) createVectorIndex(ctx context.Context, vectorscaleInstalled xlog.Warn("Failed to create DiskANN index, trying HNSW", "error", err) } - _, err := p.pool.Exec(ctx, fmt.Sprintf(` + err := p.execNoStatementTimeout(ctx, fmt.Sprintf(` CREATE INDEX IF NOT EXISTS %s ON %s USING hnsw(embedding vector_cosine_ops) `, indexName, p.tableName)) diff --git a/rag/engine/postgres_timeouts_test.go b/rag/engine/postgres_timeouts_test.go new file mode 100644 index 0000000..0af5786 --- /dev/null +++ b/rag/engine/postgres_timeouts_test.go @@ -0,0 +1,63 @@ +package engine + +import ( + "github.com/jackc/pgx/v5/pgxpool" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("applyConnTimeouts", func() { + var config *pgxpool.Config + + BeforeEach(func() { + var err error + // A minimal but valid DSN so ParseConfig succeeds without a live DB. + config, err = pgxpool.ParseConfig("postgres://u:p@localhost:5432/db") + Expect(err).NotTo(HaveOccurred()) + }) + + // noEnv simulates an environment where no override variables are set. + noEnv := func(string) string { return "" } + + It("sets safe lock and idle timeouts by default", func() { + applyConnTimeouts(config, noEnv) + rp := config.ConnConfig.RuntimeParams + Expect(rp).To(HaveKeyWithValue("lock_timeout", "30s")) + Expect(rp).To(HaveKeyWithValue("idle_in_transaction_session_timeout", "300s")) + }) + + It("leaves statement_timeout unset by default (legitimate index builds may be long)", func() { + applyConnTimeouts(config, noEnv) + Expect(config.ConnConfig.RuntimeParams).NotTo(HaveKey("statement_timeout")) + }) + + It("honors env overrides for each timeout", func() { + env := map[string]string{ + "POSTGRES_LOCK_TIMEOUT": "5s", + "POSTGRES_IDLE_IN_TRANSACTION_TIMEOUT": "60s", + "POSTGRES_STATEMENT_TIMEOUT": "90s", + } + applyConnTimeouts(config, func(k string) string { return env[k] }) + rp := config.ConnConfig.RuntimeParams + Expect(rp).To(HaveKeyWithValue("lock_timeout", "5s")) + Expect(rp).To(HaveKeyWithValue("idle_in_transaction_session_timeout", "60s")) + Expect(rp).To(HaveKeyWithValue("statement_timeout", "90s")) + }) + + It("treats 0 or off as an explicit opt-out (does not set the param)", func() { + env := map[string]string{ + "POSTGRES_LOCK_TIMEOUT": "0", + "POSTGRES_IDLE_IN_TRANSACTION_TIMEOUT": "off", + } + applyConnTimeouts(config, func(k string) string { return env[k] }) + rp := config.ConnConfig.RuntimeParams + Expect(rp).NotTo(HaveKey("lock_timeout")) + Expect(rp).NotTo(HaveKey("idle_in_transaction_session_timeout")) + }) + + It("does not panic when RuntimeParams is nil", func() { + config.ConnConfig.RuntimeParams = nil + Expect(func() { applyConnTimeouts(config, noEnv) }).NotTo(Panic()) + Expect(config.ConnConfig.RuntimeParams).To(HaveKey("lock_timeout")) + }) +})