feat: support statement_timeout and transaction_timeout property

Add a statement_timeout connection property that is used as the default timeout
for all statements that are executed on a connection. If both a context timeout
and a statement_timeout is specified, then the lower of the two will be used.

Also adds a transaction_timeout property that is additionally used for all
statements in a read/write transaction. The deadline of the transaction is
calculated at the start of the transaction, and all statements in the transaction
get this deadline, unless the statement already has an earlier deadline from
for example a statement_timeout or a context deadline.

Fixes #574
Fixes #575

Author: olavloite

conn.go

@@ -831,6 +831,51 @@ func (c *conn) PrepareContext(_ context.Context, query string) (driver.Stmt, err
 	return &stmt{conn: c, query: parsedSQL, numArgs: len(args), execOptions: execOptions}, nil
 }
 
+func (c *conn) addStatementAndTransactionTimeout(ctx context.Context) (context.Context, context.CancelFunc, error) {
+	var statementDeadline time.Time
+	var transactionDeadline time.Time
+	var deadline time.Time
+	var hasStatementDeadline bool
+	var hasTransactionDeadline bool
+
+	statementTimeout := propertyStatementTimeout.GetValueOrDefault(c.state)
+	if statementTimeout != time.Duration(0) {
+		hasStatementDeadline = true
+		statementDeadline = time.Now().Add(statementTimeout)
+	}
+	transactionDeadline, hasTransactionDeadline, err := c.transactionDeadline()
+	if err != nil {
+		return nil, nil, err
+	}
+	if !hasStatementDeadline && !hasTransactionDeadline {
+		return ctx, func() {}, nil
+	}
+	if hasTransactionDeadline && hasStatementDeadline {
+		if statementDeadline.Before(transactionDeadline) {
+			deadline = statementDeadline
+		} else {
+			deadline = transactionDeadline
+		}
+	} else if hasStatementDeadline {
+		deadline = statementDeadline
+	} else {
+		deadline = transactionDeadline
+	}
+	newCtx, cancel := context.WithDeadline(ctx, deadline)
+	return newCtx, cancel, nil
+}
+
+func (c *conn) transactionDeadline() (time.Time, bool, error) {
+	if c.tx == nil {
+		return time.Time{}, false, nil
+	}
+	if err := c.tx.ensureActivated(); err != nil {
+		return time.Time{}, false, err
+	}
+	deadline, hasDeadline := c.tx.deadline()
+	return deadline, hasDeadline, nil
+}
+
 func (c *conn) QueryContext(ctx context.Context, query string, args []driver.NamedValue) (driver.Rows, error) {
 	// Execute client side statement if it is one.
 	clientStmt, err := c.parser.ParseClientSideStatement(query)
@@ -849,13 +894,22 @@ func (c *conn) QueryContext(ctx context.Context, query string, args []driver.Nam
 	return c.queryContext(ctx, query, execOptions, args)
 }
 
-func (c *conn) queryContext(ctx context.Context, query string, execOptions *ExecOptions, args []driver.NamedValue) (driver.Rows, error) {
+func (c *conn) queryContext(ctx context.Context, query string, execOptions *ExecOptions, args []driver.NamedValue) (returnedRows driver.Rows, returnedErr error) {
+	ctx, cancel, err := c.addStatementAndTransactionTimeout(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if returnedErr != nil {
+			cancel()
+		}
+	}()
 	// Clear the commit timestamp of this connection before we execute the query.
 	c.clearCommitResponse()
 	// Check if the execution options contains an instruction to execute
 	// a specific partition of a PartitionedQuery.
 	if pq := execOptions.PartitionedQueryOptions.ExecutePartition.PartitionedQuery; pq != nil {
-		return pq.execute(ctx, execOptions.PartitionedQueryOptions.ExecutePartition.Index)
+		return pq.execute(ctx, cancel, execOptions.PartitionedQueryOptions.ExecutePartition.Index)
 	}
 
 	stmt, err := prepareSpannerStmt(c.parser, query, args)
@@ -869,7 +923,7 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions *Exec
 		if err != nil {
 			return nil, err
 		}
-		return createDriverResultRows(res, execOptions), nil
+		return createDriverResultRows(res, cancel, execOptions), nil
 	}
 	var iter rowIterator
 	if c.tx == nil {
@@ -884,7 +938,7 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions *Exec
 		} else if execOptions.PartitionedQueryOptions.PartitionQuery {
 			return nil, spanner.ToSpannerError(status.Errorf(codes.FailedPrecondition, "PartitionQuery is only supported in batch read-only transactions"))
 		} else if execOptions.PartitionedQueryOptions.AutoPartitionQuery {
-			return c.executeAutoPartitionedQuery(ctx, query, execOptions, args)
+			return c.executeAutoPartitionedQuery(ctx, cancel, query, execOptions, args)
 		} else {
 			// The statement was either detected as being a query, or potentially not recognized at all.
 			// In that case, just default to using a single-use read-only transaction and let Spanner
@@ -893,14 +947,15 @@ func (c *conn) queryContext(ctx context.Context, query string, execOptions *Exec
 		}
 	} else {
 		if execOptions.PartitionedQueryOptions.PartitionQuery {
+			defer cancel()
 			return c.tx.partitionQuery(ctx, stmt, execOptions)
 		}
 		iter, err = c.tx.Query(ctx, stmt, statementInfo.StatementType, execOptions)
 		if err != nil {
 			return nil, err
 		}
 	}
-	res := createRows(iter, execOptions)
+	res := createRows(iter, cancel, execOptions)
 	if execOptions.DirectExecuteQuery {
 		// This call to res.getColumns() triggers the execution of the statement, as it needs to fetch the metadata.
 		res.getColumns()
@@ -929,7 +984,16 @@ func (c *conn) ExecContext(ctx context.Context, query string, args []driver.Name
 	return c.execContext(ctx, query, execOptions, args)
 }
 
-func (c *conn) execContext(ctx context.Context, query string, execOptions *ExecOptions, args []driver.NamedValue) (driver.Result, error) {
+func (c *conn) execContext(ctx context.Context, query string, execOptions *ExecOptions, args []driver.NamedValue) (returnedResult driver.Result, returnedErr error) {
+	ctx, cancel, err := c.addStatementAndTransactionTimeout(ctx)
+	if err != nil {
+		return nil, err
+	}
+	defer func() {
+		if returnedErr != nil {
+			cancel()
+		}
+	}()
 	// Clear the commit timestamp of this connection before we execute the statement.
 	c.clearCommitResponse()
 
@@ -1041,6 +1105,17 @@ func (c *conn) resetTransactionForRetry(ctx context.Context, errDuringCommit boo
 			return noTransaction()
 		}
 		c.tx = c.prevTx
+		// If the aborted error happened during the Commit, then the transaction context has been cancelled,
+		// and we need to create a new one.
+		if rwTx, ok := c.tx.contextTransaction.(*readWriteTransaction); ok {
+			newCtx, cancel := c.addTransactionTimeout(c.tx.ctx)
+			rwTx.ctx = newCtx
+			origClose := rwTx.close
+			rwTx.close = func(result txResult, commitResponse *spanner.CommitResponse, commitErr error) {
+				origClose(result, commitResponse, commitErr)
+				cancel()
+			}
+		}
 		c.resetForRetry = true
 	} else if c.tx == nil {
 		return noTransaction()
@@ -1248,6 +1323,14 @@ func (c *conn) beginTx(ctx context.Context, driverOpts driver.TxOptions, closeFu
 	return c.tx, nil
 }
 
+func (c *conn) addTransactionTimeout(ctx context.Context) (context.Context, context.CancelFunc) {
+	timeout := propertyTransactionTimeout.GetValueOrDefault(c.state)
+	if timeout == time.Duration(0) {
+		return ctx, func() {}
+	}
+	return context.WithTimeout(ctx, timeout)
+}
+
 func (c *conn) activateTransaction() (contextTransaction, error) {
 	closeFunc := c.tx.close
 	if propertyTransactionReadOnly.GetValueOrDefault(c.state) {
@@ -1283,19 +1366,21 @@ func (c *conn) activateTransaction() (contextTransaction, error) {
 	opts := spanner.TransactionOptions{}
 	opts.BeginTransactionOption = c.convertDefaultBeginTransactionOption(propertyBeginTransactionOption.GetValueOrDefault(c.state))
 
-	tx, err := spanner.NewReadWriteStmtBasedTransactionWithCallbackForOptions(c.tx.ctx, c.client, opts, func() spanner.TransactionOptions {
+	ctx, cancel := c.addTransactionTimeout(c.tx.ctx)
+	tx, err := spanner.NewReadWriteStmtBasedTransactionWithCallbackForOptions(ctx, c.client, opts, func() spanner.TransactionOptions {
 		defer func() {
 			// Reset the transaction_tag after starting the transaction.
 			_ = propertyTransactionTag.ResetValue(c.state, connectionstate.ContextUser)
 		}()
 		return c.effectiveTransactionOptions(spannerpb.TransactionOptions_ISOLATION_LEVEL_UNSPECIFIED, c.options( /*reset=*/ true))
 	})
 	if err != nil {
+		cancel()
 		return nil, err
 	}
 	logger := c.logger.With("tx", "rw")
 	return &readWriteTransaction{
-		ctx:    c.tx.ctx,
+		ctx:    ctx,
 		conn:   c,
 		logger: logger,
 		rwTx:   tx,
@@ -1307,6 +1392,7 @@ func (c *conn) activateTransaction() (contextTransaction, error) {
 			} else {
 				closeFunc(txResultRollback)
 			}
+			cancel()
 		},
 		retryAborts: sync.OnceValue(func() bool {
 			return c.RetryAbortsInternally()
@@ -1371,18 +1457,20 @@ func queryInSingleUse(ctx context.Context, c *spanner.Client, statement spanner.
 	return c.Single().WithTimestampBound(tb).QueryWithOptions(ctx, statement, options.QueryOptions)
 }
 
-func (c *conn) executeAutoPartitionedQuery(ctx context.Context, query string, execOptions *ExecOptions, args []driver.NamedValue) (driver.Rows, error) {
+func (c *conn) executeAutoPartitionedQuery(ctx context.Context, cancel context.CancelFunc, query string, execOptions *ExecOptions, args []driver.NamedValue) (driver.Rows, error) {
 	tx, err := c.BeginTx(ctx, driver.TxOptions{ReadOnly: true, Isolation: withBatchReadOnly(driver.IsolationLevel(sql.LevelDefault))})
 	if err != nil {
 		return nil, err
 	}
 	r, err := c.queryContext(ctx, query, execOptions, args)
 	if err != nil {
 		_ = tx.Rollback()
+		cancel()
 		return nil, err
 	}
 	if rows, ok := r.(*rows); ok {
 		rows.close = func() error {
+			defer cancel()
 			return tx.Commit()
 		}
 	}

connection_leak_test.go

@@ -24,22 +24,23 @@ import (
 
 	"cloud.google.com/go/spanner"
 	"github.com/googleapis/go-sql-spanner/testutil"
+	"go.uber.org/goleak"
 	"google.golang.org/grpc/codes"
 	gstatus "google.golang.org/grpc/status"
 )
 
 func TestNoLeak(t *testing.T) {
-	t.Parallel()
+	// Not parallel, as it checks for leaked goroutines.
 
-	db, server, teardown := setupTestDBConnection(t)
+	db, server, teardown := setupTestDBConnectionWithParams(t, "statement_timeout=10s;transaction_timeout=20s")
 	defer teardown()
 	// Set MaxOpenConns to 1 to force an error if anything leaks a connection.
 	db.SetMaxOpenConns(1)
 
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
 	defer cancel()
 
-	for i := 0; i < 2; i++ {
+	runTests := func() {
 		pingContext(ctx, t, db)
 		pingFailed(ctx, t, server, db)
 		simpleQuery(ctx, t, db)
@@ -50,8 +51,28 @@ func TestNoLeak(t *testing.T) {
 		readOnlyTxWithStaleness(ctx, t, db)
 		simpleReadWriteTx(ctx, t, db)
 		runTransactionRetry(ctx, t, server, db)
+		runTransactionRetryAbortedHalfway(ctx, t, server, db)
 		readOnlyTxWithOptions(ctx, t, db)
 	}
+
+	for i := 0; i < 2; i++ {
+		runTests()
+	}
+	ignoreCurrent := goleak.IgnoreCurrent()
+
+	for i := 0; i < 10; i++ {
+		runTests()
+	}
+	goleak.VerifyNone(t, ignoreCurrent,
+		goleak.IgnoreTopFunction("cloud.google.com/go/spanner.(*healthChecker).worker"),
+		goleak.IgnoreTopFunction("cloud.google.com/go/spanner.(*healthChecker).multiplexSessionWorker"),
+		goleak.IgnoreTopFunction("cloud.google.com/go/spanner.(*healthChecker).maintainer"),
+		goleak.IgnoreTopFunction("google.golang.org/grpc/internal/transport.(*controlBuffer).get"),
+		goleak.IgnoreTopFunction("google.golang.org/grpc/internal/transport.(*http2Server).keepalive"),
+		goleak.IgnoreTopFunction("internal/poll.runtime_pollWait"),
+		goleak.IgnoreTopFunction("google.golang.org/grpc/internal/grpcsync.(*CallbackSerializer).run"),
+		goleak.IgnoreTopFunction("cloud.google.com/go/spanner.(*sessionPool).createMultiplexedSession"),
+	)
 }
 
 func pingContext(ctx context.Context, t *testing.T, db *sql.DB) {
@@ -308,6 +329,50 @@ func runTransactionRetry(ctx context.Context, t *testing.T, server *testutil.Moc
 	}
 }
 
+func runTransactionRetryAbortedHalfway(ctx context.Context, t *testing.T, server *testutil.MockedSpannerInMemTestServer, db *sql.DB) {
+	var attempts int
+	err := RunTransactionWithOptions(ctx, db, &sql.TxOptions{}, func(ctx context.Context, tx *sql.Tx) error {
+		attempts++
+		rows, err := tx.QueryContext(ctx, testutil.SelectFooFromBar, ExecOptions{QueryOptions: spanner.QueryOptions{RequestTag: "tag_1"}})
+		if err != nil {
+			t.Fatal(err)
+		}
+		for rows.Next() {
+		}
+		if err := rows.Close(); err != nil {
+			t.Fatal(err)
+		}
+
+		if attempts == 1 {
+			server.TestSpanner.PutExecutionTime(testutil.MethodExecuteStreamingSql, testutil.SimulatedExecutionTime{
+				Errors: []error{gstatus.Error(codes.Aborted, "Aborted")},
+			})
+		}
+		if _, err := tx.ExecContext(ctx, testutil.UpdateBarSetFoo, ExecOptions{QueryOptions: spanner.QueryOptions{RequestTag: "tag_2"}}); err != nil {
+			return err
+		}
+
+		if attempts == 2 {
+			server.TestSpanner.PutExecutionTime(testutil.MethodExecuteBatchDml, testutil.SimulatedExecutionTime{
+				Errors: []error{gstatus.Error(codes.Aborted, "Aborted")},
+			})
+		}
+		if _, err := tx.ExecContext(ctx, "start batch dml", ExecOptions{QueryOptions: spanner.QueryOptions{RequestTag: "tag_3"}}); err != nil {
+			return err
+		}
+		if _, err := tx.ExecContext(ctx, testutil.UpdateBarSetFoo); err != nil {
+			return err
+		}
+		if _, err := tx.ExecContext(ctx, "run batch"); err != nil {
+			return err
+		}
+		return nil
+	}, spanner.TransactionOptions{TransactionTag: "my_transaction_tag"})
+	if err != nil {
+		t.Fatalf("failed to run transaction: %v", err)
+	}
+}
+
 func readOnlyTxWithOptions(ctx context.Context, t *testing.T, db *sql.DB) {
 	tx, err := BeginReadOnlyTransaction(ctx, db,
 		ReadOnlyTransactionOptions{TimestampBound: spanner.ExactStaleness(10 * time.Second)})

connection_properties.go

@@ -305,6 +305,16 @@ var propertyTransactionBatchReadOnly = createConnectionProperty(
 	connectionstate.ContextUser,
 	connectionstate.ConvertBool,
 )
+var propertyTransactionTimeout = createConnectionProperty(
+	"transaction_timeout",
+	"The timeout to apply to all read/write transactions on this connection. "+
+		"Setting the timeout to zero means no timeout.",
+	time.Duration(0),
+	false,
+	nil,
+	connectionstate.ContextUser,
+	connectionstate.ConvertDuration,
+)
 
 // ------------------------------------------------------------------------------------------------
 // Statement connection properties.
@@ -318,6 +328,17 @@ var propertyStatementTag = createConnectionProperty(
 	connectionstate.ContextUser,
 	connectionstate.ConvertString,
 )
+var propertyStatementTimeout = createConnectionProperty(
+	"statement_timeout",
+	"The timeout to apply to all statements on this connection. "+
+		"Setting the timeout to zero means no timeout. "+
+		"Any existing context deadline will take precedence over this statement timeout.",
+	time.Duration(0),
+	false,
+	nil,
+	connectionstate.ContextUser,
+	connectionstate.ConvertDuration,
+)
 
 // ------------------------------------------------------------------------------------------------
 // Startup connection properties.

driver.go

@@ -1022,6 +1022,9 @@ func runTransactionWithOptions(ctx context.Context, db *sql.DB, opts *sql.TxOpti
 		return nil, err
 	}
 	for {
+		// Create a derived context for each new attempt
+		// ctx, cancel := context.WithCancel(ctx)
+
 		err = protected(ctx, tx, f)
 		errDuringCommit := false
 		if err == nil {

driver_with_mockserver_test.go

@@ -4572,9 +4572,9 @@ func TestRunTransaction(t *testing.T) {
 		// Verify that internal retries are disabled during RunTransaction
 		txi := reflect.ValueOf(tx).Elem().FieldByName("txi")
 		delegatingTx := (*delegatingTransaction)(txi.Elem().UnsafePointer())
-		rwTx := delegatingTx.contextTransaction.(*readWriteTransaction)
+		rwTx, ok := delegatingTx.contextTransaction.(*readWriteTransaction)
 		// Verify that getting the transaction through reflection worked.
-		if g, w := rwTx.ctx, ctx; g != w {
+		if !ok {
 			return fmt.Errorf("getting the transaction through reflection failed")
 		}
 		if rwTx.retryAborts() {
@@ -5034,9 +5034,9 @@ func TestBeginReadWriteTransaction(t *testing.T) {
 		// Verify that internal retries are disabled during this transaction.
 		txi := reflect.ValueOf(tx).Elem().FieldByName("txi")
 		delegatingTx := (*delegatingTransaction)(txi.Elem().UnsafePointer())
-		rwTx := delegatingTx.contextTransaction.(*readWriteTransaction)
+		rwTx, ok := delegatingTx.contextTransaction.(*readWriteTransaction)
 		// Verify that getting the transaction through reflection worked.
-		if g, w := rwTx.ctx, ctx; g != w {
+		if !ok {
 			t.Fatal("getting the transaction through reflection failed")
 		}
 		if rwTx.retryAborts() {