isession: create internal session

This introduces the InternalSession. The immediate motivation for this
is to improve the performance of LDR. LDR currently depends on the
internal executor. Every statement executed with the internal executor
creates its own connExecutor instance. This prevents LDR from taking
advantage of prepared statements and generic query plans.

Long term, the idea is to have InternalSession replace the internal
executor. But it was developed as a stand alone conn executor wrapper in
order to minimize the risk of breaking existing clients of the internal
executor.

One key design feature of the InternalSession is the state machine was
moved from a separate goroutine to the goroutine that is requesting the
statement execution. This makes CPU profiles for internal processes much
easier to read because it is obvious which call stack is using the SQL
layer. It may also help with integrating LDR's SQL CPU usage with
elastic admission control since the CPU time is accounted to the calling
goroutine.

Release note: none
Informs: #148310

Author: jeffswenson

pkg/BUILD.bazel

@@ -512,6 +512,7 @@ ALL_TESTS = [
     "//pkg/sql/inspect:inspect_test",
     "//pkg/sql/inverted:inverted_disallowed_imports_test",
     "//pkg/sql/inverted:inverted_test",
+    "//pkg/sql/isession:isession_test",
     "//pkg/sql/lex:lex_disallowed_imports_test",
     "//pkg/sql/lex:lex_test",
     "//pkg/sql/lexbase:lexbase_test",
@@ -2076,6 +2077,8 @@ GO_TARGETS = [
     "//pkg/sql/inspect:inspect_test",
     "//pkg/sql/inverted:inverted",
     "//pkg/sql/inverted:inverted_test",
+    "//pkg/sql/isession:isession",
+    "//pkg/sql/isession:isession_test",
     "//pkg/sql/isql:isql",
     "//pkg/sql/lex:lex",
     "//pkg/sql/lex:lex_test",

pkg/server/BUILD.bazel

@@ -245,6 +245,7 @@ go_library(
         "//pkg/sql/idxusage",
         "//pkg/sql/importer",
         "//pkg/sql/inspect",
+        "//pkg/sql/isession",
         "//pkg/sql/isql",
         "//pkg/sql/lexbase",
         "//pkg/sql/optionalnodeliveness",

pkg/server/server.go

@@ -101,6 +101,7 @@ import (
 	_ "github.com/cockroachdb/cockroach/pkg/sql/gcjob"    // register jobs declared outside of pkg/sql
 	_ "github.com/cockroachdb/cockroach/pkg/sql/importer" // register jobs/planHooks declared outside of pkg/sql
 	_ "github.com/cockroachdb/cockroach/pkg/sql/inspect"  // register job and planHook declared outside of pkg/sql
+	_ "github.com/cockroachdb/cockroach/pkg/sql/isession" // register isession constructor hook
 	"github.com/cockroachdb/cockroach/pkg/sql/optionalnodeliveness"
 	"github.com/cockroachdb/cockroach/pkg/sql/pgwire"
 	_ "github.com/cockroachdb/cockroach/pkg/sql/schemachanger/scjob" // register jobs declared outside of pkg/sql

pkg/sql/BUILD.bazel

@@ -145,6 +145,7 @@ go_library(
         "instrumentation.go",
         "internal.go",
         "internal_result_channel.go",
+        "internal_session.go",
         "inverted_filter.go",
         "inverted_join.go",
         "job_exec_context.go",

pkg/sql/conn_executor.go

@@ -3503,9 +3503,9 @@ var retryableMinTimestampBoundUnsatisfiableError = errors.Newf(
 	"retryable MinTimestampBoundUnsatisfiableError",
 )
 
-// errIsRetryable is true if the error is a client-visible retry error
+// ErrIsRetryable is true if the error is a client-visible retry error
 // or the error is a special error that is handled internally and retried.
-func errIsRetryable(err error) bool {
+func ErrIsRetryable(err error) bool {
 	return errors.HasInterface(err, (*pgerror.ClientVisibleRetryError)(nil)) ||
 		errors.Is(err, retryableMinTimestampBoundUnsatisfiableError) ||
 		descs.IsTwoVersionInvariantViolationError(err)
@@ -3563,7 +3563,7 @@ func (ex *connExecutor) makeErrEvent(err error, stmt tree.Statement) (fsm.Event,
 		}
 	}
 
-	retryable := errIsRetryable(err)
+	retryable := ErrIsRetryable(err)
 	if retryable {
 		var rc rewindCapability
 		var canAutoRetry bool

pkg/sql/conn_executor_prepare.go

@@ -388,13 +388,13 @@ func (ex *connExecutor) execBind(
 	// Decode the arguments, except for internal queries for which we just verify
 	// that the arguments match what's expected.
 	qargs := make(tree.QueryArguments, numQArgs)
-	if bindCmd.internalArgs != nil {
-		if len(bindCmd.internalArgs) != int(numQArgs) {
+	if bindCmd.InternalArgs != nil {
+		if len(bindCmd.InternalArgs) != int(numQArgs) {
 			return retErr(
 				pgwirebase.NewProtocolViolationErrorf(
-					"expected %d arguments, got %d", numQArgs, len(bindCmd.internalArgs)))
+					"expected %d arguments, got %d", numQArgs, len(bindCmd.InternalArgs)))
 		}
-		for i, datum := range bindCmd.internalArgs {
+		for i, datum := range bindCmd.InternalArgs {
 			t := ps.InferredTypes[i]
 			if oid := datum.ResolvedType().Oid(); datum != tree.DNull && oid != t {
 				return retErr(

pkg/sql/conn_executor_savepoints.go

@@ -147,7 +147,7 @@ func (ex *connExecutor) execRelease(
 		}
 		// Committing the transaction failed. We'll go to state RestartWait if
 		// it's a retryable error, or to state RollbackWait otherwise.
-		if errIsRetryable(err) {
+		if ErrIsRetryable(err) {
 			// For certain retryable errors, we should turn them into client visible
 			// errors, since the client needs to retry now.
 			var conversionError error

pkg/sql/conn_executor_show_commit_timestamp.go

@@ -83,7 +83,7 @@ func (ex *connExecutor) execShowCommitTimestampInOpenState(
 
 	// Committing the transaction failed. We'll go to state RestartWait if
 	// it's a retryable error, or to state RollbackWait otherwise.
-	if errIsRetryable(err) {
+	if ErrIsRetryable(err) {
 		rc, canAutoRetry := ex.getRewindTxnCapability()
 		ev := eventRetryableErr{
 			IsCommit:     fsm.FromBool(false /* isCommit */),

pkg/sql/conn_io.go

@@ -281,14 +281,14 @@ type BindStmt struct {
 	// code, in which case that code will be applied to all arguments.
 	ArgFormatCodes []pgwirebase.FormatCode
 
-	// internalArgs, if not nil, represents the arguments for the prepared
+	// InternalArgs, if not nil, represents the arguments for the prepared
 	// statements as produced by the internal clients. These don't need to go
 	// through encoding/decoding of the args. However, the types of the datums
 	// must correspond exactly to the inferred types (but note that the types of
 	// the datums are passes as type hints to the PrepareStmt command, so the
 	// inferred types should reflect that).
-	// If internalArgs is specified, Args and ArgFormatCodes are ignored.
-	internalArgs []tree.Datum
+	// If InternalArgs is specified, Args and ArgFormatCodes are ignored.
+	InternalArgs []tree.Datum
 }
 
 // command implements the Command interface.
@@ -558,6 +558,23 @@ func (buf *StmtBuf) CurCmd() (Command, CmdPos, error) {
 	}
 }
 
+// Empty returns true if there are no unprocessed commands in the buffer.
+// If the buffer is closed, it returns io.EOF.
+func (buf *StmtBuf) Empty() (bool, error) {
+	buf.mu.Lock()
+	defer 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
+	}
+	return !(cmdIdx < buf.mu.data.Len()), 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/conn_io_test.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/util/leaktest"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
+	"github.com/stretchr/testify/require"
 )
 
 func assertStmt(t *testing.T, cmd Command, exp string) {
@@ -342,3 +343,52 @@ func TestStmtBufBatching(t *testing.T) {
 		t.Fatalf("expected pos to be %d, got: %d", 9, pos)
 	}
 }
+
+func TestStmtBufEmpty(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+
+	t.Run("EmptyBuffer", func(t *testing.T) {
+		buf := NewStmtBuf(0)
+		empty, err := buf.Empty()
+		require.NoError(t, err)
+		require.True(t, empty)
+	})
+
+	t.Run("NonEmptyBuffer", func(t *testing.T) {
+		buf := NewStmtBuf(0)
+		s1, err := parser.ParseOne("SELECT 1")
+		require.NoError(t, err)
+		mustPush(ctx, t, buf, ExecStmt{Statement: s1})
+
+		empty, err := buf.Empty()
+		require.NoError(t, err)
+		require.False(t, empty)
+	})
+
+	t.Run("EmptyAfterAdvancing", func(t *testing.T) {
+		buf := NewStmtBuf(0)
+		s1, err := parser.ParseOne("SELECT 1")
+		require.NoError(t, err)
+		mustPush(ctx, t, buf, ExecStmt{Statement: s1})
+
+		empty, err := buf.Empty()
+		require.NoError(t, err)
+		require.False(t, empty)
+
+		buf.AdvanceOne()
+		empty, err = buf.Empty()
+		require.NoError(t, err)
+		require.True(t, empty)
+	})
+
+	t.Run("ClosedBuffer", func(t *testing.T) {
+		buf := NewStmtBuf(0)
+		buf.Close()
+
+		_, err := buf.Empty()
+		require.Equal(t, io.EOF, err)
+	})
+}