isession: add a helper for mutating the session

This adds a MutateSession method that allows modifying the session
without going through the SQL layer. It is possible to modify the
session using `SET ...` statements, but LDR must mutate the session
before every statement and modifying the session consumes ~10% of all
CPU cycles in a cluster that is using the crud writer to replicate data.

An alternative to adding MutateSession is we could optimize the Set
statement. The main reasons its inefficient for LDR are:
1. It only supports strings, so we can set the hlc origin timestamp
	 without roundtripping through a string.
2. The optimizer does not support generic query plans for prepared
	 statements that include a set statement.

Release note: none
Epic:	CRDB-48647

Author: jeffswenson

pkg/sql/internal_session.go

@@ -113,6 +113,11 @@ func (c *ConnectionStateMachine) Push(ctx context.Context, cmd Command) error {
 	return c.buffer.Push(ctx, cmd)
 }
 
+// SessionDataMutatorIterator returns the session data mutator iterator for this connection.
+func (c *ConnectionStateMachine) SessionDataMutatorIterator() *sessionmutator.SessionDataMutatorIterator {
+	return c.config.sdMutIterator
+}
+
 // NewInternalSession constructs an internal session instance.
 func (s *Server) NewInternalSession(
 	ctx context.Context,

pkg/sql/isession/BUILD.bazel

@@ -20,6 +20,7 @@ go_library(
         "//pkg/sql/pgwire/pgwirebase",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondatapb",
+        "//pkg/sql/sessionmutator",
         "//pkg/sql/types",
         "//pkg/util/retry",
         "@com_github_cockroachdb_crlib//crtime",
@@ -45,6 +46,7 @@ go_test(
         "//pkg/sql/isql",
         "//pkg/sql/parser",
         "//pkg/sql/sem/tree",
+        "//pkg/sql/sessionmutator",
         "//pkg/sql/types",
         "//pkg/testutils/serverutils",
         "//pkg/testutils/sqlutils",

pkg/sql/isession/internal_session.go

@@ -14,6 +14,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessionmutator"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
 	"github.com/cockroachdb/crlib/crtime"
@@ -325,6 +326,20 @@ func (i *InternalSession) readResults(ctx context.Context) ([]tree.Datums, int,
 	return rows, rowCount, resultErr
 }
 
+func (i *InternalSession) ModifySession(
+	ctx context.Context, mutate func(mutator sessionmutator.SessionDataMutator),
+) error {
+	if i.poison != nil {
+		return i.poison
+	}
+
+	sdMutIterator := i.csm.SessionDataMutatorIterator()
+	return sdMutIterator.ApplyOnTopMutator(func(m sessionmutator.SessionDataMutator) error {
+		mutate(m)
+		return nil
+	})
+}
+
 func (i *InternalSession) Close(ctx context.Context) {
 	i.csm.Close(ctx)
 }

pkg/sql/isession/internal_session_test.go

@@ -16,6 +16,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/sql/isql"
 	"github.com/cockroachdb/cockroach/pkg/sql/parser"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessionmutator"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 	"github.com/cockroachdb/cockroach/pkg/testutils/serverutils"
 	"github.com/cockroachdb/cockroach/pkg/testutils/sqlutils"
@@ -612,3 +613,35 @@ func TestSavepoint(t *testing.T) {
 		{tree.NewDInt(tree.DInt(1)), tree.NewDInt(tree.DInt(10))},
 	}, rows)
 }
+
+func TestModifySession(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+	defer log.Scope(t).Close(t)
+
+	ctx := context.Background()
+	s := serverutils.StartServerOnly(t, base.TestServerArgs{})
+	defer s.Stopper().Stop(ctx)
+
+	db := s.InternalDB().(descs.DB)
+	session, err := db.Session(ctx, "test-session")
+	require.NoError(t, err)
+	defer session.Close(ctx)
+
+	// Modify the application name using ModifySession.
+	err = session.ModifySession(ctx, func(m sessionmutator.SessionDataMutator) {
+		m.SetApplicationName("my_test_app")
+	})
+	require.NoError(t, err)
+
+	// Verify the change by querying the session variable.
+	stmt, err := parser.ParseOne("SHOW application_name")
+	require.NoError(t, err)
+
+	prepared, err := session.Prepare(ctx, "show_app", stmt, nil)
+	require.NoError(t, err)
+
+	rows, err := session.QueryPrepared(ctx, prepared, nil)
+	require.NoError(t, err)
+	require.Len(t, rows, 1)
+	require.Equal(t, "my_test_app", string(tree.MustBeDString(rows[0][0])))
+}

pkg/sql/isql/BUILD.bazel

@@ -18,6 +18,7 @@ go_library(
         "//pkg/sql/parser/statements",
         "//pkg/sql/sem/tree",
         "//pkg/sql/sessiondata",
+        "//pkg/sql/sessionmutator",
         "//pkg/sql/types",
         "//pkg/util/admission/admissionpb",
         "@com_github_cockroachdb_redact//:redact",

pkg/sql/isql/isql_session.go

@@ -10,6 +10,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/sql/parser/statements"
 	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
+	"github.com/cockroachdb/cockroach/pkg/sql/sessionmutator"
 	"github.com/cockroachdb/cockroach/pkg/sql/types"
 )
 
@@ -80,6 +81,19 @@ type Session interface {
 	// }
 	Savepoint(ctx context.Context, do func(context.Context) error) error
 
+	// ModifySession executes a function that mutates the session using the
+	// sessionmutator.SessionDataMutator argument.
+	//
+	// Example:
+	// err := session.ModifySession(ctx, func(mutator sessionmutator.SessionDataMutator) {
+	// 	mutator.SetApplicationName("my_app")
+	// 	mutator.SetDatabase("my_database")
+	// })
+	// if err != nil {
+	// 	return err
+	// }
+	ModifySession(ctx context.Context, mutate func(mutator sessionmutator.SessionDataMutator)) error
+
 	// Close closes the session and cleans up internal resources.
 	Close(ctx context.Context)
 }