Skip to content

Commit 3f94b95

Browse files
sheetalkamatsheetalkamat
committed
Parallel
1 parent 01fb67d commit 3f94b95

File tree

2 files changed

+128
-74
lines changed

2 files changed

+128
-74
lines changed

‎internal/lsp/server.go‎

Lines changed: 127 additions & 65 deletions
Original file line numberDiff line numberDiff line change
@@ -592,6 +592,11 @@ type projectAndTextDocumentPosition struct {
592592
Position lsproto.Position
593593
}
594594

595+
type response[Resp any] struct {
596+
complete bool
597+
result Resp
598+
}
599+
595600
func registerMultiProjectReferenceRequestHandler[Req lsproto.HasTextDocumentPosition, Resp any](
596601
handlers handlerMap,
597602
info lsproto.RequestInfo[Req, Resp],
@@ -611,51 +616,85 @@ func registerMultiProjectReferenceRequestHandler[Req lsproto.HasTextDocumentPosi
611616
}
612617
defer s.recover(req)
613618

614-
var queue []projectAndTextDocumentPosition
615-
var results collections.OrderedMap[tspath.Path, Resp]
619+
var results collections.SyncMap[tspath.Path, *response[Resp]]
616620
var defaultDefinition *ls.NonLocalDefinition
621+
canSearchProject := func(project *project.Project) bool {
622+
_, searched := results.Load(project.Id())
623+
return !searched
624+
}
625+
wg := core.NewWorkGroup(false)
626+
var errMu sync.Mutex
627+
var enqueueItem func(item projectAndTextDocumentPosition)
628+
enqueueItem = func(item projectAndTextDocumentPosition) {
629+
var response response[Resp]
630+
if _, loaded := results.LoadOrStore(item.project.Id(), &response); loaded {
631+
return
632+
}
633+
wg.Queue(func() {
634+
if ctx.Err() != nil {
635+
return
636+
}
617637

618-
searchPosition := func(project *project.Project, ls *ls.LanguageService, uri lsproto.DocumentUri, position lsproto.Position) (Resp, error) {
619-
originalNode, symbolsAndEntries, ok := ls.ProvideSymbolsAndEntries(ctx, uri, position, info.Method == lsproto.MethodTextDocumentRename)
620-
if ok {
621-
for _, entry := range symbolsAndEntries {
622-
// Find the default definition that can be in another project
623-
// Later we will use this load ancestor tree that references this location and expand search
624-
if project == defaultProject && defaultDefinition == nil {
625-
defaultDefinition = ls.GetNonLocalDefinition(ctx, entry)
638+
// Process the item
639+
ls := item.ls
640+
if ls == nil {
641+
// Get it now
642+
ls = s.session.GetLanguageServiceForProjectWithFile(ctx, item.project, item.Uri)
643+
if ls == nil {
644+
return
626645
}
627-
ls.ForEachOriginalDefinitionLocation(ctx, entry, func(uri lsproto.DocumentUri, position lsproto.Position) {
628-
// Get default configured project for this file
629-
defProjects, errProjects := s.session.GetProjectsForFile(ctx, uri)
630-
if errProjects != nil {
631-
return
646+
}
647+
originalNode, symbolsAndEntries, ok := ls.ProvideSymbolsAndEntries(ctx, item.Uri, item.Position, info.Method == lsproto.MethodTextDocumentRename)
648+
if ok {
649+
for _, entry := range symbolsAndEntries {
650+
// Find the default definition that can be in another project
651+
// Later we will use this load ancestor tree that references this location and expand search
652+
if item.project == defaultProject && defaultDefinition == nil {
653+
defaultDefinition = ls.GetNonLocalDefinition(ctx, entry)
632654
}
633-
for _, defProject := range defProjects {
634-
// Optimization: don't enqueue if will be discarded
635-
if !results.Has(defProject.Id()) {
636-
queue = append(queue, projectAndTextDocumentPosition{
637-
project: defProject,
638-
Uri: uri,
639-
Position: position,
640-
})
655+
ls.ForEachOriginalDefinitionLocation(ctx, entry, func(uri lsproto.DocumentUri, position lsproto.Position) {
656+
// Get default configured project for this file
657+
defProjects, errProjects := s.session.GetProjectsForFile(ctx, uri)
658+
if errProjects != nil {
659+
return
641660
}
642-
}
643-
})
661+
for _, defProject := range defProjects {
662+
// Optimization: don't enqueue if will be discarded
663+
if canSearchProject(defProject) {
664+
enqueueItem(projectAndTextDocumentPosition{
665+
project: defProject,
666+
Uri: uri,
667+
Position: position,
668+
})
669+
}
670+
}
671+
})
672+
}
644673
}
645-
}
646-
return fn(s, ctx, ls, params, originalNode, symbolsAndEntries)
674+
675+
if result, errSearch := fn(s, ctx, ls, params, originalNode, symbolsAndEntries); errSearch == nil {
676+
response.complete = true
677+
response.result = result
678+
} else {
679+
errMu.Lock()
680+
defer errMu.Unlock()
681+
if err != nil {
682+
err = errSearch
683+
}
684+
}
685+
})
647686
}
648687

649688
// Initial set of projects and locations in the queue, starting with default project
650-
queue = append(queue, projectAndTextDocumentPosition{
689+
enqueueItem(projectAndTextDocumentPosition{
651690
project: defaultProject,
652691
ls: defaultLs,
653692
Uri: params.TextDocumentURI(),
654693
Position: params.TextDocumentPosition(),
655694
})
656695
for _, project := range allProjects {
657696
if project != defaultProject {
658-
queue = append(queue, projectAndTextDocumentPosition{
697+
enqueueItem(projectAndTextDocumentPosition{
659698
project: project,
660699
// TODO!! symlinks need to change the URI
661700
Uri: params.TextDocumentURI(),
@@ -664,69 +703,92 @@ func registerMultiProjectReferenceRequestHandler[Req lsproto.HasTextDocumentPosi
664703
}
665704
}
666705

667-
// Outer loop - to complete work if more is added after completing existing queue
668-
for len(queue) > 0 {
669-
// Process existing known projects first
670-
for len(queue) > 0 {
671-
if ctx.Err() != nil {
672-
return ctx.Err()
673-
}
674-
675-
item := queue[0]
676-
queue = queue[1:]
677-
678-
if results.Has(item.project.Id()) {
679-
continue
680-
}
681-
682-
// Process the item
683-
ls := item.ls
684-
if ls == nil {
685-
// Get it now
686-
ls = s.session.GetLanguageServiceForProjectWithFile(ctx, item.project, item.Uri)
687-
if ls == nil {
688-
continue
706+
getResultsIterator := func() iter.Seq[Resp] {
707+
return func(yield func(Resp) bool) {
708+
var seenProjects collections.SyncSet[tspath.Path]
709+
if response, loaded := results.Load(defaultProject.Id()); loaded && response.complete {
710+
if !yield(response.result) {
711+
return
689712
}
690713
}
691-
if result, err := searchPosition(item.project, ls, item.Uri, item.Position); err == nil {
692-
results.Set(item.project.Id(), result)
693-
} else {
694-
return err
714+
seenProjects.Add(defaultProject.Id())
715+
for _, project := range allProjects {
716+
if seenProjects.AddIfAbsent(project.Id()) {
717+
if response, loaded := results.Load(project.Id()); loaded && response.complete {
718+
if !yield(response.result) {
719+
return
720+
}
721+
}
722+
}
695723
}
724+
results.Range(func(key tspath.Path, response *response[Resp]) bool {
725+
if seenProjects.AddIfAbsent(key) && response.complete {
726+
return yield(response.result)
727+
}
728+
return true
729+
})
696730
}
731+
}
697732

733+
// Outer loop - to complete work if more is added after completing existing queue
734+
for {
735+
// Process existing known projects first
736+
wg.RunAndWait()
737+
if ctx.Err() != nil {
738+
return ctx.Err()
739+
}
740+
// No need to use mu here since we are not in parallel at this point
741+
if err != nil {
742+
return err
743+
}
744+
745+
wg = core.NewWorkGroup(false)
746+
hasMoreWork := false
698747
if defaultDefinition != nil {
699-
if err := s.session.ForEachProjectLocationLoadingProjectTree(
748+
requestedProjectTrees := make(map[tspath.Path]struct{})
749+
results.Range(func(key tspath.Path, response *response[Resp]) bool {
750+
if response.complete {
751+
requestedProjectTrees[key] = struct{}{}
752+
}
753+
return true
754+
})
755+
756+
// Load more projects based on default definition found
757+
if errFromLoadingTree := s.session.ForEachProjectLocationLoadingProjectTree(
700758
ctx,
701-
results.Keys(),
759+
requestedProjectTrees,
702760
defaultDefinition,
703761
// Can loop forever without this (enqueue here, dequeue above, repeat)
704-
func(projectFromSnapshot *project.Project) bool {
705-
return !results.Has(projectFromSnapshot.Id())
706-
},
762+
canSearchProject,
707763
func(project *project.Project, uri lsproto.DocumentUri, position lsproto.Position) {
708764
// Enqueue the project and location for further processing
709-
queue = append(queue, projectAndTextDocumentPosition{
765+
enqueueItem(projectAndTextDocumentPosition{
710766
project: project,
711767
Uri: uri,
712768
Position: position,
713769
})
770+
hasMoreWork = true
714771
},
715-
); err != nil {
716-
return err
772+
); errFromLoadingTree != nil {
773+
return errFromLoadingTree
717774
}
718775
}
776+
if !hasMoreWork {
777+
break
778+
}
719779
}
720780

721781
var resp Resp
722782
if results.Size() > 1 {
723-
resp = combineResults(results.Values())
783+
resp = combineResults(getResultsIterator())
724784
} else {
725-
for value := range results.Values() {
785+
// Single result, return that directly
786+
for value := range getResultsIterator() {
726787
resp = value
727788
break
728789
}
729790
}
791+
730792
s.sendResult(req.ID, resp)
731793
return nil
732794
}

‎internal/project/session.go‎

Lines changed: 1 addition & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@ package project
33
import (
44
"context"
55
"fmt"
6-
"iter"
76
"slices"
87
"strings"
98
"sync"
@@ -490,18 +489,11 @@ func (s *Session) projectContainsFile(project *Project, uri lsproto.DocumentUri)
490489

491490
func (s *Session) ForEachProjectLocationLoadingProjectTree(
492491
ctx context.Context,
493-
projectsReferenced iter.Seq[tspath.Path],
492+
requestedProjectTrees map[tspath.Path]struct{},
494493
defaultDefinition *ls.NonLocalDefinition,
495494
canIterateProject func(project *Project) bool,
496495
handleLocation func(*Project, lsproto.DocumentUri, lsproto.Position),
497496
) error {
498-
var requestedProjectTrees map[tspath.Path]struct{}
499-
for path := range projectsReferenced {
500-
if requestedProjectTrees == nil {
501-
requestedProjectTrees = make(map[tspath.Path]struct{})
502-
}
503-
requestedProjectTrees[path] = struct{}{}
504-
}
505497
snapshot := s.getSnapshot(
506498
ctx,
507499
snapshotChangeRequest{requestedProjectTrees: requestedProjectTrees},

0 commit comments

Comments
 (0)