Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

/go/libraries/doltcore/sql/dsess: parallelize sql.NewDatabase work #8740

Merged
merged 10 commits into from
Jan 14, 2025
Merged

/go/libraries/doltcore/sql/dsess: parallelize sql.NewDatabase work #8740

Show file tree
Hide file tree
Changes from 9 commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
d869c70
/go/libraries/doltcore/sqle/dsess: parallelize new sqle db stuff
coffeegoddd Jan 14, 2025
eb23dcd
/go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go: wip, nil …
coffeegoddd Jan 14, 2025
fa6c12f
/go/libraries/doltcore/sqle/dsess/globalstate.go: try just changing ai
coffeegoddd Jan 14, 2025
53ba7ca
temp: add some timing logs
coffeegoddd Jan 14, 2025
4035348
/go/libraries/doltcore/sqle/dsess/globalstate.go: parallelize global …
coffeegoddd Jan 14, 2025
bff342a
/go/libraries/doltcore/sqle/dsess: remove logs
coffeegoddd Jan 14, 2025
ddecc47
/go/libraries/doltcore/sqle/dsess/globalstate.go: check for ctx error
coffeegoddd Jan 14, 2025
8c1eda8
/go/libraries/doltcore/sqle/dsess/globalstate.go: trying to fix close…
coffeegoddd Jan 14, 2025
a819aa2
/go/libraries/doltcore/sqle/dsess/globalstate.go: fix again
coffeegoddd Jan 14, 2025
93355c5
/go/libraries/doltcore/sqle/dsess/globalstate.go: pr feedback
coffeegoddd Jan 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 27 additions & 21 deletions go/libraries/doltcore/sqle/dsess/autoincrement_tracker.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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()
}
73 changes: 50 additions & 23 deletions go/libraries/doltcore/sqle/dsess/globalstate.go
100755 → 100644
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand All @@ -40,35 +41,61 @@ func NewGlobalStateStoreForDb(ctx context.Context, dbName string, db *doltdb.Dol
rootRefs = append(rootRefs, branches...)
rootRefs = append(rootRefs, remotes...)

var roots []doltdb.Rootish
rootRefsChan := make(chan doltdb.Rootish, len(rootRefs))
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Instead of a channel here, just allocate roots here:

roots := make([]doltdb.Rootish, len(rootRefs))

and then have the loop capture the iteration variable and write the root into the right place:

for i, b := range rootRefs {
    eg.Go(func() error {
        ...
            roots[i] = cm
        ...
            roots[i] = ws
        ...
    })
}

Probably best to eg.SetLimit() on this errgroup as well.

eg, egCtx := errgroup.WithContext(ctx)
wg := sync.WaitGroup{}

for _, b := range rootRefs {
switch b.GetType() {
case ref.BranchRefType:
wsRef, err := ref.WorkingSetRefForHead(b)
if err != nil {
return GlobalStateImpl{}, err
wg.Add(1)
eg.Go(func() error {
defer wg.Done()
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
}
rootRefsChan <- cm
} else if err != nil {
return err
} else {
rootRefsChan <- ws
}
case ref.RemoteRefType:
cm, err := db.ResolveCommitRef(egCtx, b)
if err != nil {
return err
}
rootRefsChan <- cm
}
roots = append(roots, cm)
}
return nil
})
}

// prevent sending on closed channel
wg.Wait()
close(rootRefsChan)

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

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

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