diff --git a/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go b/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go index 4dda09c08a3..10415505e18 100644 --- a/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go +++ b/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go @@ -23,6 +23,7 @@ import ( "github.com/dolthub/go-mysql-server/sql" gmstypes "github.com/dolthub/go-mysql-server/sql/types" + "golang.org/x/sync/errgroup" "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/doltdb/durable" @@ -398,34 +399,39 @@ func (a *AutoIncrementTracker) AcquireTableLock(ctx *sql.Context, tableName stri } func (a *AutoIncrementTracker) InitWithRoots(ctx context.Context, roots ...doltdb.Rootish) error { - for _, root := range roots { - r, err := root.ResolveRootValue(ctx) - if err != nil { - return err - } + eg, egCtx := errgroup.WithContext(ctx) + eg.SetLimit(128) - err = r.IterTables(ctx, func(tableName doltdb.TableName, table *doltdb.Table, sch schema.Schema) (bool, error) { - if !schema.HasAutoIncrement(sch) { - return false, nil + for _, root := range roots { + eg.Go(func() error { + if egCtx.Err() != nil { + return egCtx.Err() } - seq, err := table.GetAutoIncrementValue(ctx) - if err != nil { - return true, err + r, rerr := root.ResolveRootValue(egCtx) + if rerr != nil { + return rerr } - tableNameStr := tableName.ToLower().Name - if oldValue, loaded := a.sequences.LoadOrStore(tableNameStr, seq); loaded && seq > oldValue.(uint64) { - a.sequences.Store(tableNameStr, seq) - } + return r.IterTables(egCtx, func(tableName doltdb.TableName, table *doltdb.Table, sch schema.Schema) (bool, error) { + if !schema.HasAutoIncrement(sch) { + return false, nil + } - return false, nil - }) + seq, iErr := table.GetAutoIncrementValue(egCtx) + if iErr != nil { + return true, iErr + } - if err != nil { - return err - } + tableNameStr := tableName.ToLower().Name + if oldValue, loaded := a.sequences.LoadOrStore(tableNameStr, seq); loaded && seq > oldValue.(uint64) { + a.sequences.Store(tableNameStr, seq) + } + + return false, nil + }) + }) } - return nil + return eg.Wait() } diff --git a/go/libraries/doltcore/sqle/dsess/globalstate.go b/go/libraries/doltcore/sqle/dsess/globalstate.go old mode 100755 new mode 100644 index dd675d099bc..8315997460b --- a/go/libraries/doltcore/sqle/dsess/globalstate.go +++ b/go/libraries/doltcore/sqle/dsess/globalstate.go @@ -19,6 +19,7 @@ import ( "sync" "github.com/dolthub/go-mysql-server/sql" + "golang.org/x/sync/errgroup" "github.com/dolthub/dolt/go/libraries/doltcore/doltdb" "github.com/dolthub/dolt/go/libraries/doltcore/ref" @@ -40,35 +41,51 @@ func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.Dol rootRefs = append(rootRefs, branches...) rootRefs = append(rootRefs, remotes...) - var roots []doltdb.Rootish - for _, b := range rootRefs { - switch b.GetType() { - case ref.BranchRefType: - wsRef, err := ref.WorkingSetRefForHead(b) - if err != nil { - return GlobalStateImpl{}, err + roots := make([]doltdb.Rootish, len(rootRefs)) + eg, egCtx := errgroup.WithContext(ctx) + eg.SetLimit(128) + + for idx, b := range rootRefs { + idx, b := idx, b + eg.Go(func() error { + if egCtx.Err() != nil { + return egCtx.Err() } - ws, err := db.ResolveWorkingSet(ctx, wsRef) - if err == doltdb.ErrWorkingSetNotFound { - // use the branch head if there isn't a working set for it - cm, err := db.ResolveCommitRef(ctx, b) + switch b.GetType() { + case ref.BranchRefType: + wsRef, err := ref.WorkingSetRefForHead(b) if err != nil { - return GlobalStateImpl{}, err + return err } - roots = append(roots, cm) - } else if err != nil { - return GlobalStateImpl{}, err - } else { - roots = append(roots, ws) - } - case ref.RemoteRefType: - cm, err := db.ResolveCommitRef(ctx, b) - if err != nil { - return GlobalStateImpl{}, err + + ws, err := db.ResolveWorkingSet(egCtx, wsRef) + if err == doltdb.ErrWorkingSetNotFound { + // use the branch head if there isn't a working set for it + cm, err := db.ResolveCommitRef(egCtx, b) + if err != nil { + return err + } + roots[idx] = cm + } else if err != nil { + return err + } else { + roots[idx] = ws + } + case ref.RemoteRefType: + cm, err := db.ResolveCommitRef(egCtx, b) + if err != nil { + return err + } + roots[idx] = cm } - roots = append(roots, cm) - } + return nil + }) + } + + err = eg.Wait() + if err != nil { + return GlobalStateImpl{}, err } tracker, err := NewAutoIncrementTracker(ctx, dbName, roots...)