Skip to content

Commit

Permalink
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
/go/libraries/doltcore/sqle/dsess: parallelize new sqle db stuff
Browse files Browse the repository at this point in the history
coffeegoddd committed Jan 14, 2025
1 parent 581c2f9 commit 1939e96
Showing 2 changed files with 72 additions and 45 deletions.
47 changes: 26 additions & 21 deletions go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go
Original file line number Diff line number Diff line change
@@ -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,38 @@ 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
}
return r.IterTables(ctx, func(tableName doltdb.TableName, table *doltdb.Table, sch schema.Schema) (bool, error) {
if !schema.HasAutoIncrement(sch) {
return false, nil
}

tableNameStr := tableName.ToLower().Name
if oldValue, loaded := a.sequences.LoadOrStore(tableNameStr, seq); loaded && seq > oldValue.(uint64) {
a.sequences.Store(tableNameStr, seq)
}
seq, iErr := table.GetAutoIncrementValue(egCtx)
if iErr != nil {
return true, iErr
}

return false, nil
})
tableNameStr := tableName.ToLower().Name
if oldValue, loaded := a.sequences.LoadOrStore(tableNameStr, seq); loaded && seq > oldValue.(uint64) {
a.sequences.Store(tableNameStr, seq)
}

if err != nil {
return err
}
return false, nil
})
})
}

return nil
return eg.Wait()
}
70 changes: 46 additions & 24 deletions go/libraries/doltcore/sqle/dsess/globalstate.go
100755 → 100644
Original file line number Diff line number Diff line change
@@ -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"
@@ -36,39 +37,60 @@ func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.Dol
return GlobalStateImpl{}, err
}

eg, egCtx := errgroup.WithContext(ctx)
eg.SetLimit(128)

rootRefs := make([]ref.DoltRef, 0, len(branches)+len(remotes))
rootRefs = append(rootRefs, branches...)
rootRefs = append(rootRefs, remotes...)

var roots []doltdb.Rootish
rootChan := make(chan doltdb.Rootish, len(rootRefs))
for _, b := range rootRefs {
switch b.GetType() {
case ref.BranchRefType:
wsRef, err := ref.WorkingSetRefForHead(b)
if err != nil {
return GlobalStateImpl{}, err
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)
if err != nil {
return GlobalStateImpl{}, err
switch b.GetType() {
case ref.BranchRefType:
wsRef, rerr := ref.WorkingSetRefForHead(b)
if rerr != nil {
return rerr
}
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, rerr := db.ResolveWorkingSet(egCtx, wsRef)
if rerr == doltdb.ErrWorkingSetNotFound {
// use the branch head if there isn't a working set for it
cm, rerr := db.ResolveCommitRef(egCtx, b)
if rerr != nil {
return rerr
}
rootChan <- cm
} else if err != nil {
return rerr
} else {
rootChan <- ws
}
case ref.RemoteRefType:
cm, rerr := db.ResolveCommitRef(egCtx, b)
if rerr != nil {
return rerr
}
rootChan <- cm
}
roots = append(roots, cm)
}
return nil
})
}

close(rootChan)
err = eg.Wait()
if err != nil {
return GlobalStateImpl{}, err
}

var roots []doltdb.Rootish
for root := range rootChan {
roots = append(roots, root)
}

tracker, err := NewAutoIncrementTracker(ctx, dbName, roots...)

0 comments on commit 1939e96

Please sign in to comment.