diff --git a/go/libraries/doltcore/sqle/database_provider.go b/go/libraries/doltcore/sqle/database_provider.go index 55b0331a63..ad44a87c4f 100644 --- a/go/libraries/doltcore/sqle/database_provider.go +++ b/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 diff --git a/go/libraries/doltcore/sqle/dsess/session.go b/go/libraries/doltcore/sqle/dsess/session.go index 75a3422b21..e65613f892 100644 --- a/go/libraries/doltcore/sqle/dsess/session.go +++ b/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 diff --git a/go/libraries/doltcore/sqle/dsess/session_cache.go b/go/libraries/doltcore/sqle/dsess/session_cache.go index f4e0889110..c42902c714 100755 --- a/go/libraries/doltcore/sqle/dsess/session_cache.go +++ b/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 } \ No newline at end of file