Merge pull request #8740 from dolthub/db/parallelize

/go/libraries/doltcore/sql/dsess: parallelize sql.NewDatabase work
dolthub · Jan 14, 2025 · 99c7cca · 99c7cca
2 parents d05b6a8 + 93355c5
commit 99c7cca
Show file tree

Hide file tree

Showing 2 changed files with 68 additions and 45 deletions.
diff --git 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
@@ -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...)