From a84d2ae3d21a86915cc63560d432189f524f7d12 Mon Sep 17 00:00:00 2001 From: qinling0210 <88864212+qinling0210@users.noreply.github.com> Date: Thu, 23 Jul 2026 20:37:38 +0800 Subject: [PATCH] fix(infinity): avoid hangs with pooled connections and bounded timeouts (#17287) ### Summary **Issue** When calling Infinity under concurrent load, one shared connection was reused across concurrent operations, it just hangs. **Solution** Avoid hangs with pooled connections and bounded timeouts --- go.mod | 4 +- go.sum | 4 +- internal/admin/handler.go | 8 +- internal/common/error_code.go | 110 +++++------ internal/engine/infinity/chunk.go | 42 +++-- internal/engine/infinity/client.go | 234 ++++++++++++++++++++++-- internal/engine/infinity/client_test.go | 75 ++++++++ internal/engine/infinity/common.go | 26 ++- internal/engine/infinity/document.go | 9 +- internal/engine/infinity/metadata.go | 78 +++++--- 10 files changed, 461 insertions(+), 129 deletions(-) create mode 100644 internal/engine/infinity/client_test.go diff --git a/go.mod b/go.mod index 7bffd8cd7b..5918c2ce24 100644 --- a/go.mod +++ b/go.mod @@ -8,6 +8,7 @@ require ( github.com/alibabacloud-go/agentrun-20250910/v5 v5.8.4 github.com/alibabacloud-go/darabonba-openapi/v2 v2.2.1 github.com/alicebob/miniredis/v2 v2.38.0 + github.com/apache/thrift v0.23.0 github.com/aws/aws-sdk-go-v2 v1.41.3 github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream v1.7.6 github.com/aws/aws-sdk-go-v2/config v1.32.11 @@ -86,7 +87,6 @@ require ( github.com/alibabacloud-go/tea-utils/v2 v2.0.9 // indirect github.com/aliyun/credentials-go v1.4.5 // indirect github.com/antithesishq/antithesis-sdk-go v0.7.0-default-no-op // indirect - github.com/apache/thrift v0.23.0 // indirect github.com/apapsch/go-jsonmerge/v2 v2.0.0 // indirect github.com/aws/aws-sdk-go-v2/feature/ec2/imds v1.18.19 // indirect github.com/aws/aws-sdk-go-v2/internal/configsources v1.4.19 // indirect @@ -220,4 +220,4 @@ require ( modernc.org/sqlite v1.23.1 // indirect ) -replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929 +replace github.com/infiniflow/infinity-go-sdk => github.com/infiniflow/infinity/go v0.0.0-20260723093510-ceb4bc518010 diff --git a/go.sum b/go.sum index 50d1828244..9d00b41570 100644 --- a/go.sum +++ b/go.sum @@ -302,8 +302,8 @@ github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/ad github.com/hashicorp/hcl v1.0.0 h1:0Anlzjpi4vEasTeNFn2mLJgTSwt0+6sfsiTG8qcWGx4= github.com/hashicorp/hcl v1.0.0/go.mod h1:E5yfLk+7swimpb2L/Alb/PJmXilQ/rhwaUYs4T20WEQ= github.com/hpcloud/tail v1.0.0/go.mod h1:ab1qPbhIpdTxEkNHXyeSf5vhxWSCs/tWer42PpOxQnU= -github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929 h1:0M1BNouFVpnF12XEmF/42aR8CRU0bt/rMEVEsRUtSfQ= -github.com/infiniflow/infinity/go v0.0.0-20260424025959-72028e662929/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= +github.com/infiniflow/infinity/go v0.0.0-20260723093510-ceb4bc518010 h1:TKHL99MrpO2d+q6zRUSPRWsxTXo+pojFr3teiK8HLWw= +github.com/infiniflow/infinity/go v0.0.0-20260723093510-ceb4bc518010/go.mod h1:hw3z5AwNFsGy1cdrE0Mfjot2y9jqVHTxBufUx9VzZ+0= github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a h1:Inib12UR9HAfBubrGNraPjKt/Cu8xPbTJbC50+0wP5U= github.com/iromli/go-itsdangerous v0.0.0-20220223194502-9c8bef8dac6a/go.mod h1:8N0Hlye5Lzw+H/yHWpZMkT0QLA+iOHG7KLdvAm95DZg= github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E= diff --git a/internal/admin/handler.go b/internal/admin/handler.go index ff7ffaa78d..0f8f108d6e 100644 --- a/internal/admin/handler.go +++ b/internal/admin/handler.go @@ -1245,10 +1245,16 @@ func (h *Handler) PingCache(c *gin.Context) { } func (h *Handler) PingEngine(c *gin.Context) { - docEngine := engine.Get() ctx := context.Background() if err := docEngine.Ping(ctx); err != nil { + var coded interface { + Code() common.ErrorCode + } + if errors.As(err, &coded) { + common.ResponseWithHttpCodeData(c, http.StatusServiceUnavailable, coded.Code(), nil, err.Error()) + return + } common.ErrorWithCode(c, common.CodeServerError, err.Error()) return } diff --git a/internal/common/error_code.go b/internal/common/error_code.go index 2faba7c06a..2d35306c39 100644 --- a/internal/common/error_code.go +++ b/internal/common/error_code.go @@ -21,63 +21,65 @@ import "errors" type ErrorCode int const ( - CodeSuccess ErrorCode = 0 - CodeLackResources ErrorCode = 1 - CodeNotEffective ErrorCode = 10 - CodeExceptionError ErrorCode = 100 - CodeArgumentError ErrorCode = 101 - CodeDataError ErrorCode = 102 - CodeOperatingError ErrorCode = 103 - CodeTimeoutError ErrorCode = 104 - CodeConnectionError ErrorCode = 105 - CodeRunning ErrorCode = 106 - CodeResourceExhausted ErrorCode = 107 - CodePermissionError ErrorCode = 108 - CodeAuthenticationError ErrorCode = 109 - CodeParamError ErrorCode = 110 - CodeLicenseValid ErrorCode = 320 - CodeLicenseInactiveError ErrorCode = 321 - CodeLicenseExpiredError ErrorCode = 322 - CodeLicenseDigestError ErrorCode = 323 - CodeLicenseTimeRollback ErrorCode = 324 - CodeLicenseNotFound ErrorCode = 325 - CodeLicenseUnexpectedError ErrorCode = 326 - CodeBadRequest ErrorCode = 400 - CodeUnauthorized ErrorCode = 401 - CodeForbidden ErrorCode = 403 - CodeNotFound ErrorCode = 404 - CodeConflict ErrorCode = 409 - CodeServerError ErrorCode = 500 - CodeNotImplemented ErrorCode = 501 + CodeSuccess ErrorCode = 0 + CodeLackResources ErrorCode = 1 + CodeNotEffective ErrorCode = 10 + CodeExceptionError ErrorCode = 100 + CodeArgumentError ErrorCode = 101 + CodeDataError ErrorCode = 102 + CodeOperatingError ErrorCode = 103 + CodeTimeoutError ErrorCode = 104 + CodeConnectionError ErrorCode = 105 + CodeRunning ErrorCode = 106 + CodeResourceExhausted ErrorCode = 107 + CodePermissionError ErrorCode = 108 + CodeAuthenticationError ErrorCode = 109 + CodeParamError ErrorCode = 110 + CodeConnectionPoolExhausted ErrorCode = 111 + CodeLicenseValid ErrorCode = 320 + CodeLicenseInactiveError ErrorCode = 321 + CodeLicenseExpiredError ErrorCode = 322 + CodeLicenseDigestError ErrorCode = 323 + CodeLicenseTimeRollback ErrorCode = 324 + CodeLicenseNotFound ErrorCode = 325 + CodeLicenseUnexpectedError ErrorCode = 326 + CodeBadRequest ErrorCode = 400 + CodeUnauthorized ErrorCode = 401 + CodeForbidden ErrorCode = 403 + CodeNotFound ErrorCode = 404 + CodeConflict ErrorCode = 409 + CodeServerError ErrorCode = 500 + CodeNotImplemented ErrorCode = 501 ) var errorMessages = map[ErrorCode]string{ - CodeSuccess: "Success", - CodeNotEffective: "Not effective", - CodeExceptionError: "System exception", - CodeArgumentError: "Invalid argument", - CodeDataError: "Data error", - CodeOperatingError: "Operation error", - CodeTimeoutError: "Timeout", - CodeConnectionError: "Connection error", - CodeRunning: "System running", - CodeResourceExhausted: "Resource exhausted", - CodePermissionError: "Permission denied", - CodeAuthenticationError: "Authentication failed", - CodeParamError: "Invalid parameters", - CodeLicenseValid: "License valid", - CodeLicenseInactiveError: "License inactive", - CodeLicenseExpiredError: "License expired", - CodeLicenseDigestError: "License digest error", - CodeLicenseTimeRollback: "License time rollback detected", - CodeLicenseNotFound: "License not found", - CodeLicenseUnexpectedError: "Unexpected license error", - CodeBadRequest: "Bad request", - CodeUnauthorized: "Unauthorized", - CodeForbidden: "Forbidden", - CodeNotFound: "Resource not found", - CodeConflict: "Resource conflict", - CodeServerError: "Internal server error", + CodeSuccess: "Success", + CodeNotEffective: "Not effective", + CodeExceptionError: "System exception", + CodeArgumentError: "Invalid argument", + CodeDataError: "Data error", + CodeOperatingError: "Operation error", + CodeTimeoutError: "Timeout", + CodeConnectionError: "Connection error", + CodeRunning: "System running", + CodeResourceExhausted: "Resource exhausted", + CodePermissionError: "Permission denied", + CodeAuthenticationError: "Authentication failed", + CodeParamError: "Invalid parameters", + CodeConnectionPoolExhausted: "Connection pool exhausted", + CodeLicenseValid: "License valid", + CodeLicenseInactiveError: "License inactive", + CodeLicenseExpiredError: "License expired", + CodeLicenseDigestError: "License digest error", + CodeLicenseTimeRollback: "License time rollback detected", + CodeLicenseNotFound: "License not found", + CodeLicenseUnexpectedError: "Unexpected license error", + CodeBadRequest: "Bad request", + CodeUnauthorized: "Unauthorized", + CodeForbidden: "Forbidden", + CodeNotFound: "Resource not found", + CodeConflict: "Resource conflict", + CodeServerError: "Internal server error", } func (e ErrorCode) Message() string { diff --git a/internal/engine/infinity/chunk.go b/internal/engine/infinity/chunk.go index 5f824685c9..30e0b0c564 100644 --- a/internal/engine/infinity/chunk.go +++ b/internal/engine/infinity/chunk.go @@ -51,6 +51,16 @@ var pagerankAdjustLocks [pagerankAdjustLockCount]sync.Mutex // The full table name is built as "{baseName}_{datasetID}" // For skill index (datasetID="skill"), tableName is just baseName and uses skill_infinity_mapping.json func (e *infinityEngine) CreateChunkStore(ctx context.Context, baseName, datasetID string, vectorSize int, parserID string) error { + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") + if err != nil { + return fmt.Errorf("failed to get database: %w", err) + } + defer release() + + return e.createChunkStoreWithDB(db, baseName, datasetID, vectorSize, parserID) +} + +func (e *infinityEngine) createChunkStoreWithDB(db *infinity.Database, baseName, datasetID string, vectorSize int, parserID string) error { vecSize := vectorSize // Determine table name and mapping file based on index type @@ -82,17 +92,11 @@ func (e *infinityEngine) CreateChunkStore(ctx context.Context, baseName, dataset return fmt.Errorf("failed to parse mapping file: %w", err) } - // Get database - db, err := e.client.conn.GetDatabase(e.client.dbName) - if err != nil { - return fmt.Errorf("failed to get database: %w", err) - } - // Determine vector column name vectorColName := fmt.Sprintf("q_%d_vec", vecSize) // Check if table already exists - exists, err := e.tableExists(ctx, tableName) + exists, err := e.tableExistsWithDB(db, tableName) if err != nil { return fmt.Errorf("failed to check if table exists: %w", err) } @@ -268,10 +272,11 @@ func (e *infinityEngine) InsertChunks(ctx context.Context, chunks []map[string]i tableName := buildChunkTableName(baseName, datasetID) common.Info("InfinityConnection.InsertChunks called", zap.String("tableName", tableName), zap.Int("chunkCount", len(chunks))) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return nil, fmt.Errorf("Failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -307,7 +312,7 @@ func (e *infinityEngine) InsertChunks(ctx context.Context, chunks []map[string]i } // Create table - if err := e.CreateChunkStore(ctx, baseName, datasetID, vectorSize, parserID); err != nil { + if err := e.createChunkStoreWithDB(db, baseName, datasetID, vectorSize, parserID); err != nil { return nil, fmt.Errorf("Failed to create table: %w", err) } @@ -387,10 +392,11 @@ func (e *infinityEngine) UpdateChunks(ctx context.Context, condition map[string] tableName := buildChunkTableName(baseName, datasetID) common.Info("InfinityConnection.UpdateChunks called", zap.String("tableName", tableName), zap.Any("condition", condition)) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return fmt.Errorf("Failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -540,7 +546,7 @@ func (e *infinityEngine) AdjustChunkPagerank(ctx context.Context, baseName, chun if ctx == nil { ctx = context.Background() } - if e.client == nil || e.client.conn == nil { + if e.client == nil || e.client.pool == nil { return fmt.Errorf("Infinity client not initialized") } @@ -549,10 +555,11 @@ func (e *infinityEngine) AdjustChunkPagerank(ctx context.Context, baseName, chun lock.Lock() defer lock.Unlock() - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { return fmt.Errorf("failed to get table %s: %w", tableName, err) @@ -626,10 +633,11 @@ func (e *infinityEngine) AdjustChunkPagerank(ctx context.Context, baseName, chun func (e *infinityEngine) DeleteChunks(ctx context.Context, condition map[string]interface{}, baseName string, datasetID string) (int64, error) { tableName := buildChunkTableName(baseName, datasetID) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return 0, fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -700,10 +708,11 @@ func (e *infinityEngine) Search(ctx context.Context, req *types.SearchRequest) ( offset = 0 } - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return nil, fmt.Errorf("failed to get database: %w", err) } + defer release() isSkillIndex := false for _, idx := range req.IndexNames { @@ -1201,7 +1210,7 @@ func (e *infinityEngine) Search(ctx context.Context, req *types.SearchRequest) ( // GetChunk gets a chunk by ID func (e *infinityEngine) GetChunk(ctx context.Context, tableName, chunkID string, datasetIDs []string) (interface{}, error) { - if e.client == nil || e.client.conn == nil { + if e.client == nil || e.client.pool == nil { return nil, fmt.Errorf("Infinity client not initialized") } @@ -1217,10 +1226,11 @@ func (e *infinityEngine) GetChunk(ctx context.Context, tableName, chunkID string } // Try each table and collect results from all tables - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "chunk.go") if err != nil { return nil, fmt.Errorf("failed to get database: %w", err) } + defer release() // Collect chunks from all tables (same as Python's concat_dataframes) allChunks := make(map[string]map[string]interface{}) diff --git a/internal/engine/infinity/client.go b/internal/engine/infinity/client.go index ed141d300e..43a8196467 100644 --- a/internal/engine/infinity/client.go +++ b/internal/engine/infinity/client.go @@ -18,21 +18,31 @@ package infinity import ( "context" + "errors" "fmt" + "os" "ragflow/internal/common" "reflect" "strconv" "strings" + "sync/atomic" "time" + "github.com/apache/thrift/lib/go/thrift" + "ragflow/internal/server" infinity "github.com/infiniflow/infinity-go-sdk" + "go.uber.org/zap" ) type infinityClient struct { - conn *infinity.InfinityConnection - dbName string + pool *infinity.ConnectionPool + poolMaxOpen int + dbName string + + getDatabaseSeq atomic.Uint64 + getDatabaseInflight atomic.Int64 // Original URI from config, used by RunSQL to extract the host. hostURI string @@ -44,7 +54,67 @@ type infinityClient struct { mappingFileName string } +// defaultPoolMaxSize is the hard upper bound (MaxOpen) for the Go Infinity +// connection pool. Invalid or missing INFINITY_POOL_MAX_SIZE values fall back +// to this rather than becoming unbounded. +const defaultPoolMaxSize = 4 + +// defaultMaxIdleConnections caps how many established connections the pool +// keeps idle for reuse, independent of the pool's hard open cap (MaxOpen). +// It bounds idle socket retention on the Infinity server when MaxOpen is large. +const defaultMaxIdleConnections = 2 + +// defaultOperationTimeout bounds the total duration of a single pool operation +// (applied via ensureDeadline when the caller context has no deadline). +const defaultOperationTimeout = 30 * time.Second + +// defaultSocketTimeout bounds individual thrift read/write waits on a checked-out +// connection. +const defaultSocketTimeout = 30 * time.Second + +// resolvePoolMaxSize reads INFINITY_POOL_MAX_SIZE and falls back to defaultPoolMaxSize. A value < 1 is +// rejected so the pool never becomes unbounded. +func resolvePoolMaxSize() int { + raw := os.Getenv("INFINITY_POOL_MAX_SIZE") + if raw == "" { + return defaultPoolMaxSize + } + v, err := strconv.Atoi(raw) + if err != nil || v < 1 { + common.Warn("INFINITY_POOL_MAX_SIZE must be a positive integer; using default", + zap.String("value", raw), zap.Int("default", defaultPoolMaxSize)) + return defaultPoolMaxSize + } + return v +} + // NewInfinityClient creates a new Infinity client using the SDK +type poolExhaustedError struct { + caller string +} + +func (e *poolExhaustedError) Error() string { + return fmt.Sprintf("Infinity connection pool exhausted while %s", e.caller) +} + +func (e *poolExhaustedError) Code() common.ErrorCode { + return common.CodeConnectionPoolExhausted +} + +func (e *poolExhaustedError) Message() string { + return e.Error() +} + +func ensureDeadline(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) { + if ctx == nil { + return context.WithTimeout(context.Background(), timeout) + } + if _, ok := ctx.Deadline(); ok { + return ctx, func() {} + } + return context.WithTimeout(ctx, timeout) +} + func NewInfinityClient(cfg *server.InfinityConfig) (*infinityClient, error) { // Parse URI like "localhost:23817" to get IP and port host := "127.0.0.1" @@ -60,12 +130,40 @@ func NewInfinityClient(cfg *server.InfinityConfig) (*infinityClient, error) { } } - // Retry connecting for up to 120 seconds (24 attempts * 5 seconds) + // Build the connection pool object. Connections are created lazily on + // first use (InitialSize=0), so this does not verify reachability; the + // actual connectivity check happens later in WaitForHealthy. The retry + // loop only re-attempts pool construction against transient config errors. common.Info("Connecting to Infinity") - var conn *infinity.InfinityConnection + uri := infinity.NetworkAddress{IP: host, Port: port} + maxOpen := resolvePoolMaxSize() + var pool *infinity.ConnectionPool var err error for i := 0; i < 24; i++ { - conn, err = infinity.Connect(infinity.NetworkAddress{IP: host, Port: port}) + // Cap idle connections retained for reuse; never exceed the pool's + // hard open cap (MaxOpen). + maxIdle := maxOpen + if maxIdle > defaultMaxIdleConnections { + maxIdle = defaultMaxIdleConnections + } + poolCfg := infinity.DefaultConnectionPoolConfig(uri) + // Don't pre-open connections at process startup (InitialSize=0); + // connections are created lazily on first use. Keep MaxIdle so that + // established connections are retained for reuse instead of being + // reopened on every request. + poolCfg.InitialSize = 0 + poolCfg.MaxOpen = maxOpen + poolCfg.MaxIdle = maxIdle + pool, err = infinity.NewConnectionPool(poolCfg, func(u infinity.URI) (*infinity.InfinityConnection, error) { + networkAddress, ok := u.(infinity.NetworkAddress) + if !ok { + return nil, fmt.Errorf("unexpected URI type: %T", u) + } + return infinity.NewInfinityConnectionWithConfig(networkAddress, &thrift.TConfiguration{ + ConnectTimeout: server.DefaultConnectTimeout, + SocketTimeout: defaultSocketTimeout, + }) + }) if err == nil { break } @@ -78,7 +176,8 @@ func NewInfinityClient(cfg *server.InfinityConfig) (*infinityClient, error) { } client := &infinityClient{ - conn: conn, + pool: pool, + poolMaxOpen: maxOpen, dbName: cfg.DBName, hostURI: cfg.URI, postgresPort: cfg.PostgresPort, @@ -99,7 +198,13 @@ func (c *infinityClient) WaitForHealthy(ctx context.Context, timeout time.Durati default: } - res, err := c.conn.ShowCurrentNode() + conn, release, err := c.checkoutConn(ctx, "WaitForHealthy") + if err != nil { + time.Sleep(5 * time.Second) + continue + } + res, err := conn.ShowCurrentNode() + release() if err != nil { time.Sleep(5 * time.Second) continue @@ -108,6 +213,10 @@ func (c *infinityClient) WaitForHealthy(ctx context.Context, timeout time.Durati // since ShowCurrentNodeResponse is in an internal package v := reflect.ValueOf(res) if v.Kind() != reflect.Ptr { + common.Warn("Infinity health check returned unexpected response kind", + zap.String("type", fmt.Sprintf("%T", res)), + zap.String("kind", v.Kind().String()), + ) time.Sleep(5 * time.Second) continue } @@ -115,6 +224,11 @@ func (c *infinityClient) WaitForHealthy(ctx context.Context, timeout time.Durati errorCode := v.FieldByName("ErrorCode") serverStatus := v.FieldByName("ServerStatus") if !errorCode.IsValid() || !serverStatus.IsValid() { + common.Warn("Infinity health check response missing expected fields", + zap.String("type", fmt.Sprintf("%T", res)), + zap.Bool("has_error_code", errorCode.IsValid()), + zap.Bool("has_server_status", serverStatus.IsValid()), + ) time.Sleep(5 * time.Second) continue } @@ -131,6 +245,91 @@ func (c *infinityClient) WaitForHealthy(ctx context.Context, timeout time.Durati return fmt.Errorf("Infinity not healthy after %v", timeout) } +func (c *infinityClient) checkoutConn(ctx context.Context, caller string) (*infinity.InfinityConnection, func(), error) { + if c == nil || c.pool == nil { + return nil, nil, fmt.Errorf("Infinity client not initialized") + } + ctx, cancel := ensureDeadline(ctx, defaultOperationTimeout) + conn, err := c.pool.GetContext(ctx) + if err != nil { + cancel() + if (errors.Is(err, context.DeadlineExceeded) || errors.Is(err, context.Canceled)) && c.poolMaxOpen > 0 { + stats := c.pool.Stats() + if stats.TotalConnections >= c.poolMaxOpen && stats.AvailableConnections == 0 { + return nil, nil, &poolExhaustedError{caller: caller} + } + } + return nil, nil, err + } + released := false + release := func() { + if released { + return + } + released = true + defer cancel() + if err := c.pool.Put(conn); err != nil { + common.Warn("Infinity connection release failed", + zap.String("caller", caller), + zap.String("conn_ptr", fmt.Sprintf("%p", conn)), + zap.Error(err), + ) + } + } + return conn, release, nil +} + +func (c *infinityClient) checkoutDatabase(ctx context.Context, caller string) (*infinity.Database, func(), error) { + conn, release, err := c.checkoutConn(ctx, caller) + if err != nil { + return nil, nil, err + } + reqID := c.getDatabaseSeq.Add(1) + inflight := c.getDatabaseInflight.Add(1) + start := time.Now() + stats := c.pool.Stats() + common.Info("Infinity GetDatabase begin", + zap.Uint64("req_id", reqID), + zap.String("caller", caller), + zap.String("db_name", c.dbName), + zap.Bool("is_connected", conn.IsConnected()), + zap.Int64("inflight", inflight), + zap.String("conn_ptr", fmt.Sprintf("%p", conn)), + zap.Int("pool_in_use", stats.InUseConnections), + zap.Int("pool_available", stats.AvailableConnections), + ) + db, err := conn.GetDatabase(c.dbName) + if err != nil { + inflight = c.getDatabaseInflight.Add(-1) + elapsed := time.Since(start) + common.Warn("Infinity GetDatabase failed", + zap.Uint64("req_id", reqID), + zap.String("caller", caller), + zap.String("db_name", c.dbName), + zap.Bool("is_connected", conn.IsConnected()), + zap.Int64("inflight", inflight), + zap.Duration("elapsed", elapsed), + zap.Error(err), + ) + release() + return nil, nil, err + } + wrappedRelease := func() { + inflight := c.getDatabaseInflight.Add(-1) + elapsed := time.Since(start) + common.Info("Infinity GetDatabase done", + zap.Uint64("req_id", reqID), + zap.String("caller", caller), + zap.String("db_name", c.dbName), + zap.Bool("is_connected", conn.IsConnected()), + zap.Int64("inflight", inflight), + zap.Duration("elapsed", elapsed), + ) + release() + } + return db, wrappedRelease, nil +} + // Engine Infinity engine implementation using Go SDK type infinityEngine struct { config *server.InfinityConfig @@ -198,10 +397,15 @@ func (e *infinityEngine) SupportsPageRank() bool { // Ping checks if Infinity is accessible func (e *infinityEngine) Ping(ctx context.Context) error { - if e.client == nil || e.client.conn == nil { + if e.client == nil || e.client.pool == nil { return fmt.Errorf("Infinity client not initialized") } - if !e.client.conn.IsConnected() { + conn, release, err := e.client.checkoutConn(ctx, "Ping") + if err != nil { + return err + } + defer release() + if !conn.IsConnected() { return fmt.Errorf("Infinity not connected") } return nil @@ -209,16 +413,20 @@ func (e *infinityEngine) Ping(ctx context.Context) error { // Close closes the Infinity connection func (e *infinityEngine) Close() error { - if e.client != nil && e.client.conn != nil { - _, err := e.client.conn.Disconnect() - return err + if e.client != nil && e.client.pool != nil { + return e.client.pool.Close() } return nil } // MigrateDB creates the database if it doesn't exist func (e *infinityEngine) MigrateDB(ctx context.Context) error { - _, err := e.client.conn.CreateDatabase(e.client.dbName, infinity.ConflictTypeIgnore, "") + conn, release, err := e.client.checkoutConn(ctx, "MigrateDB") + if err != nil { + return fmt.Errorf("failed to get connection: %w", err) + } + defer release() + _, err = conn.CreateDatabase(e.client.dbName, infinity.ConflictTypeIgnore, "") if err != nil { return fmt.Errorf("failed to create database: %w", err) } diff --git a/internal/engine/infinity/client_test.go b/internal/engine/infinity/client_test.go new file mode 100644 index 0000000000..7cf2c2b262 --- /dev/null +++ b/internal/engine/infinity/client_test.go @@ -0,0 +1,75 @@ +package infinity + +import ( + "context" + "errors" + "fmt" + "testing" + "time" + + "ragflow/internal/common" +) + +// TestPoolExhaustedError_Contract exercises the real poolExhaustedError type +// that checkoutConn returns when the connection pool is saturated. It is the +// exact value the admin PingEngine handler surfaces as HTTP 503, so it must +// keep reporting CodeConnectionPoolExhausted. No live Infinity is required. +func TestPoolExhaustedError_Contract(t *testing.T) { + pe := &poolExhaustedError{caller: "SearchMetadata"} + + if pe.Code() != common.CodeConnectionPoolExhausted { + t.Fatalf("Code() = %v, want %v", pe.Code(), common.CodeConnectionPoolExhausted) + } + + wantMsg := "Infinity connection pool exhausted while SearchMetadata" + if pe.Error() != wantMsg { + t.Fatalf("Error() = %q, want %q", pe.Error(), wantMsg) + } + if pe.Message() != wantMsg { + t.Fatalf("Message() = %q, want %q", pe.Message(), wantMsg) + } + + // When wrapped, errors.As must still extract it and keep the same code. + wrapped := fmt.Errorf("outer: %w", pe) + var extracted *poolExhaustedError + if !errors.As(wrapped, &extracted) { + t.Fatalf("errors.As failed to extract *poolExhaustedError from %v", wrapped) + } + if extracted.Code() != common.CodeConnectionPoolExhausted { + t.Fatalf("extracted Code() = %v, want %v", extracted.Code(), common.CodeConnectionPoolExhausted) + } +} + +// TestEnsureDeadline covers the core of the bounded-latency fix: callers that +// pass context.Background() (no deadline) must get one applied, while callers +// that already carry a deadline are left untouched and their cancel is not +// hijacked by the no-op cancel returned for them. +func TestEnsureDeadline(t *testing.T) { + // nil ctx -> gets a deadline and a real cancel func. + ctx, cancel := ensureDeadline(nil, 50*time.Millisecond) + defer cancel() + if _, ok := ctx.Deadline(); !ok { + t.Fatal("expected deadline on ctx built from nil") + } + + // ctx already carrying a deadline -> returned unchanged with a no-op cancel. + base, baseCancel := context.WithTimeout(context.Background(), time.Second) + defer baseCancel() + got, noop := ensureDeadline(base, 50*time.Millisecond) + if got != base { + t.Fatal("expected same context when a deadline already exists") + } + noop() // must be safe to call and must not cancel the caller's context. + select { + case <-got.Done(): + t.Fatal("no-op cancel must not cancel the caller's context") + default: + } + + // plain Background -> gets a deadline applied. + bctx, bcancel := ensureDeadline(context.Background(), 50*time.Millisecond) + defer bcancel() + if _, ok := bctx.Deadline(); !ok { + t.Fatal("expected deadline on Background-derived ctx") + } +} diff --git a/internal/engine/infinity/common.go b/internal/engine/infinity/common.go index 720e554bbd..97754434c1 100644 --- a/internal/engine/infinity/common.go +++ b/internal/engine/infinity/common.go @@ -35,8 +35,14 @@ func (e *infinityEngine) dropTable(ctx context.Context, tableName string) error return fmt.Errorf("table name cannot be empty") } + db, release, err := e.client.checkoutDatabase(ctx, "common.go") + if err != nil { + return fmt.Errorf("failed to get database: %w", err) + } + defer release() + // Check if table exists - exists, err := e.tableExists(ctx, tableName) + exists, err := e.tableExistsWithDB(db, tableName) if err != nil { return fmt.Errorf("failed to check table existence: %w", err) } @@ -44,11 +50,6 @@ func (e *infinityEngine) dropTable(ctx context.Context, tableName string) error return fmt.Errorf("table '%s' does not exist", tableName) } - db, err := e.client.conn.GetDatabase(e.client.dbName) - if err != nil { - return fmt.Errorf("failed to get database: %w", err) - } - _, err = db.DropTable(tableName, infinity.ConflictTypeError) if err != nil { return fmt.Errorf("failed to drop table: %w", err) @@ -64,13 +65,22 @@ func (e *infinityEngine) tableExists(ctx context.Context, tableName string) (boo return false, fmt.Errorf("table name cannot be empty") } - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "common.go") if err != nil { return false, fmt.Errorf("failed to get database: %w", err) } + defer release() + + return e.tableExistsWithDB(db, tableName) +} + +func (e *infinityEngine) tableExistsWithDB(db *infinity.Database, tableName string) (bool, error) { + if db == nil { + return false, fmt.Errorf("database is nil") + } // Try to get the table - if it exists, no error - _, err = db.GetTable(tableName) + _, err := db.GetTable(tableName) if err != nil { errMsg := strings.ToLower(err.Error()) if strings.Contains(errMsg, "not found") || strings.Contains(errMsg, "doesn't exist") { diff --git a/internal/engine/infinity/document.go b/internal/engine/infinity/document.go index 1cd1ee8ce2..10369377d8 100644 --- a/internal/engine/infinity/document.go +++ b/internal/engine/infinity/document.go @@ -39,10 +39,11 @@ func (e *infinityEngine) IndexDocument(ctx context.Context, tableName, docID str // InsertSkill inserts a skill document into skill index // Auto-creates the table if it doesn't exist func (e *infinityEngine) InsertSkill(ctx context.Context, tableName, docID string, doc interface{}) error { - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "document.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -107,10 +108,11 @@ func (e *infinityEngine) BulkIndex(ctx context.Context, tableName string, docs [ // matching the behavior of InsertSkill. Creates shallow copies of input maps to // avoid mutating caller data. func (e *infinityEngine) BulkInsertSkill(ctx context.Context, tableName string, docs []interface{}) (int, error) { - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "document.go") if err != nil { return 0, fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -205,10 +207,11 @@ func (e *infinityEngine) DeleteDocument(ctx context.Context, tableName, docID st return fmt.Errorf("document id cannot be empty") } - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "document.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { diff --git a/internal/engine/infinity/metadata.go b/internal/engine/infinity/metadata.go index a146458a52..79bd26516a 100644 --- a/internal/engine/infinity/metadata.go +++ b/internal/engine/infinity/metadata.go @@ -36,16 +36,20 @@ import ( // CreateMetadataStore creates a metadata table in Infinity // tenantID is the tenant identifier used to build the table name func (e *infinityEngine) CreateMetadataStore(ctx context.Context, tenantID string) error { - tableName := buildMetadataTableName(tenantID) - - // Get database - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() + + return e.createMetadataStoreWithDB(db, tenantID) +} + +func (e *infinityEngine) createMetadataStoreWithDB(db *infinity.Database, tenantID string) error { + tableName := buildMetadataTableName(tenantID) // Check if table already exists - exists, err := e.tableExists(ctx, tableName) + exists, err := e.tableExistsWithDB(db, tableName) if err != nil { return fmt.Errorf("failed to check if table exists: %w", err) } @@ -136,10 +140,11 @@ func (e *infinityEngine) InsertMetadata(ctx context.Context, metadata []map[stri tableName := buildMetadataTableName(tenantID) common.Info("InfinityConnection.InsertMetadata called", zap.String("tableName", tableName), zap.Int("metaCount", len(metadata))) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return nil, fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -150,7 +155,7 @@ func (e *infinityEngine) InsertMetadata(ctx context.Context, metadata []map[stri } // Create metadata table - if createErr := e.CreateMetadataStore(ctx, tenantID); createErr != nil { + if createErr := e.createMetadataStoreWithDB(db, tenantID); createErr != nil { return nil, fmt.Errorf("failed to create metadata table: %w", createErr) } @@ -233,10 +238,11 @@ func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, datas tableName := buildMetadataTableName(tenantID) common.Info("InfinityConnection.UpdateMetadata called", zap.String("tableName", tableName), zap.String("docID", docID), zap.String("datasetID", datasetID)) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -338,10 +344,11 @@ func (e *infinityEngine) UpdateMetadata(ctx context.Context, docID string, datas func (e *infinityEngine) DeleteMetadata(ctx context.Context, condition map[string]interface{}, tenantID string) (int64, error) { tableName := buildMetadataTableName(tenantID) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return 0, fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -349,6 +356,14 @@ func (e *infinityEngine) DeleteMetadata(ctx context.Context, condition map[strin return 0, nil } + return e.deleteMetadataWithTable(table, condition) +} + +func (e *infinityEngine) deleteMetadataWithTable(table *infinity.Table, condition map[string]interface{}) (int64, error) { + if table == nil { + return 0, fmt.Errorf("metadata table is nil") + } + // Get table columns for building filter clmns := make(map[string]struct { Type string @@ -397,10 +412,11 @@ func (e *infinityEngine) DeleteMetadataKeys(ctx context.Context, docID string, d tableName := buildMetadataTableName(tenantID) common.Info("InfinityConnection.DeleteMetadataKeys called", zap.String("tableName", tableName), zap.String("docID", docID), zap.Any("keys", keys)) - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return fmt.Errorf("failed to get database: %w", err) } + defer release() table, err := db.GetTable(tableName) if err != nil { @@ -500,7 +516,7 @@ func (e *infinityEngine) DeleteMetadataKeys(ctx context.Context, docID string, d "id": docID, "kb_id": datasetID, } - _, err := e.DeleteMetadata(ctx, condition, tenantID) + _, err := e.deleteMetadataWithTable(table, condition) if err != nil { return fmt.Errorf("failed to delete document: %w", err) } @@ -554,23 +570,6 @@ func (e *infinityEngine) SearchMetadata(ctx context.Context, req *types.SearchMe // Build table name from tenantID tableName := buildMetadataTableName(tenantID) - exists, err := e.tableExists(ctx, tableName) - if err != nil { - common.Warn("Infinity SearchMetadata table existence check failed", zap.String("table", tableName), zap.Error(err)) - return nil, fmt.Errorf("failed to check metadata table existence: %w", err) - } - if !exists { - common.Debug("Infinity SearchMetadata table absent, returning empty result", zap.String("table", tableName)) - // Return an empty (non-nil) slice — Python returns `[]`, and a - // nil slice is read by callers as "fall back to in-memory". A - // zero-match against an absent table is a definitive answer, - // not a missing-data condition. - return &types.SearchMetadataResult{ - MetadataRecords: []map[string]interface{}{}, - Total: 0, - }, nil - } - // Build output columns: use caller-specified fields, or "*" for all columns var outputColumns []string if len(req.SelectFields) > 0 { @@ -596,10 +595,28 @@ func (e *infinityEngine) SearchMetadata(ctx context.Context, req *types.SearchMe } // Get database and table - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil { return nil, fmt.Errorf("failed to get database: %w", err) } + defer release() + + exists, err := e.tableExistsWithDB(db, tableName) + if err != nil { + common.Warn("Infinity SearchMetadata table existence check failed", zap.String("table", tableName), zap.Error(err)) + return nil, fmt.Errorf("failed to check metadata table existence: %w", err) + } + if !exists { + common.Debug("Infinity SearchMetadata table absent, returning empty result", zap.String("table", tableName)) + // Return an empty (non-nil) slice — Python returns `[]`, and a + // nil slice is read by callers as "fall back to in-memory". A + // zero-match against an absent table is a definitive answer, + // not a missing-data condition. + return &types.SearchMetadataResult{ + MetadataRecords: []map[string]interface{}{}, + Total: 0, + }, nil + } tbl, err := db.GetTable(tableName) if err != nil { @@ -822,10 +839,11 @@ func (e *infinityEngine) FilterDocIdsByMetaPushdown(ctx context.Context, kbIDs [ whereClause = kbFilter + " AND (" + whereClause + ")" // Use Infinity connection to execute query - db, err := e.client.conn.GetDatabase(e.client.dbName) + db, release, err := e.client.checkoutDatabase(ctx, "metadata.go") if err != nil || db == nil { return nil } + defer release() table, err := db.GetTable(tableName) if err != nil || table == nil {