Merge pull request #7024 from nustiueudinastea/concurrent-branches-map

Concurrent branches map

go/libraries/doltcore/env/actions/branch.go

@@ -149,7 +149,7 @@ func DeleteBranchOnDB(ctx context.Context, dbdata env.DbData, branchRef ref.Dolt
 			return err
 		}
 
-		trackedBranch, hasUpstream := trackedBranches[branchRef.GetPath()]
+		trackedBranch, hasUpstream := trackedBranches.Get(branchRef.GetPath())
 		if hasUpstream {
 			err = validateBranchMergedIntoUpstream(ctx, dbdata, branchRef, trackedBranch.Remote, pro)
 			if err != nil {

go/libraries/doltcore/env/environment.go

@@ -994,7 +994,7 @@ func (dEnv *DoltEnv) RemoveBackup(ctx context.Context, name string) error {
 	return nil
 }
 
-func (dEnv *DoltEnv) GetBranches() (map[string]BranchConfig, error) {
+func (dEnv *DoltEnv) GetBranches() (*concurrentmap.Map[string, BranchConfig], error) {
 	if dEnv.RSLoadErr != nil {
 		return nil, dEnv.RSLoadErr
 	}
@@ -1007,7 +1007,7 @@ func (dEnv *DoltEnv) UpdateBranch(name string, new BranchConfig) error {
 		return dEnv.RSLoadErr
 	}
 
-	dEnv.RepoState.Branches[name] = new
+	dEnv.RepoState.Branches.Set(name, new)
 
 	err := dEnv.RepoState.Save(dEnv.FS)
 	if err != nil {

go/libraries/doltcore/env/environment_test.go

@@ -53,7 +53,7 @@ func createTestEnv(isInitialized bool, hasLocalConfig bool) (*DoltEnv, *filesys.
 		initialDirs = append(initialDirs, doltDataDir)
 
 		mainRef := ref.NewBranchRef(DefaultInitBranch)
-		repoState := &RepoState{Head: ref.MarshalableRef{Ref: mainRef}, Remotes: concurrentmap.New[string, Remote](), Backups: concurrentmap.New[string, Remote]()}
+		repoState := &RepoState{Head: ref.MarshalableRef{Ref: mainRef}, Remotes: concurrentmap.New[string, Remote](), Backups: concurrentmap.New[string, Remote](), Branches: concurrentmap.New[string, BranchConfig]()}
 		repoStateData, err := json.Marshal(repoState)
 
 		if err != nil {

go/libraries/doltcore/env/memory.go

@@ -216,8 +216,8 @@ func (m MemoryRepoState) AddRemote(r Remote) error {
 	return fmt.Errorf("cannot insert a remote in a memory database")
 }
 
-func (m MemoryRepoState) GetBranches() (map[string]BranchConfig, error) {
-	return make(map[string]BranchConfig), nil
+func (m MemoryRepoState) GetBranches() (*concurrentmap.Map[string, BranchConfig], error) {
+	return concurrentmap.New[string, BranchConfig](), nil
 }
 
 func (m MemoryRepoState) UpdateBranch(name string, new BranchConfig) error {

go/libraries/doltcore/env/remotes.go

@@ -30,6 +30,7 @@ import (
 	"github.com/dolthub/dolt/go/libraries/doltcore/doltdb"
 	"github.com/dolthub/dolt/go/libraries/doltcore/ref"
 	"github.com/dolthub/dolt/go/libraries/utils/argparser"
+	"github.com/dolthub/dolt/go/libraries/utils/concurrentmap"
 	"github.com/dolthub/dolt/go/libraries/utils/config"
 	"github.com/dolthub/dolt/go/libraries/utils/earl"
 	filesys2 "github.com/dolthub/dolt/go/libraries/utils/filesys"
@@ -281,7 +282,7 @@ func getPushTargetsAndRemoteFromNoArg(ctx context.Context, rsr RepoStateReader,
 	}
 }
 
-func getPushTargetsAndRemoteForAllBranches(ctx context.Context, rsrBranches map[string]BranchConfig, currentBranch ref.DoltRef, remote *Remote, ddb *doltdb.DoltDB, force, setUpstream bool) ([]*PushTarget, *Remote, error) {
+func getPushTargetsAndRemoteForAllBranches(ctx context.Context, rsrBranches *concurrentmap.Map[string, BranchConfig], currentBranch ref.DoltRef, remote *Remote, ddb *doltdb.DoltDB, force, setUpstream bool) ([]*PushTarget, *Remote, error) {
 	localBranches, err := ddb.GetBranches(ctx)
 	if err != nil {
 		return nil, nil, err
@@ -293,7 +294,7 @@ func getPushTargetsAndRemoteForAllBranches(ctx context.Context, rsrBranches map[
 	return getPushTargetsAndRemoteForBranchRefs(ctx, rsrBranches, lbNames, currentBranch, remote, ddb, force, setUpstream)
 }
 
-func getPushTargetsAndRemoteForBranchRefs(ctx context.Context, rsrBranches map[string]BranchConfig, localBranches []string, currentBranch ref.DoltRef, remote *Remote, ddb *doltdb.DoltDB, force, setUpstream bool) ([]*PushTarget, *Remote, error) {
+func getPushTargetsAndRemoteForBranchRefs(ctx context.Context, rsrBranches *concurrentmap.Map[string, BranchConfig], localBranches []string, currentBranch ref.DoltRef, remote *Remote, ddb *doltdb.DoltDB, force, setUpstream bool) ([]*PushTarget, *Remote, error) {
 	var pushOptsList []*PushTarget
 	for _, refSpecName := range localBranches {
 		refSpec, err := getRefSpecFromStr(ctx, ddb, refSpecName)
@@ -303,7 +304,7 @@ func getPushTargetsAndRemoteForBranchRefs(ctx context.Context, rsrBranches map[s
 
 		// if the remote of upstream does not match the remote given,
 		// it should push to the given remote creating new remote branch
-		upstream, hasUpstream := rsrBranches[refSpecName]
+		upstream, hasUpstream := rsrBranches.Get(refSpecName)
 		hasUpstream = hasUpstream && upstream.Remote == remote.Name
 
 		opts, err := getPushTargetFromRefSpec(refSpec, currentBranch, remote, force, setUpstream, hasUpstream)
@@ -354,15 +355,15 @@ func getPushTargetFromRefSpec(refSpec ref.RefSpec, currentBranch ref.DoltRef, re
 // If there is no remote specified, the current branch needs to have upstream set to push; otherwise, returns error.
 // This function returns |refSpec| for current branch, name of the remote the branch is associated with and
 // whether the current branch has upstream set.
-func getCurrentBranchRefSpec(ctx context.Context, branches map[string]BranchConfig, rsr RepoStateReader, ddb *doltdb.DoltDB, remoteName string, isDefaultRemote, remoteSpecified, setUpstream, pushAutoSetupRemote bool) (ref.RefSpec, string, bool, error) {
+func getCurrentBranchRefSpec(ctx context.Context, branches *concurrentmap.Map[string, BranchConfig], rsr RepoStateReader, ddb *doltdb.DoltDB, remoteName string, isDefaultRemote, remoteSpecified, setUpstream, pushAutoSetupRemote bool) (ref.RefSpec, string, bool, error) {
 	var refSpec ref.RefSpec
 	currentBranch, err := rsr.CWBHeadRef()
 	if err != nil {
 		return nil, "", false, err
 	}
 
 	currentBranchName := currentBranch.GetPath()
-	upstream, hasUpstream := branches[currentBranchName]
+	upstream, hasUpstream := branches.Get(currentBranchName)
 
 	if remoteSpecified || pushAutoSetupRemote {
 		if isDefaultRemote && !pushAutoSetupRemote {
@@ -610,7 +611,7 @@ func NewPullSpec(
 			return nil, err
 		}
 
-		trackedBranch, hasUpstream := trackedBranches[branch.GetPath()]
+		trackedBranch, hasUpstream := trackedBranches.Get(branch.GetPath())
 		if !hasUpstream {
 			if remoteOnly {
 				return nil, fmt.Errorf(ErrPullWithRemoteNoUpstream.Error(), remoteName)

go/libraries/doltcore/env/repo_state.go

@@ -33,7 +33,7 @@ type RepoStateReader interface {
 	CWBHeadSpec() (*doltdb.CommitSpec, error)
 	GetRemotes() (*concurrentmap.Map[string, Remote], error)
 	GetBackups() (*concurrentmap.Map[string, Remote], error)
-	GetBranches() (map[string]BranchConfig, error)
+	GetBranches() (*concurrentmap.Map[string, BranchConfig], error)
 }
 
 type RepoStateWriter interface {
@@ -69,10 +69,10 @@ type BranchConfig struct {
 }
 
 type RepoState struct {
-	Head     ref.MarshalableRef                 `json:"head"`
-	Remotes  *concurrentmap.Map[string, Remote] `json:"remotes"`
-	Backups  *concurrentmap.Map[string, Remote] `json:"backups"`
-	Branches map[string]BranchConfig            `json:"branches"`
+	Head     ref.MarshalableRef                       `json:"head"`
+	Remotes  *concurrentmap.Map[string, Remote]       `json:"remotes"`
+	Backups  *concurrentmap.Map[string, Remote]       `json:"backups"`
+	Branches *concurrentmap.Map[string, BranchConfig] `json:"branches"`
 	// |staged|, |working|, and |merge| are legacy fields left over from when Dolt repos stored this info in the repo
 	// state file, not in the DB directly. They're still here so that we can migrate existing repositories forward to the
 	// new storage format, but they should be used only for this purpose and are no longer written.
@@ -84,13 +84,13 @@ type RepoState struct {
 // repoStateLegacy only exists to unmarshall legacy repo state files, since the JSON marshaller can't work with
 // unexported fields
 type repoStateLegacy struct {
-	Head     ref.MarshalableRef                 `json:"head"`
-	Remotes  *concurrentmap.Map[string, Remote] `json:"remotes"`
-	Backups  *concurrentmap.Map[string, Remote] `json:"backups"`
-	Branches map[string]BranchConfig            `json:"branches"`
-	Staged   string                             `json:"staged,omitempty"`
-	Working  string                             `json:"working,omitempty"`
-	Merge    *mergeState                        `json:"merge,omitempty"`
+	Head     ref.MarshalableRef                       `json:"head"`
+	Remotes  *concurrentmap.Map[string, Remote]       `json:"remotes"`
+	Backups  *concurrentmap.Map[string, Remote]       `json:"backups"`
+	Branches *concurrentmap.Map[string, BranchConfig] `json:"branches"`
+	Staged   string                                   `json:"staged,omitempty"`
+	Working  string                                   `json:"working,omitempty"`
+	Merge    *mergeState                              `json:"merge,omitempty"`
 }
 
 // repoStateLegacyFromRepoState creates a new repoStateLegacy from a RepoState file. Only for testing.
@@ -128,6 +128,9 @@ func (rs *repoStateLegacy) toRepoState() *RepoState {
 	if newRS.Backups == nil {
 		newRS.Backups = concurrentmap.New[string, Remote]()
 	}
+	if newRS.Branches == nil {
+		newRS.Branches = concurrentmap.New[string, BranchConfig]()
+	}
 
 	return newRS
 }
@@ -170,7 +173,7 @@ func CloneRepoState(fs filesys.ReadWriteFS, r Remote) (*RepoState, error) {
 		staged:   hashStr,
 		working:  hashStr,
 		Remotes:  remotes,
-		Branches: make(map[string]BranchConfig),
+		Branches: concurrentmap.New[string, BranchConfig](),
 		Backups:  concurrentmap.New[string, Remote](),
 	}
 
@@ -192,7 +195,7 @@ func CreateRepoState(fs filesys.ReadWriteFS, br string) (*RepoState, error) {
 	rs := &RepoState{
 		Head:     ref.MarshalableRef{Ref: headRef},
 		Remotes:  concurrentmap.New[string, Remote](),
-		Branches: make(map[string]BranchConfig),
+		Branches: concurrentmap.New[string, BranchConfig](),
 		Backups:  concurrentmap.New[string, Remote](),
 	}
 

go/libraries/doltcore/sqle/database.go

@@ -418,7 +418,7 @@ func (db Database) getTableInsensitive(ctx *sql.Context, head *doltdb.Commit, ds
 		adapter := dsess.NewSessionStateAdapter(
 			sess, db.RevisionQualifiedName(),
 			concurrentmap.New[string, env.Remote](),
-			map[string]env.BranchConfig{},
+			concurrentmap.New[string, env.BranchConfig](),
 			concurrentmap.New[string, env.Remote]())
 		ws, err := sess.WorkingSet(ctx, db.RevisionQualifiedName())
 		if err != nil {

go/libraries/doltcore/sqle/dsess/database_session_state.go

@@ -42,7 +42,7 @@ type InitialDbState struct {
 	ReadOnly bool
 	DbData   env.DbData
 	Remotes  *concurrentmap.Map[string, env.Remote]
-	Branches map[string]env.BranchConfig
+	Branches *concurrentmap.Map[string, env.BranchConfig]
 	Backups  *concurrentmap.Map[string, env.Remote]
 
 	// If err is set, this InitialDbState is partially invalid, but may be

go/libraries/doltcore/sqle/dsess/session_state_adapter.go

@@ -34,7 +34,7 @@ type SessionStateAdapter struct {
 	dbName   string
 	remotes  *concurrentmap.Map[string, env.Remote]
 	backups  *concurrentmap.Map[string, env.Remote]
-	branches map[string]env.BranchConfig
+	branches *concurrentmap.Map[string, env.BranchConfig]
 }
 
 func (s SessionStateAdapter) SetCWBHeadRef(ctx context.Context, newRef ref.MarshalableRef) error {
@@ -45,9 +45,9 @@ var _ env.RepoStateReader = SessionStateAdapter{}
 var _ env.RepoStateWriter = SessionStateAdapter{}
 var _ env.RootsProvider = SessionStateAdapter{}
 
-func NewSessionStateAdapter(session *DoltSession, dbName string, remotes *concurrentmap.Map[string, env.Remote], branches map[string]env.BranchConfig, backups *concurrentmap.Map[string, env.Remote]) SessionStateAdapter {
+func NewSessionStateAdapter(session *DoltSession, dbName string, remotes *concurrentmap.Map[string, env.Remote], branches *concurrentmap.Map[string, env.BranchConfig], backups *concurrentmap.Map[string, env.Remote]) SessionStateAdapter {
 	if branches == nil {
-		branches = make(map[string]env.BranchConfig)
+		branches = concurrentmap.New[string, env.BranchConfig]()
 	}
 	return SessionStateAdapter{session: session, dbName: dbName, remotes: remotes, branches: branches, backups: backups}
 }
@@ -96,12 +96,12 @@ func (s SessionStateAdapter) GetBackups() (*concurrentmap.Map[string, env.Remote
 	return s.backups, nil
 }
 
-func (s SessionStateAdapter) GetBranches() (map[string]env.BranchConfig, error) {
+func (s SessionStateAdapter) GetBranches() (*concurrentmap.Map[string, env.BranchConfig], error) {
 	return s.branches, nil
 }
 
 func (s SessionStateAdapter) UpdateBranch(name string, new env.BranchConfig) error {
-	s.branches[name] = new
+	s.branches.Set(name, new)
 
 	fs, err := s.session.Provider().FileSystemForDatabase(s.dbName)
 	if err != nil {
@@ -112,7 +112,7 @@ func (s SessionStateAdapter) UpdateBranch(name string, new env.BranchConfig) err
 	if err != nil {
 		return err
 	}
-	repoState.Branches[name] = new
+	repoState.Branches.Set(name, new)
 
 	return repoState.Save(fs)
 }

go/libraries/doltcore/sqle/dtables/branches_table.go

@@ -218,7 +218,7 @@ func (itr *BranchItr) Next(ctx *sql.Context) (sql.Row, error) {
 
 		remoteName := ""
 		branchName := ""
-		branch, ok := branches[name]
+		branch, ok := branches.Get(name)
 		if ok {
 			remoteName = branch.Remote
 			branchName = branch.Merge.Ref.GetPath()