fix(vastbase): reset timeout for each query page

2026-10-21 06:00:39 +00:00 · 2026-09-10 12:45:02 +08:00
parent d48167ef66
commit 86eabb7224
3 changed files with 200 additions and 2 deletions
@@ -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() {
@@ -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
@@ -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) {