perf: inline BeginTransaction with first statement (#466)

* perf: inline BeginTransaction with first statement

Inline the BeginTransaction option with the first statement in the transaction,
instead of executing a separate BeginTransaction RPC. This reduces the number
of round-trips to Spanner by one for all transactions that have at least one
SQL statement.

Using line-begin improves performance for most transaction shapes, as it requires
one less round-trip to Spanner. Some transaction shapes do not benefit from this.
These are:
1. Transactions that only write mutations still need an explicit BeginTransaction
   RPC to be executed, as mutations are included in the Commit RPC. The Commit
   RPC can also start a transaction, but such transactions are not guaranteed to
   be applied only once to Spanner.
2. Transactions that execute multiple parallel queries at the start of the
   transaction can see higher end-to-end execution times, as only one query
   can include the BeginTransaction option. All other queries must wait for
   the first query to return at least one result, which also includes the
   transaction identifier, before they can proceed.

The default for the database/sql driver is to use inline-begin.
A follow-up pull request will add an option to the driver to set a different
default for a connection.

* feat: add BeginTransactionOption (#467)

Adds a BeginTransactionOption configuration field that can be used to
determine how the database/sql driver should begin transactions. The
default is to inline the BeginTransaction option with the first SQL
statement of the transaction. This reduces the number of round-trips
needed per transaction by one for most transaction shapes.

Author: olavloite

aborted_transactions_test.go

@@ -325,7 +325,7 @@ func TestQueryWithError_CommitAborted(t *testing.T) {
 		server.PutExecutionTime(testutil.MethodCommitTransaction, testutil.SimulatedExecutionTime{
 			Errors: []error{status.Error(codes.Aborted, "Aborted")},
 		})
-	}, codes.NotFound, 0, 2, 2)
+	}, codes.NotFound, 0, 3, 2)
 }
 
 func TestQueryWithErrorHalfway_CommitAborted(t *testing.T) {
@@ -1080,7 +1080,7 @@ func TestBatchUpdateAbortedWithError_DifferentErrorDuringRetry(t *testing.T) {
 		t.Fatalf("dml statement failed: %v", err)
 	}
 	if _, err := tx.ExecContext(ctx, "RUN BATCH"); spanner.ErrCode(err) != codes.NotFound {
-		t.Fatalf("error code mismatch\nGot: %v\nWant: %v", spanner.ErrCode(err), codes.NotFound)
+		t.Fatalf("error code mismatch\n Got: %v\nWant: %v", spanner.ErrCode(err), codes.NotFound)
 	}
 
 	// Remove the error for the DML statement and cause a retry. The missing
@@ -1094,19 +1094,37 @@ func TestBatchUpdateAbortedWithError_DifferentErrorDuringRetry(t *testing.T) {
 	})
 	err = tx.Commit()
 	if err != ErrAbortedDueToConcurrentModification {
-		t.Fatalf("commit error mismatch\nGot: %v\nWant: %v", err, ErrAbortedDueToConcurrentModification)
+		t.Fatalf("commit error mismatch\n Got: %v\nWant: %v", err, ErrAbortedDueToConcurrentModification)
 	}
 	reqs := drainRequestsFromServer(server.TestSpanner)
 	execReqs := requestsOfType(reqs, reflect.TypeOf(&sppb.ExecuteBatchDmlRequest{}))
-	if g, w := len(execReqs), 2; g != w {
-		t.Fatalf("batch request count mismatch\nGot: %v\nWant: %v", g, w)
+	// There are 3 ExecuteBatchDmlRequests sent to Spanner:
+	// 1. An initial attempt with a BeginTransaction RPC, but this returns a NotFound error.
+	//    This causes the transaction to be retried with an explicit BeginTransaction request.
+	// 2. Another attempt with a transaction ID.
+	// 3. A third attempt after the initial transaction is aborted.
+	if g, w := len(execReqs), 3; g != w {
+		t.Fatalf("batch request count mismatch\n Got: %v\nWant: %v", g, w)
 	}
 	commitReqs := requestsOfType(reqs, reflect.TypeOf(&sppb.CommitRequest{}))
 	// The commit should be attempted only once.
 	if g, w := len(commitReqs), 1; g != w {
-		t.Fatalf("commit request count mismatch\nGot: %v\nWant: %v", g, w)
+		t.Fatalf("commit request count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	// The first ExecuteBatchDml request should try to use an inline-begin.
+	// After that, we should have two BeginTransaction requests.
+	req1 := execReqs[0].(*sppb.ExecuteBatchDmlRequest)
+	if req1.GetTransaction() == nil || req1.GetTransaction().GetBegin() == nil {
+		t.Fatal("the first ExecuteBatchDmlRequest should have a BeginTransaction")
+	}
+	req2 := execReqs[1].(*sppb.ExecuteBatchDmlRequest)
+	if req2.GetTransaction() == nil || req2.GetTransaction().GetId() == nil {
+		t.Fatal("the second ExecuteBatchDmlRequest should have a transaction id")
+	}
+	beginRequests := requestsOfType(reqs, reflect.TypeOf(&sppb.BeginTransactionRequest{}))
+	if g, w := len(beginRequests), 2; g != w {
+		t.Fatalf("begin request count mismatch\n Got: %v\nWant: %v", g, w)
 	}
-
 	// Verify that the db is still usable.
 	if _, err := db.ExecContext(ctx, testutil.UpdateSingersSetLastName); err != nil {
 		t.Fatalf("failed to execute statement after transaction: %v", err)

auto_dml_batch_test.go

@@ -303,8 +303,13 @@ func TestAutoBatchDml_FollowedByRollback(t *testing.T) {
 		if g, w := len(commitRequests), 0; g != w {
 			t.Fatalf("num commit requests mismatch\n Got: %v\nWant: %v", g, w)
 		}
+		beginRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.BeginTransactionRequest{}))
+		if g, w := len(beginRequests), 0; g != w {
+			t.Fatalf("num BeginTransaction requests mismatch\n Got: %v\nWant: %v", g, w)
+		}
 		rollbackRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.RollbackRequest{}))
-		if g, w := len(rollbackRequests), 1; g != w {
+		// There are no rollback requests sent to Spanner, as the transaction is never started.
+		if g, w := len(rollbackRequests), 0; g != w {
 			t.Fatalf("num rollback requests mismatch\n Got: %v\nWant: %v", g, w)
 		}
 	}

conn.go

@@ -248,6 +248,8 @@ type conn struct {
 	// transactions on this connection. This default is ignored if the BeginTx function is
 	// called with an isolation level other than sql.LevelDefault.
 	isolationLevel sql.IsolationLevel
+	// beginTransactionOption determines the default transactions start mode.
+	beginTransactionOption spanner.BeginTransactionOption
 
 	// execOptions are applied to the next statement or transaction that is executed
 	// on this connection. It can also be set by passing it in as an argument to
@@ -660,6 +662,7 @@ func (c *conn) ResetSession(_ context.Context) error {
 	c.autoBatchDmlUpdateCountVerification = !c.connector.connectorConfig.DisableAutoBatchDmlUpdateCountVerification
 	c.retryAborts = c.connector.retryAbortsInternally
 	c.isolationLevel = c.connector.connectorConfig.IsolationLevel
+	c.beginTransactionOption = c.connector.connectorConfig.BeginTransactionOption
 	// TODO: Reset the following fields to the connector default
 	c.autocommitDMLMode = Transactional
 	c.readOnlyStaleness = spanner.TimestampBound{}
@@ -927,7 +930,9 @@ func (c *conn) withTempTransactionOptions(options *ReadWriteTransactionOptions)
 func (c *conn) getTransactionOptions() ReadWriteTransactionOptions {
 	if c.tempTransactionOptions != nil {
 		defer func() { c.tempTransactionOptions = nil }()
-		return *c.tempTransactionOptions
+		opts := *c.tempTransactionOptions
+		opts.TransactionOptions.BeginTransactionOption = c.convertDefaultBeginTransactionOption(opts.TransactionOptions.BeginTransactionOption)
+		return opts
 	}
 	// Clear the transaction tag that has been set on the connection after returning
 	// from this function.
@@ -947,6 +952,9 @@ func (c *conn) getTransactionOptions() ReadWriteTransactionOptions {
 			txOpts.TransactionOptions.IsolationLevel = level
 		}
 	}
+	if txOpts.TransactionOptions.BeginTransactionOption == spanner.DefaultBeginTransaction {
+		txOpts.TransactionOptions.BeginTransactionOption = c.convertDefaultBeginTransactionOption(c.beginTransactionOption)
+	}
 	return txOpts
 }
 
@@ -957,9 +965,11 @@ func (c *conn) withTempReadOnlyTransactionOptions(options *ReadOnlyTransactionOp
 func (c *conn) getReadOnlyTransactionOptions() ReadOnlyTransactionOptions {
 	if c.tempReadOnlyTransactionOptions != nil {
 		defer func() { c.tempReadOnlyTransactionOptions = nil }()
-		return *c.tempReadOnlyTransactionOptions
+		opts := *c.tempReadOnlyTransactionOptions
+		opts.BeginTransactionOption = c.convertDefaultBeginTransactionOption(opts.BeginTransactionOption)
+		return opts
 	}
-	return ReadOnlyTransactionOptions{TimestampBound: c.readOnlyStaleness}
+	return ReadOnlyTransactionOptions{TimestampBound: c.readOnlyStaleness, BeginTransactionOption: c.convertDefaultBeginTransactionOption(c.beginTransactionOption)}
 }
 
 func (c *conn) withTempBatchReadOnlyTransactionOptions(options *BatchReadOnlyTransactionOptions) {
@@ -1034,7 +1044,7 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 			ro = &bo.ReadOnlyTransaction
 		} else {
 			logger = c.logger.With("tx", "ro")
-			ro = c.client.ReadOnlyTransaction().WithTimestampBound(readOnlyTxOpts.TimestampBound)
+			ro = c.client.ReadOnlyTransaction().WithBeginTransactionOption(readOnlyTxOpts.BeginTransactionOption).WithTimestampBound(readOnlyTxOpts.TimestampBound)
 		}
 		c.tx = &readOnlyTransaction{
 			roTx:   ro,
@@ -1080,6 +1090,16 @@ func (c *conn) BeginTx(ctx context.Context, opts driver.TxOptions) (driver.Tx, e
 	return c.tx, nil
 }
 
+func (c *conn) convertDefaultBeginTransactionOption(opt spanner.BeginTransactionOption) spanner.BeginTransactionOption {
+	if opt == spanner.DefaultBeginTransaction {
+		if c.beginTransactionOption == spanner.DefaultBeginTransaction {
+			return spanner.InlinedBeginTransaction
+		}
+		return c.beginTransactionOption
+	}
+	return opt
+}
+
 func (c *conn) inTransaction() bool {
 	return c.tx != nil
 }

conn_with_mockserver_test.go

@@ -39,15 +39,69 @@ func TestBeginTx(t *testing.T) {
 
 	requests := drainRequestsFromServer(server.TestSpanner)
 	beginRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.BeginTransactionRequest{}))
-	if g, w := len(beginRequests), 1; g != w {
+	if g, w := len(beginRequests), 0; g != w {
 		t.Fatalf("begin requests count mismatch\n Got: %v\nWant: %v", g, w)
 	}
-	request := beginRequests[0].(*spannerpb.BeginTransactionRequest)
-	if g, w := request.Options.GetIsolationLevel(), spannerpb.TransactionOptions_ISOLATION_LEVEL_UNSPECIFIED; g != w {
+	executeRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+	if g, w := len(executeRequests), 1; g != w {
+		t.Fatalf("execute requests count mismatch\n Got: %v\nWant: %v", g, w)
+	}
+	request := executeRequests[0].(*spannerpb.ExecuteSqlRequest)
+	if request.GetTransaction() == nil || request.GetTransaction().GetBegin() == nil {
+		t.Fatal("missing begin transaction on ExecuteSqlRequest")
+	}
+	if g, w := request.GetTransaction().GetBegin().GetIsolationLevel(), spannerpb.TransactionOptions_ISOLATION_LEVEL_UNSPECIFIED; g != w {
 		t.Fatalf("begin isolation level mismatch\n Got: %v\nWant: %v", g, w)
 	}
 }
 
+func TestExplicitBeginTx(t *testing.T) {
+	t.Parallel()
+
+	db, server, teardown := setupTestDBConnectionWithConnectorConfig(t, ConnectorConfig{
+		Project:  "p",
+		Instance: "i",
+		Database: "d",
+
+		BeginTransactionOption: spanner.ExplicitBeginTransaction,
+	})
+	defer teardown()
+	ctx := context.Background()
+
+	for _, readOnly := range []bool{true, false} {
+		tx, err := db.BeginTx(ctx, &sql.TxOptions{ReadOnly: readOnly})
+		if err != nil {
+			t.Fatal(err)
+		}
+		res, err := tx.QueryContext(ctx, testutil.SelectFooFromBar)
+		if err != nil {
+			t.Fatal(err)
+		}
+		for res.Next() {
+		}
+		if err := res.Err(); err != nil {
+			t.Fatal(err)
+		}
+		if err := tx.Rollback(); err != nil {
+			t.Fatal(err)
+		}
+
+		requests := drainRequestsFromServer(server.TestSpanner)
+		beginRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.BeginTransactionRequest{}))
+		if g, w := len(beginRequests), 1; g != w {
+			t.Fatalf("begin requests count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+		executeRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+		if g, w := len(executeRequests), 1; g != w {
+			t.Fatalf("execute requests count mismatch\n Got: %v\nWant: %v", g, w)
+		}
+		request := executeRequests[0].(*spannerpb.ExecuteSqlRequest)
+		if request.GetTransaction() == nil || request.GetTransaction().GetId() == nil {
+			t.Fatal("missing transaction id on ExecuteSqlRequest")
+		}
+	}
+}
+
 func TestBeginTxWithIsolationLevel(t *testing.T) {
 	t.Parallel()
 
@@ -76,12 +130,19 @@ func TestBeginTxWithIsolationLevel(t *testing.T) {
 
 			requests := drainRequestsFromServer(server.TestSpanner)
 			beginRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.BeginTransactionRequest{}))
-			if g, w := len(beginRequests), 1; g != w {
+			if g, w := len(beginRequests), 0; g != w {
 				t.Fatalf("begin requests count mismatch\n Got: %v\nWant: %v", g, w)
 			}
-			request := beginRequests[0].(*spannerpb.BeginTransactionRequest)
+			executeRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+			if g, w := len(executeRequests), 1; g != w {
+				t.Fatalf("execute requests count mismatch\n Got: %v\nWant: %v", g, w)
+			}
+			request := executeRequests[0].(*spannerpb.ExecuteSqlRequest)
+			if request.GetTransaction() == nil || request.GetTransaction().GetBegin() == nil {
+				t.Fatalf("execute request does not have a begin transaction")
+			}
 			wantIsolationLevel, _ := toProtoIsolationLevel(originalLevel)
-			if g, w := request.Options.GetIsolationLevel(), wantIsolationLevel; g != w {
+			if g, w := request.GetTransaction().GetBegin().GetIsolationLevel(), wantIsolationLevel; g != w {
 				t.Fatalf("begin isolation level mismatch\n Got: %v\nWant: %v", g, w)
 			}
 		}
@@ -162,12 +223,19 @@ func TestDefaultIsolationLevel(t *testing.T) {
 
 			requests := drainRequestsFromServer(server.TestSpanner)
 			beginRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.BeginTransactionRequest{}))
-			if g, w := len(beginRequests), 1; g != w {
+			if g, w := len(beginRequests), 0; g != w {
 				t.Fatalf("begin requests count mismatch\n Got: %v\nWant: %v", g, w)
 			}
-			request := beginRequests[0].(*spannerpb.BeginTransactionRequest)
+			executeRequests := requestsOfType(requests, reflect.TypeOf(&spannerpb.ExecuteSqlRequest{}))
+			if g, w := len(executeRequests), 1; g != w {
+				t.Fatalf("execute requests count mismatch\n Got: %v\nWant: %v", g, w)
+			}
+			request := executeRequests[0].(*spannerpb.ExecuteSqlRequest)
+			if request.GetTransaction() == nil || request.GetTransaction().GetBegin() == nil {
+				t.Fatalf("ExecuteSqlRequest should have a Begin transaction")
+			}
 			wantIsolationLevel, _ := toProtoIsolationLevel(originalLevel)
-			if g, w := request.Options.GetIsolationLevel(), wantIsolationLevel; g != w {
+			if g, w := request.GetTransaction().GetBegin().GetIsolationLevel(), wantIsolationLevel; g != w {
 				t.Fatalf("begin isolation level mismatch\n Got: %v\nWant: %v", g, w)
 			}
 		}

driver.go

@@ -305,6 +305,11 @@ type ConnectorConfig struct {
 	// IsolationLevel is the default isolation level for read/write transactions.
 	IsolationLevel sql.IsolationLevel
 
+	// BeginTransactionOption determines the default for how to begin transactions.
+	// The Spanner database/sql driver uses spanner.InlinedBeginTransaction by default
+	// for both read-only and read/write transactions.
+	BeginTransactionOption spanner.BeginTransactionOption
+
 	// DecodeToNativeArrays determines whether arrays that have a Go native
 	// type should be decoded to those types rather than the corresponding
 	// spanner.NullTypeName type.
@@ -552,6 +557,11 @@ func createConnector(d *Driver, connectorConfig ConnectorConfig) (*connector, er
 			connectorConfig.IsolationLevel = val
 		}
 	}
+	if strval, ok := connectorConfig.Params[strings.ToLower("BeginTransactionOption")]; ok {
+		if val, err := parseBeginTransactionOption(strval); err == nil {
+			connectorConfig.BeginTransactionOption = val
+		}
+	}
 	if strval, ok := connectorConfig.Params[strings.ToLower("StatementCacheSize")]; ok {
 		if val, err := strconv.Atoi(strval); err == nil {
 			connectorConfig.StatementCacheSize = val
@@ -1038,7 +1048,8 @@ func clearTempReadWriteTransactionOptions(conn *sql.Conn) {
 // ReadOnlyTransactionOptions can be used to create a read-only transaction
 // on a Spanner connection.
 type ReadOnlyTransactionOptions struct {
-	TimestampBound spanner.TimestampBound
+	TimestampBound         spanner.TimestampBound
+	BeginTransactionOption spanner.BeginTransactionOption
 
 	close func()
 }
@@ -1295,6 +1306,18 @@ func checkIsValidType(v driver.Value) bool {
 	return true
 }
 
+func parseBeginTransactionOption(val string) (spanner.BeginTransactionOption, error) {
+	switch strings.ToLower(val) {
+	case strings.ToLower("DefaultBeginTransaction"):
+		return spanner.DefaultBeginTransaction, nil
+	case strings.ToLower("InlinedBeginTransaction"):
+		return spanner.InlinedBeginTransaction, nil
+	case strings.ToLower("ExplicitBeginTransaction"):
+		return spanner.ExplicitBeginTransaction, nil
+	}
+	return spanner.DefaultBeginTransaction, spanner.ToSpannerError(status.Errorf(codes.InvalidArgument, "invalid or unsupported BeginTransactionOption: %v", val))
+}
+
 func parseIsolationLevel(val string) (sql.IsolationLevel, error) {
 	switch strings.Replace(strings.ToLower(strings.TrimSpace(val)), " ", "_", 1) {
 	case "default":