sql: allow the 1PC optimization in some cases in the extended protocol

A previous commit (7e2cbf5) fixed our
pgwire implementation so that it does not auto-commit a statement
executed in the extended protocol until a Sync message is received. That
change also had the undesired effect of disabling the 1PC ("insert fast
path") optimization that is critical for write-heavy workloads.

With this current commit, the 1PC optimization is allowed again, as long
as the statement execution is immediately followed by a Sync message.
This still has the correct bugfix semantics, but allows the optimization
for the common case of how the extended protocol is used.

No release note since this only affects unreleased versions.

Release note: None

pkg/sql/conn_executor_exec.go

@@ -262,22 +262,42 @@ func (ex *connExecutor) execStmtInOpenState(
 	ast := parserStmt.AST
 	ctx = withStatement(ctx, ast)
 
+	makeErrEvent := func(err error) (fsm.Event, fsm.EventPayload, error) {
+		ev, payload := ex.makeErrEvent(err, ast)
+		return ev, payload, nil
+	}
+
 	var stmt Statement
 	queryID := ex.generateID()
 	// Update the deadline on the transaction based on the collections.
 	err := ex.extraTxnState.descCollection.MaybeUpdateDeadline(ctx, ex.state.mu.txn)
 	if err != nil {
-		ev, pl := ex.makeErrEvent(err, ast)
-		return ev, pl, nil
+		return makeErrEvent(err)
 	}
+	os := ex.machine.CurState().(stateOpen)
 
+	isNextCmdSync := false
 	isExtendedProtocol := prepared != nil
 	if isExtendedProtocol {
 		stmt = makeStatementFromPrepared(prepared, queryID)
+		// Only check for Sync in the extended protocol. In the simple protocol,
+		// Sync is meaningless, so it's OK to let isNextCmdSync default to false.
+		isNextCmdSync, err = ex.stmtBuf.isNextCmdSync()
+		if err != nil {
+			return makeErrEvent(err)
+		}
 	} else {
 		stmt = makeStatement(parserStmt, queryID)
 	}
 
+	// In some cases, we need to turn off autocommit behavior here. The postgres
+	// docs say that commands in the extended protocol are all treated as an
+	// implicit transaction that does not get committed until a Sync message is
+	// received. However, if we are executing a statement that is immediately
+	// followed by Sync (which is the common case), then we still can auto-commit,
+	// which allows the "insert fast path" (1PC optimization) to be used.
+	canAutoCommit := os.ImplicitTxn.Get() && (!isExtendedProtocol || isNextCmdSync)
+
 	ex.incrementStartedStmtCounter(ast)
 	defer func() {
 		if retErr == nil && !payloadHasError(retPayload) {
@@ -289,8 +309,6 @@ func (ex *connExecutor) execStmtInOpenState(
 	ex.state.mu.stmtCount++
 	ex.state.mu.Unlock()
 
-	os := ex.machine.CurState().(stateOpen)
-
 	var timeoutTicker *time.Timer
 	queryTimedOut := false
 	// doneAfterFunc will be allocated only when timeoutTicker is non-nil.
@@ -355,11 +373,6 @@ func (ex *connExecutor) execStmtInOpenState(
 		}
 	}(ctx, res)
 
-	makeErrEvent := func(err error) (fsm.Event, fsm.EventPayload, error) {
-		ev, payload := ex.makeErrEvent(err, ast)
-		return ev, payload, nil
-	}
-
 	p := &ex.planner
 	stmtTS := ex.server.cfg.Clock.PhysicalTime()
 	ex.statsCollector.Reset(ex.applicationStats, ex.phaseTimes)
@@ -497,12 +510,7 @@ func (ex *connExecutor) execStmtInOpenState(
 		if retEv != nil || retErr != nil {
 			return
 		}
-		// The postgres docs say that commands in the extended protocol are
-		// all treated as an implicit transaction that does not get committed
-		// until a Sync message is received. The prepared statement will only be
-		// nil if we are in the simple protocol; for the extended protocol the
-		// commit occurs when Sync is received.
-		if os.ImplicitTxn.Get() && !isExtendedProtocol {
+		if canAutoCommit {
 			retEv, retPayload = ex.handleAutoCommit(ctx, ast)
 			return
 		}
@@ -701,11 +709,7 @@ func (ex *connExecutor) execStmtInOpenState(
 	p.stmt = stmt
 	p.cancelChecker.Reset(ctx)
 
-	// We need to turn off autocommit behavior here so that the "insert fast path"
-	// does not get triggered. The postgres docs say that commands in the extended
-	// protocol are all treated as an implicit transaction that does not get
-	// committed until a Sync message is received.
-	p.autoCommit = os.ImplicitTxn.Get() && !isExtendedProtocol && !ex.server.cfg.TestingKnobs.DisableAutoCommit
+	p.autoCommit = canAutoCommit && !ex.server.cfg.TestingKnobs.DisableAutoCommit
 
 	var stmtThresholdSpan *tracing.Span
 	alreadyRecording := ex.transitionCtx.sessionTracing.Enabled()

pkg/sql/conn_io.go

@@ -442,6 +442,36 @@ func (buf *StmtBuf) CurCmd() (Command, CmdPos, error) {
 	}
 }
 
+// isNextCmdSync peeks at the next command, and returns true if it is a Sync
+// message. If there is not another command in the buffer already, then false
+// is returned. The position is reset before returning, so this does not
+// affect the curPos.
+//
+// If the buffer has previously been Close()d, or is closed while this is
+// blocked, io.EOF is returned.
+func (buf *StmtBuf) isNextCmdSync() (bool, error) {
+	buf.mu.Lock()
+	prev := buf.mu.curPos
+	buf.mu.curPos++
+	defer func() {
+		buf.mu.curPos = prev
+		buf.mu.Unlock()
+	}()
+	if buf.mu.closed {
+		return false, io.EOF
+	}
+	curPos := buf.mu.curPos
+	cmdIdx, err := buf.translatePosLocked(curPos)
+	if err != nil {
+		return false, err
+	}
+	if cmdIdx < buf.mu.data.Len() {
+		_, isSync := buf.mu.data.Get(cmdIdx).(Sync)
+		return isSync, nil
+	}
+	return false, nil
+}
+
 // translatePosLocked translates an absolute position of a command (counting
 // from the connection start) to the index of the respective command in the
 // buffer (so, it returns an index relative to the start of the buffer).

pkg/sql/tests/BUILD.bazel

@@ -32,6 +32,7 @@ go_test(
         "enum_test.go",
         "hash_sharded_test.go",
         "impure_builtin_test.go",
+        "insert_fast_path_test.go",
         "inverted_index_test.go",
         "kv_test.go",
         "main_test.go",

pkg/sql/tests/insert_fast_path_test.go

@@ -0,0 +1,73 @@
+// Copyright 2022 The Cockroach Authors.
+//
+// Use of this software is governed by the Business Source License
+// included in the file licenses/BSL.txt.
+//
+// As of the Change Date specified in that file, in accordance with
+// the Business Source License, use of this software will be governed
+// by the Apache License, Version 2.0, included in the file
+// licenses/APL.txt.
+
+package tests
+
+import (
+	"context"
+	gosql "database/sql"
+	"testing"
+
+	"github.com/cockroachdb/cockroach/pkg/base"
+	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
+	"github.com/cockroachdb/cockroach/pkg/testutils/testcluster"
+	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
+	"github.com/stretchr/testify/require"
+)
+
+// TestInsertFastPathExtendedProtocol verifies that the 1PC "insert fast path"
+// optimization is applied when doing a simple INSERT with a prepared statement.
+func TestInsertFastPathExtendedProtocol(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	ctx := context.Background()
+
+	var db *gosql.DB
+
+	params, _ := CreateTestServerParams()
+	params.Settings = cluster.MakeTestingClusterSettings()
+
+	tc := testcluster.StartTestCluster(t, 1, base.TestClusterArgs{ServerArgs: params})
+	defer tc.Stopper().Stop(ctx)
+	db = tc.ServerConn(0)
+	_, err := db.Exec(`CREATE TABLE fast_path_test(val int);`)
+	require.NoError(t, err)
+
+	conn, err := db.Conn(ctx)
+	require.NoError(t, err)
+	_, err = conn.ExecContext(ctx, "SET tracing = 'on'")
+	require.NoError(t, err)
+	// Use placeholders to force usage of extended protocol.
+	_, err = conn.ExecContext(ctx, "INSERT INTO fast_path_test VALUES($1)", 1)
+	require.NoError(t, err)
+
+	fastPathEnabled := false
+	rows, err := conn.QueryContext(ctx, "SELECT message, operation FROM [SHOW TRACE FOR SESSION]")
+	require.NoError(t, err)
+	for rows.Next() {
+		var msg, operation string
+		err = rows.Scan(&msg, &operation)
+		require.NoError(t, err)
+		if msg == "autocommit enabled" && operation == "batch flow coordinator" {
+			fastPathEnabled = true
+		}
+	}
+	require.NoError(t, rows.Err())
+	require.True(t, fastPathEnabled)
+	_, err = conn.ExecContext(ctx, "SET tracing = 'off'")
+	require.NoError(t, err)
+	err = conn.Close()
+	require.NoError(t, err)
+
+	// Verify that the insert committed successfully.
+	var c int
+	err = db.QueryRow("SELECT count(*) FROM fast_path_test").Scan(&c)
+	require.NoError(t, err)
+	require.Equal(t, 1, c, "expected 1 row, got %d", c)
+}