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/LegendServerGlobalState.java

@@ -26,6 +26,7 @@
 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.stream.Stream;
@@ -64,6 +65,16 @@ 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)
+        );
+    }
+
     @Override
     public void logInfo(String message)
     {