From 86eabb7224c251522d62ea95d73d771b6f204b48 Mon Sep 17 00:00:00 2001 From: zipg Date: Thu, 10 Sep 2026 12:45:02 +0800 Subject: [PATCH] fix(vastbase): reset timeout for each query page --- .../drivers/vastbase-go/integration_test.go | 11 ++ agents/drivers/vastbase-go/main.go | 71 ++++++++++- agents/drivers/vastbase-go/main_test.go | 120 ++++++++++++++++++ 3 files changed, 200 insertions(+), 2 deletions(-) diff --git a/agents/drivers/vastbase-go/integration_test.go b/agents/drivers/vastbase-go/integration_test.go index 348ec31b16..85e4746c51 100644 --- a/agents/drivers/vastbase-go/integration_test.go +++ b/agents/drivers/vastbase-go/integration_test.go @@ -142,6 +142,17 @@ func TestVastbaseIntegration(t *testing.T) { t.Fatalf("third page failed: page=%v err=%v", third, err) } + idlePage, err := server.executeQueryPage(queryOptions{SQL: "SELECT generate_series(1, 3)", MaxRows: 3, TimeoutSecs: 1}, 1) + if err != nil || idlePage.SessionID == nil || !idlePage.HasMore || len(idlePage.Rows) != 1 { + t.Fatalf("idle pagination first page failed: page=%v err=%v", idlePage, err) + } + time.Sleep(1100 * time.Millisecond) + idleSecond, err := server.fetchQueryPage(*idlePage.SessionID, 1) + if err != nil || !idleSecond.HasMore || len(idleSecond.Rows) != 1 { + t.Fatalf("idle time consumed the next page timeout: page=%v err=%v", idleSecond, err) + } + server.closeQuerySession(*idlePage.SessionID) + cancelStart := time.Now() cancelResult := make(chan error, 1) go func() { diff --git a/agents/drivers/vastbase-go/main.go b/agents/drivers/vastbase-go/main.go index 1fe5c23587..0440e4b521 100644 --- a/agents/drivers/vastbase-go/main.go +++ b/agents/drivers/vastbase-go/main.go @@ -126,9 +126,15 @@ type querySession struct { pending []any pendingSpatial []*uint32 remaining int + timeoutSecs int cancel context.CancelFunc } +type queryPageTimeout struct { + timer *time.Timer + fired chan struct{} +} + type rowScanner struct { values []any destinations []any @@ -719,6 +725,55 @@ func (s *server) queryRows(sqlText string, schema string, timeoutSecs int) (*sql return rows, conn, cancel, nil } +func (s *server) queryRowsForPage(sqlText string, schema string, timeoutSecs int) (*sql.Rows, *sql.Conn, context.CancelFunc, *queryPageTimeout, error) { + ctx, cancel := s.beginOperation(0) + timeout := startQueryPageTimeout(timeoutSecs, cancel) + conn, err := s.schemaConn(ctx, schema) + if err != nil { + timedOut := timeout.stop() + s.endOperation(cancel) + if timedOut { + return nil, nil, nil, nil, context.DeadlineExceeded + } + return nil, nil, nil, nil, err + } + rows, err := conn.QueryContext(ctx, sqlText) + if err != nil { + timedOut := timeout.stop() + _ = conn.Close() + s.endOperation(cancel) + if timedOut { + return nil, nil, nil, nil, context.DeadlineExceeded + } + return nil, nil, nil, nil, err + } + return rows, conn, cancel, timeout, nil +} + +func startQueryPageTimeout(timeoutSecs int, cancel context.CancelFunc) *queryPageTimeout { + if timeoutSecs <= 0 { + return nil + } + fired := make(chan struct{}) + timeout := &queryPageTimeout{fired: fired} + timeout.timer = time.AfterFunc(time.Duration(timeoutSecs)*time.Second, func() { + cancel() + close(fired) + }) + return timeout +} + +func (timeout *queryPageTimeout) stop() bool { + if timeout == nil { + return false + } + if timeout.timer.Stop() { + return false + } + <-timeout.fired + return true +} + func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageResult, error) { start := time.Now() sqlText := trimStatementSQL(opts.SQL) @@ -727,15 +782,19 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes result, err := s.executeQuery(opts) return queryPageResult{Columns: result.Columns, ColumnTypes: result.ColumnTypes, SpatialColumns: result.SpatialColumns, SpatialValues: result.SpatialValues, Rows: result.Rows, AffectedRows: result.AffectedRows, ExecutionTimeMS: result.ExecutionTimeMS, Truncated: result.Truncated}, err } - rows, conn, cancel, err := s.queryRows(sqlText, opts.Schema, opts.TimeoutSecs) + rows, conn, cancel, timeout, err := s.queryRowsForPage(sqlText, opts.Schema, opts.TimeoutSecs) if err != nil { return queryPageResult{}, err } columns, err := rows.Columns() if err != nil { + timedOut := timeout.stop() _ = rows.Close() _ = conn.Close() s.endOperation(cancel) + if timedOut { + return queryPageResult{}, context.DeadlineExceeded + } return queryPageResult{}, err } maxRows := opts.MaxRows @@ -743,9 +802,13 @@ func (s *server) executeQueryPage(opts queryOptions, pageSize int) (queryPageRes maxRows = defaultMaxRows } columnTypes := columnTypeNames(rows) - session := &querySession{rows: rows, conn: conn, columns: columns, columnTypes: columnTypes, scanner: newRowScanner(len(columns), newSpatialDecoder(columnTypes)), remaining: maxRows, cancel: cancel} + session := &querySession{rows: rows, conn: conn, columns: columns, columnTypes: columnTypes, scanner: newRowScanner(len(columns), newSpatialDecoder(columnTypes)), remaining: maxRows, timeoutSecs: opts.TimeoutSecs, cancel: cancel} result, err := readQuerySessionPage(session, pageSize) + timedOut := timeout.stop() result.ExecutionTimeMS = time.Since(start).Milliseconds() + if timedOut { + err = context.DeadlineExceeded + } if err != nil { _ = rows.Close() _ = conn.Close() @@ -770,7 +833,11 @@ func (s *server) fetchQueryPage(id string, pageSize int) (queryPageResult, error if session == nil { return queryPageResult{Columns: []string{}, ColumnTypes: []string{}, Rows: [][]any{}}, nil } + timeout := startQueryPageTimeout(session.timeoutSecs, session.cancel) result, err := readQuerySessionPage(session, pageSize) + if timeout.stop() { + err = context.DeadlineExceeded + } if err != nil { s.closeQuerySession(id) return queryPageResult{}, err diff --git a/agents/drivers/vastbase-go/main_test.go b/agents/drivers/vastbase-go/main_test.go index 675953b3ae..9bb31d7965 100644 --- a/agents/drivers/vastbase-go/main_test.go +++ b/agents/drivers/vastbase-go/main_test.go @@ -5,10 +5,13 @@ import ( "database/sql" "database/sql/driver" "encoding/json" + "errors" + "io" "strings" "sync" "sync/atomic" "testing" + "time" pq "gitcode.com/opengauss/openGauss-connector-go-pq" ) @@ -265,6 +268,74 @@ func TestValidateConnectionRecoversAfterCanceledDriverConnection(t *testing.T) { } } +func TestPagedQueryTimeoutDoesNotExpireWhilePageIsIdle(t *testing.T) { + registerVastbasePaginationDriver.Do(func() { + sql.Register("vastbase-pagination-timeout-test", &paginationTimeoutDriver{}) + }) + db, err := sql.Open("vastbase-pagination-timeout-test", "") + if err != nil { + t.Fatal(err) + } + defer db.Close() + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + + server := newServer() + server.db = db + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT id FROM rows", MaxRows: 3, TimeoutSecs: 1}, 1) + if err != nil { + t.Fatal(err) + } + if first.SessionID == nil || !first.HasMore || len(first.Rows) != 1 { + t.Fatalf("unexpected first page: %#v", first) + } + + time.Sleep(1100 * time.Millisecond) + second, err := server.fetchQueryPage(*first.SessionID, 1) + if err != nil { + t.Fatalf("idle time must not consume the next page timeout: %v", err) + } + if !second.HasMore || len(second.Rows) != 1 || second.Rows[0][0] != int64(2) { + t.Fatalf("unexpected second page: %#v", second) + } + server.closeAllQuerySessions() +} + +func TestPagedQueryTimeoutStillAppliesToEachFetch(t *testing.T) { + registerVastbasePaginationDriver.Do(func() { + sql.Register("vastbase-pagination-timeout-test", &paginationTimeoutDriver{}) + }) + db, err := sql.Open("vastbase-pagination-timeout-test", "") + if err != nil { + t.Fatal(err) + } + defer db.Close() + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + + server := newServer() + server.db = db + first, err := server.executeQueryPage(queryOptions{SQL: "SELECT id FROM slow_rows", MaxRows: 3, TimeoutSecs: 1}, 1) + if err != nil { + t.Fatal(err) + } + if first.SessionID == nil { + t.Fatalf("expected a paged query session: %#v", first) + } + + started := time.Now() + _, err = server.fetchQueryPage(*first.SessionID, 1) + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("expected fetch timeout, got %v", err) + } + if elapsed := time.Since(started); elapsed < time.Second || elapsed > 3*time.Second { + t.Fatalf("fetch timeout fired outside the expected window: %s", elapsed) + } + if len(server.sessions) != 0 { + t.Fatalf("timed-out query session was retained: %#v", server.sessions) + } +} + func TestDisconnectResetsInformationSchemaCapabilityCache(t *testing.T) { server := newServer() server.infoColumnTypeUnsupported = true @@ -293,10 +364,59 @@ func TestDisconnectResetsConstraintCapabilityCache(t *testing.T) { var ( registerVastbaseSchemaRetryDriver sync.Once registerVastbasePingRetryDriver sync.Once + registerVastbasePaginationDriver sync.Once schemaRetryOpens atomic.Int32 pingRetryOpens atomic.Int32 ) +type paginationTimeoutDriver struct{} + +func (*paginationTimeoutDriver) Open(string) (driver.Conn, error) { + return &paginationTimeoutConn{}, nil +} + +type paginationTimeoutConn struct{} + +func (*paginationTimeoutConn) Prepare(string) (driver.Stmt, error) { return nil, driver.ErrSkip } +func (*paginationTimeoutConn) Close() error { return nil } +func (*paginationTimeoutConn) Begin() (driver.Tx, error) { return nil, driver.ErrSkip } + +func (*paginationTimeoutConn) QueryContext(ctx context.Context, query string, _ []driver.NamedValue) (driver.Rows, error) { + rows := &paginationTimeoutRows{ctx: ctx, values: []int64{1, 2, 3}, blockAt: -1} + if strings.Contains(query, "slow_rows") { + rows.blockAt = 2 + } + return rows, nil +} + +type paginationTimeoutRows struct { + ctx context.Context + values []int64 + index int + blockAt int +} + +func (*paginationTimeoutRows) Columns() []string { return []string{"id"} } +func (*paginationTimeoutRows) Close() error { return nil } + +func (rows *paginationTimeoutRows) Next(dest []driver.Value) error { + if rows.index == rows.blockAt { + <-rows.ctx.Done() + return rows.ctx.Err() + } + select { + case <-rows.ctx.Done(): + return rows.ctx.Err() + default: + } + if rows.index >= len(rows.values) { + return io.EOF + } + dest[0] = rows.values[rows.index] + rows.index++ + return nil +} + type schemaRetryDriver struct{} func (*schemaRetryDriver) Open(string) (driver.Conn, error) {