Session var caching, small win

go/libraries/doltcore/sqle/database_provider.go

@@ -762,6 +762,7 @@ func (p DoltDatabaseProvider) databaseForRevision(ctx *sql.Context, revisionQual
 			return nil, false, err
 		}
 
+		dbCache.CacheRevisionDb(db)
 		return db, true, nil
 	case dsess.RevisionTypeTag:
 		// TODO: this should be an interface, not a struct
@@ -780,6 +781,8 @@ func (p DoltDatabaseProvider) databaseForRevision(ctx *sql.Context, revisionQual
 		if err != nil {
 			return nil, false, err
 		}
+
+		dbCache.CacheRevisionDb(db)
 		return db, true, nil
 	case dsess.RevisionTypeCommit:
 		// TODO: this should be an interface, not a struct
@@ -796,6 +799,8 @@ func (p DoltDatabaseProvider) databaseForRevision(ctx *sql.Context, revisionQual
 		if err != nil {
 			return nil, false, err
 		}
+
+		dbCache.CacheRevisionDb(db)
 		return db, true, nil
 	case dsess.RevisionTypeNone:
 		// Returning an error with the fully qualified db name here is our only opportunity to do so in some cases (such

go/libraries/doltcore/sqle/dsess/session.go

@@ -303,6 +303,8 @@ func (d *DoltSession) StartTransaction(ctx *sql.Context, tCharacteristic sql.Tra
 				}
 			}
 
+			// TODO: this check is relatively expensive, we should cache this value when it changes instead of looking it 
+			//  up on each transaction start
 			if _, v, ok := sql.SystemVariables.GetGlobal(ReadReplicaRemote); ok && v != "" {
 				err := ddb.Rebase(ctx)
 				if err != nil && !IgnoreReplicationErrors() {
@@ -336,7 +338,7 @@ func (d *DoltSession) StartTransaction(ctx *sql.Context, tCharacteristic sql.Tra
 			continue
 		}
 
-		_ = d.setSessionVarsForDb(ctx, db.Name(), bs)
+		_ = d.setDbSessionVars(ctx, bs, false)
 	}
 
 	return tx, nil
@@ -895,7 +897,6 @@ func (d *DoltSession) SetRoots(ctx *sql.Context, dbName string, roots doltdb.Roo
 }
 
 // SetWorkingSet sets the working set for this session.
-// Unlike setting the working root alone, this method always marks the session dirty.
 func (d *DoltSession) SetWorkingSet(ctx *sql.Context, dbName string, ws *doltdb.WorkingSet) error {
 	if ws == nil {
 		panic("attempted to set a nil working set for the session")
@@ -910,7 +911,7 @@ func (d *DoltSession) SetWorkingSet(ctx *sql.Context, dbName string, ws *doltdb.
 	}
 	branchState.workingSet = ws
 
-	err = d.setSessionVarsForDb(ctx, dbName, branchState)
+	err = d.setDbSessionVars(ctx, branchState, true)
 	if err != nil {
 		return err
 	}
@@ -965,7 +966,7 @@ func (d *DoltSession) SwitchWorkingSet(
 
 	ctx.SetCurrentDatabase(baseName)
 	
-	return d.setSessionVarsForDb(ctx, dbName, branchState)
+	return d.setDbSessionVars(ctx, branchState, false)
 }
 
 func (d *DoltSession) UseDatabase(ctx *sql.Context, db sql.Database) error {
@@ -1180,8 +1181,6 @@ func (d *DoltSession) addDB(ctx *sql.Context, db SqlDatabase) error {
 		}
 	}
 
-	d.dbCache.CacheRevisionDb(db)
-
 	branchState := sessionState.NewEmptyBranchState(rev)
 
 	// TODO: get rid of all repo state reader / writer stuff. Until we do, swap out the reader with one of our own, and
@@ -1321,10 +1320,16 @@ func (d *DoltSession) BatchMode() batchMode {
 	return d.batchMode
 }
 
-// setSessionVarsForDb updates the three session vars that track the value of the session root hashes
-func (d *DoltSession) setSessionVarsForDb(ctx *sql.Context, dbName string, state *branchState) error {
-	baseName, _ := SplitRevisionDbName(dbName)
-
+// setDbSessionVars updates the three session vars that track the value of the session root hashes
+func (d *DoltSession) setDbSessionVars(ctx *sql.Context, state *branchState, force bool) error {
+	// This check is important even when we are forcing an update, because it updates the idea of staleness  
+	varsStale := d.dbSessionVarsStale(ctx, state)
+	if !varsStale && !force {
+		return nil
+	}
+	
+	baseName := state.dbState.dbName
+	
 	// Different DBs have different requirements for what state is set, so we are maximally permissive on what's expected
 	// in the state object here
 	if state.WorkingSet() != nil {
@@ -1377,6 +1382,17 @@ func (d *DoltSession) setSessionVarsForDb(ctx *sql.Context, dbName string, state
 	return nil
 }
 
+// dbSessionVarsStale returns whether the session vars for the database with the state provided need to be updated in
+// the session
+func (d *DoltSession) dbSessionVarsStale(ctx *sql.Context, state *branchState) bool {
+	dtx, ok := ctx.GetTransaction().(*DoltTransaction)
+	if !ok {
+		return true
+	}
+	
+	return d.dbCache.CacheSessionVars(state, dtx)
+}
+
 func (d DoltSession) WithGlobals(conf config.ReadWriteConfig) *DoltSession {
 	d.globalsConf = conf
 	return &d

go/libraries/doltcore/sqle/dsess/session_cache.go

@@ -41,6 +41,8 @@ type DatabaseCache struct {
 	// initialDbStates caches the initial state of databases by name for a given noms root, which is the primary key.
 	// The secondary key is the lower-case revision-qualified database name.
 	initialDbStates map[doltdb.DataCacheKey]map[string]InitialDbState
+	// sessionVars records a key for the most recently used session vars for each database in the session 
+	sessionVars map[string]sessionVarCacheKey
 
 	mu sync.RWMutex
 }
@@ -50,6 +52,11 @@ type revisionDbCacheKey struct {
 	requestedName string
 }
 
+type sessionVarCacheKey struct {
+	root doltdb.DataCacheKey
+	head string
+}
+
 const maxCachedKeys = 64
 
 func newSessionCache() *SessionCache {
@@ -57,7 +64,9 @@ func newSessionCache() *SessionCache {
 }
 
 func newDatabaseCache() *DatabaseCache {
-	return &DatabaseCache{}
+	return &DatabaseCache{
+		sessionVars: make(map[string]sessionVarCacheKey),
+	}
 }
 
 // CacheTableIndexes caches all indexes for the table with the name given
@@ -293,4 +302,28 @@ func (c *DatabaseCache) CacheInitialDbState(key doltdb.DataCacheKey, revisionDbN
 	}
 
 	dbsForKey[revisionDbName] = state
+}
+
+// CacheSessionVars updates the session var cache for the given branch state and transaction and returns whether it 
+// was updated. If it was updated, session vars need to be set for the state and transaction given. Otherwise they 
+// haven't changed and can be reused.
+func (c *DatabaseCache) CacheSessionVars(branchState *branchState, transaction *DoltTransaction) bool {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	dbBaseName := branchState.dbState.dbName
+	
+	existingKey, found := c.sessionVars[dbBaseName]
+	root, hasRoot := transaction.GetInitialRoot(dbBaseName)
+	if !hasRoot {
+		return true
+	}
+	
+	newKey := sessionVarCacheKey{
+		root: doltdb.DataCacheKey{Hash: root},
+		head: branchState.head,
+	}
+	
+	c.sessionVars[dbBaseName] = newKey
+	return !found || existingKey != newKey
 }