Improve init processing

legend-engine-ide-lsp-default-extensions/src/main/java/org/finos/legend/engine/ide/lsp/extension/AbstractLSPGrammarExtension.java

@@ -30,6 +30,7 @@
 import java.util.Optional;
 import java.util.ServiceLoader;
 import java.util.Set;
+import java.util.concurrent.CompletableFuture;
 import java.util.function.Consumer;
 import java.util.stream.Collectors;
 import org.eclipse.collections.api.factory.Lists;
@@ -325,30 +326,42 @@ public CompileResult getCompileResult(SectionState sectionState)
     {
         DocumentState documentState = sectionState.getDocumentState();
         GlobalState globalState = documentState.getGlobalState();
-        return globalState.getProperty(COMPILE_RESULT, () -> tryCompile(globalState, documentState, sectionState));
+        // when looking for compile results, there might be another thread working on it already
+        // if current thread is the one that sets the completable future, then do the actual compilation and set it on the completable future
+        // if current thread does not set the completable future, then just join and wait for result from another thread
+        CompletableFuture<CompileResult> maybeCompileResultFuture = new CompletableFuture<>();
+        CompletableFuture<CompileResult> compileResultFuture = globalState.getProperty(COMPILE_RESULT, () -> maybeCompileResultFuture);
+
+        if (compileResultFuture == maybeCompileResultFuture)
+        {
+            compileResultFuture.complete(this.tryCompile(globalState, documentState, sectionState));
+        }
+
+        return compileResultFuture.join();
     }
 
     protected CompileResult tryCompile(GlobalState globalState, DocumentState documentState, SectionState sectionState)
     {
+        long started = System.currentTimeMillis();
         globalState.logInfo("Starting compilation");
         PureModelContextData pureModelContextData = null;
         try
         {
             pureModelContextData = buildPureModelContextData(globalState);
             PureModel pureModel = Compiler.compile(pureModelContextData, DeploymentMode.PROD, "");
-            globalState.logInfo("Compilation completed successfully");
+            globalState.logInfo("Compilation completed successfully in " + (System.currentTimeMillis() - started) + "ms");
             return new CompileResult(pureModel, pureModelContextData);
         }
         catch (EngineException e)
         {
             SourceInformation sourceInfo = e.getSourceInformation();
             if (isValidSourceInfo(sourceInfo))
             {
-                globalState.logInfo("Compilation completed with error " + "(" + sourceInfo.sourceId + " " + SourceInformationUtil.toLocation(sourceInfo) + "): " + e.getMessage());
+                globalState.logInfo("Compilation completed in " + (System.currentTimeMillis() - started) + "ms with error " + "(" + sourceInfo.sourceId + " " + SourceInformationUtil.toLocation(sourceInfo) + "): " + e.getMessage());
             }
             else
             {
-                globalState.logInfo("Compilation completed with error: " + e.getMessage());
+                globalState.logInfo("Compilation completed in " + (System.currentTimeMillis() - started) + "ms with error: " + e.getMessage());
                 globalState.logWarning("Invalid source information for compilation error");
                 LOGGER.warn("Invalid source information in exception during compilation requested for section {} of {}: {}", sectionState.getSectionNumber(), documentState.getDocumentId(), (sourceInfo == null) ? null : sourceInfo.getMessage(), e);
             }

legend-engine-ide-lsp-server/src/main/java/org/finos/legend/engine/ide/lsp/server/LegendLanguageServer.java

@@ -248,13 +248,14 @@ public String getProjectVersion()
     @Override
     public void initialized(InitializedParams params)
     {
+        long start = System.currentTimeMillis();
         checkReady();
         this.classpathFactory.initialize(this);
         CompletableFuture<Void> initializeExtensions = this.initializeExtensions();
         CompletableFuture<Void> engineServerUrl = this.initializeEngineServerUrl();
 
         CompletableFuture.allOf(initializeExtensions, engineServerUrl)
-                .thenRun(() -> this.logInfoToClient("Extension finished post-initialization"));
+                .thenRun(() -> this.logInfoToClient("Extension finished post-initialization in " + (System.currentTimeMillis() - start) + "ms"));
 
     }
 
@@ -292,10 +293,16 @@ private CompletableFuture<Void> initializeExtensions()
 
         return this.classpathFactory.create(Collections.unmodifiableSet(this.rootFolders))
                 .thenAccept(this.extensionGuard::initialize)
-                .thenRun(this.extensionGuard.wrapOnClasspath(this::reprocessDocuments))
+                .thenCompose(_x -> this.reprocessDocuments())
                 .thenRun(this.legendLanguageService::loadVirtualFileSystemContent)
                 // trigger compilation
-                .thenRun(this.extensionGuard.wrapOnClasspath(() -> this.globalState.forEachDocumentState(this.textDocumentService::getLegendDiagnostics)))
+                .thenCompose(_x -> this.globalState.forEachDocumentStateParallel(x ->
+                {
+                    long diagnosticStarted = System.currentTimeMillis();
+                    this.textDocumentService.getLegendDiagnostics(x);
+                    LOGGER.info("Diagnostics computed for {} took {}ms", x.getDocumentId(), System.currentTimeMillis() - diagnosticStarted);
+
+                }))
                 .thenRun(() ->
                 {
                     LanguageClient languageClient = this.getLanguageClient();
@@ -314,10 +321,14 @@ private CompletableFuture<Void> initializeExtensions()
                 });
     }
 
-    private void reprocessDocuments()
+    private CompletableFuture<Void> reprocessDocuments()
     {
-        this.globalState.forEachDocumentState(x -> ((LegendServerGlobalState.LegendServerDocumentState) x).recreateSectionStates());
-        this.globalState.clearProperties();
+        return this.globalState.forEachDocumentStateParallel(x ->
+        {
+            long startTime = System.currentTimeMillis();
+            ((LegendServerGlobalState.LegendServerDocumentState) x).recreateSectionStates();
+            LOGGER.info("Reprocessing {} took {}ms", x.getDocumentId(), System.currentTimeMillis() - startTime);
+        }).thenRun(this.globalState::clearProperties);
     }
 
     @Override
@@ -484,10 +495,10 @@ public <T> CompletableFuture<T> supplyPossiblyAsync(Supplier<T> supplier)
         return this.supplyPossiblyAsync_internal(this.extensionGuard.wrapOnClasspath(supplier));
     }
 
-    void runPossiblyAsync(Runnable runnable)
+    CompletableFuture<?> runPossiblyAsync(Runnable runnable)
     {
         checkReady();
-        this.runPossiblyAsync_internal(this.extensionGuard.wrapOnClasspath(runnable));
+        return this.runPossiblyAsync_internal(this.extensionGuard.wrapOnClasspath(runnable));
     }
 
     private CompletableFuture<?> runPossiblyAsync_internal(Runnable work)

legend-engine-ide-lsp-server/src/main/java/org/finos/legend/engine/ide/lsp/server/LegendLanguageService.java

@@ -120,24 +120,21 @@ public CompletableFuture<String> replClasspath()
     @Override
     public CompletableFuture<List<LegendTest>> testCases()
     {
-        return this.server.supplyPossiblyAsync(() ->
-        {
-            List<LegendTest> commands = new ArrayList<>();
-
-            this.server.getGlobalState().forEachDocumentState(docState ->
-            {
-                docState.forEachSectionState(sectionState ->
+        return this.server.getGlobalState().collectFromDocumentStateParallel(docState ->
                 {
-                    LegendLSPGrammarExtension extension = sectionState.getExtension();
-                    if (extension != null)
-                    {
-                        commands.addAll(extension.testCases(sectionState));
-                    }
-                });
-            });
+                    List<LegendTest> commands = new ArrayList<>();
 
-            return commands;
-        });
+                    docState.forEachSectionState(sectionState ->
+                    {
+                        LegendLSPGrammarExtension extension = sectionState.getExtension();
+                        if (extension != null)
+                        {
+                            commands.addAll(extension.testCases(sectionState));
+                        }
+                    });
+                    return commands;
+                }
+        );
     }
 
     @Override

legend-engine-ide-lsp-server/src/main/java/org/finos/legend/engine/ide/lsp/server/LegendServerGlobalState.java

@@ -26,8 +26,11 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.finos.legend.engine.ide.lsp.extension.LegendLSPFeature;
 import org.finos.legend.engine.ide.lsp.extension.LegendLSPGrammarExtension;
@@ -64,6 +67,25 @@ public void forEachDocumentState(Consumer<? super DocumentState> consumer)
         this.docs.values().forEach(consumer);
     }
 
+    CompletableFuture<Void> forEachDocumentStateParallel(Consumer<? super DocumentState> consumer)
+    {
+        return CompletableFuture.allOf(
+                this.docs.values()
+                        .stream()
+                        .map(x -> this.server.runPossiblyAsync(() -> consumer.accept(x)))
+                        .toArray(CompletableFuture[]::new)
+        );
+    }
+
+    <T> CompletableFuture<List<T>> collectFromDocumentStateParallel(Function<? super DocumentState, List<T>> func)
+    {
+        return this.docs.values()
+                .stream()
+                .map(x -> this.server.supplyPossiblyAsync(() -> func.apply(x)))
+                .reduce((x, y) -> x.thenCombine(y, (r, l) -> Stream.concat(r.stream(), l.stream()).collect(Collectors.toList())))
+                .orElseGet(() -> CompletableFuture.completedFuture(List.of()));
+    }
+
     @Override
     public void logInfo(String message)
     {

legend-engine-ide-lsp-server/src/main/java/org/finos/legend/engine/ide/lsp/server/LegendWorkspaceService.java

@@ -21,6 +21,7 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Collectors;
 import org.eclipse.lsp4j.CreateFilesParams;
@@ -310,15 +311,10 @@ private void fileRenamed(LegendServerGlobalState globalState, String oldUri, Str
     @Override
     public CompletableFuture<Either<List<? extends SymbolInformation>, List<? extends WorkspaceSymbol>>> symbol(WorkspaceSymbolParams params)
     {
-        return this.server.supplyPossiblyAsync(() -> getWorkspaceSymbols(params));
-    }
-
-    private Either<List<? extends SymbolInformation>, List<? extends WorkspaceSymbol>> getWorkspaceSymbols(WorkspaceSymbolParams params)
-    {
-        List<WorkspaceSymbol> symbols = new ArrayList<>();
-
-        this.server.getGlobalState().forEachDocumentState(doc ->
+        return this.server.getGlobalState().collectFromDocumentStateParallel(doc ->
         {
+            List<WorkspaceSymbol> symbols = new ArrayList<>();
+
             doc.forEachSectionState(sec ->
             {
                 if (sec.getExtension() != null)
@@ -328,9 +324,9 @@ private Either<List<? extends SymbolInformation>, List<? extends WorkspaceSymbol
                             .forEach(declaration -> toWorkspaceSymbol(params, symbols, doc, declaration, null));
                 }
             });
-        });
 
-        return Either.forRight(symbols);
+            return symbols;
+        }).thenApply(Either::forRight);
     }
 
     private static void toWorkspaceSymbol(WorkspaceSymbolParams params, List<WorkspaceSymbol> symbols, DocumentState doc, LegendDeclaration declaration, WorkspaceSymbol parent)
@@ -368,65 +364,63 @@ private static void toWorkspaceSymbol(WorkspaceSymbolParams params, List<Workspa
     @Override
     public CompletableFuture<WorkspaceDiagnosticReport> diagnostic(WorkspaceDiagnosticParams params)
     {
-        return this.server.supplyPossiblyAsync(() ->
-        {
-            Map<String, String> previousResultIds = params.getPreviousResultIds()
-                    .stream()
-                    .collect(Collectors.toMap(PreviousResultId::getUri, PreviousResultId::getValue));
-
-            Map<String, Set<LegendDiagnostic>> previousResultIdToDiagnostic = this.previousResultIdToDiagnosticReference.get();
-            Map<String, Set<LegendDiagnostic>> resultIdToDiagnostic = new HashMap<>();
+        Map<String, String> previousResultIds = params.getPreviousResultIds()
+                .stream()
+                .collect(Collectors.toMap(PreviousResultId::getUri, PreviousResultId::getValue));
 
-            List<WorkspaceDocumentDiagnosticReport> items = new ArrayList<>();
+        Map<String, Set<LegendDiagnostic>> previousResultIdToDiagnostic = this.previousResultIdToDiagnosticReference.get();
+        Map<String, Set<LegendDiagnostic>> resultIdToDiagnostic = new ConcurrentHashMap<>();
 
-            this.server.getGlobalState().forEachDocumentState(d ->
-            {
-                String previousResultId = previousResultIds.getOrDefault(d.getDocumentId(), "");
-                Set<LegendDiagnostic> prevDiagnostic = previousResultIdToDiagnostic.get(previousResultId);
+        return this.server.getGlobalState().collectFromDocumentStateParallel(d ->
+                {
+                    String previousResultId = previousResultIds.getOrDefault(d.getDocumentId(), "");
+                    Set<LegendDiagnostic> prevDiagnostic = previousResultIdToDiagnostic.get(previousResultId);
 
-                LegendServerGlobalState.LegendServerDocumentState doc = (LegendServerGlobalState.LegendServerDocumentState) d;
-                Set<LegendDiagnostic> diagnostics = this.server.getTextDocumentService().getLegendDiagnostics(doc);
+                    LegendServerGlobalState.LegendServerDocumentState doc = (LegendServerGlobalState.LegendServerDocumentState) d;
+                    Set<LegendDiagnostic> diagnostics = this.server.getTextDocumentService().getLegendDiagnostics(doc);
 
-                String publishResultId = null;
+                    String publishResultId = null;
 
-                // no previous results
-                if (prevDiagnostic == null)
-                {
-                    // there are new diagnostics
-                    if (!diagnostics.isEmpty())
+                    // no previous results
+                    if (prevDiagnostic == null)
                     {
-                        publishResultId = UUID.randomUUID().toString();
-                        resultIdToDiagnostic.put(publishResultId, diagnostics);
+                        // there are new diagnostics
+                        if (!diagnostics.isEmpty())
+                        {
+                            publishResultId = UUID.randomUUID().toString();
+                            resultIdToDiagnostic.put(publishResultId, diagnostics);
+                        }
                     }
-                }
-                // there are previous results
-                else
-                {
-                    // diagnostics are different between previous and now
-                    if (!diagnostics.equals(prevDiagnostic))
+                    // there are previous results
+                    else
                     {
-                        publishResultId = UUID.randomUUID().toString();
-                        resultIdToDiagnostic.put(publishResultId, diagnostics);
+                        // diagnostics are different between previous and now
+                        if (!diagnostics.equals(prevDiagnostic))
+                        {
+                            publishResultId = UUID.randomUUID().toString();
+                            resultIdToDiagnostic.put(publishResultId, diagnostics);
+                        }
+                        // only track old diagnostics if same as new and non-empty (ie don't track empty ones)
+                        // otherwise, next time prevDiagnostic will be null, and only we start tracking again if there are new diagnostics
+                        else if (!diagnostics.isEmpty())
+                        {
+                            resultIdToDiagnostic.put(previousResultId, diagnostics);
+                        }
                     }
-                    // only track old diagnostics if same as new and non-empty (ie don't track empty ones)
-                    // otherwise, next time prevDiagnostic will be null, and only we start tracking again if there are new diagnostics
-                    else if (!diagnostics.isEmpty())
+
+                    if (publishResultId != null)
                     {
-                        resultIdToDiagnostic.put(previousResultId, diagnostics);
+                        WorkspaceFullDocumentDiagnosticReport fullReport = new WorkspaceFullDocumentDiagnosticReport(diagnostics.stream().map(LegendToLSPUtilities::toDiagnostic).collect(Collectors.toList()), doc.getDocumentId(), doc.getVersion());
+                        fullReport.setResultId(publishResultId);
+                        return List.of(new WorkspaceDocumentDiagnosticReport(fullReport));
                     }
-                }
 
-                if (publishResultId != null)
-                {
-                    WorkspaceFullDocumentDiagnosticReport fullReport = new WorkspaceFullDocumentDiagnosticReport(diagnostics.stream().map(LegendToLSPUtilities::toDiagnostic).collect(Collectors.toList()), doc.getDocumentId(), doc.getVersion());
-                    fullReport.setResultId(publishResultId);
-                    items.add(new WorkspaceDocumentDiagnosticReport(fullReport));
+                    return List.of();
                 }
-            });
-
+        ).thenApply(x ->
+        {
             this.previousResultIdToDiagnosticReference.compareAndSet(previousResultIdToDiagnostic, resultIdToDiagnostic);
-
-            return new WorkspaceDiagnosticReport(items);
+            return new WorkspaceDiagnosticReport(x);
         });
     }
 }