Merge branch 'performance-refactor' into performance-refactor-dss

.github/workflows/ci.yml

@@ -10,6 +10,7 @@ on:
 jobs:
   build:
     strategy:
+      fail-fast: false
       matrix:
         java: [11, 17]
         os: [ubuntu-latest, macos-latest]

src/main/java/software/amazon/smithy/lsp/DocumentLifecycleManager.java

@@ -6,27 +6,41 @@
 package software.amazon.smithy.lsp;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutionException;
+import java.util.logging.Logger;
 
 /**
- * Tracks asynchronous lifecycle tasks. Allows cancelling of an ongoing task
- * if a new task needs to be started
+ * Tracks asynchronous lifecycle tasks and client-managed documents.
+ * Allows cancelling of an ongoing task if a new task needs to be started.
  */
 final class DocumentLifecycleManager {
+    private static final Logger LOGGER = Logger.getLogger(DocumentLifecycleManager.class.getName());
     private final Map<String, CompletableFuture<Void>> tasks = new HashMap<>();
+    private final Set<String> managedDocumentUris = new HashSet<>();
 
     DocumentLifecycleManager() {
     }
 
+    Set<String> getManagedDocuments() {
+        return managedDocumentUris;
+    }
+
+    boolean isManaged(String uri) {
+        return getManagedDocuments().contains(uri);
+    }
+
     CompletableFuture<Void> getTask(String uri) {
         return tasks.get(uri);
     }
 
     void cancelTask(String uri) {
         if (tasks.containsKey(uri)) {
             CompletableFuture<Void> task = tasks.get(uri);
-            if (!task.isDone() && !task.isCancelled()) {
+            if (!task.isDone()) {
                 task.cancel(true);
             }
         }
@@ -36,9 +50,25 @@ void putTask(String uri, CompletableFuture<Void> future) {
         tasks.put(uri, future);
     }
 
+    void putOrComposeTask(String uri, CompletableFuture<Void> future) {
+        if (tasks.containsKey(uri)) {
+            tasks.computeIfPresent(uri, (k, v) -> v.thenCompose((unused) -> future));
+        } else {
+            tasks.put(uri, future);
+        }
+    }
+
     void cancelAllTasks() {
         for (CompletableFuture<Void> task : tasks.values()) {
             task.cancel(true);
         }
     }
+
+    void waitForAllTasks() throws ExecutionException, InterruptedException {
+        for (CompletableFuture<Void> task : tasks.values()) {
+            if (!task.isDone()) {
+                task.get();
+            }
+        }
+    }
 }

src/main/java/software/amazon/smithy/lsp/SmithyLanguageServer.java

@@ -24,10 +24,12 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.Properties;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.logging.Logger;
 import java.util.stream.Collectors;
@@ -253,16 +255,22 @@ public CompletableFuture<InitializeResult> initialize(InitializeParams params) {
     private void tryInitProject(Path root) {
         LOGGER.info("Initializing project at " + root);
         lifecycleManager.cancelAllTasks();
-        Result<Project, List<Exception>> loadResult = ProjectLoader.load(root);
+        Result<Project, List<Exception>> loadResult = ProjectLoader.load(
+                root, projects, lifecycleManager.getManagedDocuments());
         if (loadResult.isOk()) {
+            Project updatedProject = loadResult.unwrap();
+            resolveDetachedProjects(updatedProject);
             projects.updateMainProject(loadResult.unwrap());
             LOGGER.info("Initialized project at " + root);
-            // TODO: If this is a project reload, there are open files which need to have updated diagnostics reported.
         } else {
             LOGGER.severe("Init project failed");
             // TODO: Maybe we just start with this anyways by default, and then add to it
             //  if we find a smithy-build.json, etc.
-            projects.updateMainProject(Project.empty(root));
+            // If we overwrite an existing project with an empty one, we lose track of the state of tracked
+            // files. Instead, we will just keep the original project before the reload failure.
+            if (projects.getMainProject() == null) {
+                projects.updateMainProject(Project.empty(root));
+            }
 
             String baseMessage = "Failed to load Smithy project at " + root;
             StringBuilder errorMessage = new StringBuilder(baseMessage).append(":");
@@ -278,6 +286,37 @@ private void tryInitProject(Path root) {
         }
     }
 
+    private void resolveDetachedProjects(Project updatedProject) {
+        // This is a project reload, so we need to resolve any added/removed files
+        // that need to be moved to or from detached projects.
+        if (getProject() != null) {
+            Set<String> currentProjectSmithyPaths = getProject().getSmithyFiles().keySet();
+            Set<String> updatedProjectSmithyPaths = updatedProject.getSmithyFiles().keySet();
+
+            Set<String> addedPaths = new HashSet<>(updatedProjectSmithyPaths);
+            addedPaths.removeAll(currentProjectSmithyPaths);
+            for (String addedPath : addedPaths) {
+                String addedUri = UriAdapter.toUri(addedPath);
+                if (projects.isDetached(addedUri)) {
+                    projects.removeDetachedProject(addedUri);
+                }
+            }
+
+            Set<String> removedPaths = new HashSet<>(currentProjectSmithyPaths);
+            removedPaths.removeAll(updatedProjectSmithyPaths);
+            for (String removedPath : removedPaths) {
+                String removedUri = UriAdapter.toUri(removedPath);
+                // Only move to a detached project if the file is managed
+                if (lifecycleManager.getManagedDocuments().contains(removedUri)) {
+                    // Note: This should always be non-null, since we essentially got this from the current project
+                    Document removedDocument = projects.getDocument(removedUri);
+                    // The copy here is technically unnecessary, if we make ModelAssembler support borrowed strings
+                    projects.createDetachedProject(removedUri, removedDocument.copyText());
+                }
+            }
+        }
+    }
+
     private CompletableFuture<Void> registerSmithyFileWatchers() {
         Project project = projects.getMainProject();
         List<Registration> registrations = FileWatcherRegistrationHandler.getSmithyFileWatcherRegistrations(project);
@@ -383,8 +422,8 @@ public void didChangeWatchedFiles(DidChangeWatchedFilesParams params) {
         LOGGER.info("DidChangeWatchedFiles");
         // Smithy files were added or deleted to watched sources/imports (specified by smithy-build.json),
         // or the smithy-build.json itself was changed
-        List<String> createdSmithyFiles = new ArrayList<>();
-        List<String> deletedSmithyFiles = new ArrayList<>();
+        Set<String> createdSmithyFiles = new HashSet<>(params.getChanges().size());
+        Set<String> deletedSmithyFiles = new HashSet<>(params.getChanges().size());
         boolean changedBuildFiles = false;
         for (FileEvent event : params.getChanges()) {
             String changedUri = event.getUri();
@@ -407,20 +446,22 @@ public void didChangeWatchedFiles(DidChangeWatchedFilesParams params) {
             }
         }
 
-        // TODO: Handle files being moved into projects from detached. Will need
-        //  to be able to load project with files managed by the client.
         if (changedBuildFiles) {
             client.info("Build files changed, reloading project");
             // TODO: Handle more granular updates to build files.
             tryInitProject(projects.getMainProject().getRoot());
         } else {
             client.info("Project files changed, adding files "
                         + createdSmithyFiles + " and removing files " + deletedSmithyFiles);
+            // We get this notification for watched files, which only includes project files,
+            // so we don't need to resolve detached projects.
             projects.getMainProject().updateFiles(createdSmithyFiles, deletedSmithyFiles);
         }
 
         // TODO: Update watchers based on specific changes
         unregisterSmithyFileWatchers().thenRun(this::registerSmithyFileWatchers);
+
+        sendFileDiagnosticsForManagedDocuments();
     }
 
     @Override
@@ -440,8 +481,7 @@ public void didChange(DidChangeTextDocumentParams params) {
 
         lifecycleManager.cancelTask(uri);
 
-        Project project = projects.getProject(uri);
-        Document document = project.getDocument(uri);
+        Document document = projects.getDocument(uri);
         if (document == null) {
             client.error("Attempted to change document the server isn't tracking: " + uri);
             return;
@@ -459,7 +499,15 @@ public void didChange(DidChangeTextDocumentParams params) {
             // TODO: A consequence of this is that any existing validation events are cleared, which
             //  is kinda annoying.
             // Report any parse/shape/trait loading errors
-            triggerUpdate(uri);
+            Project project = projects.getProject(uri);
+            if (project == null) {
+                client.error("Attempted to update a file the server isn't tracking: " + uri);
+                return;
+            }
+            CompletableFuture<Void> future = CompletableFuture
+                    .runAsync(() -> project.updateModelWithoutValidating(uri))
+                    .thenComposeAsync(unused -> sendFileDiagnostics(uri));
+            lifecycleManager.putTask(uri, future);
         }
     }
 
@@ -470,24 +518,25 @@ public void didOpen(DidOpenTextDocumentParams params) {
         String uri = params.getTextDocument().getUri();
 
         lifecycleManager.cancelTask(uri);
+        lifecycleManager.getManagedDocuments().add(uri);
 
         String text = params.getTextDocument().getText();
-        Project project = projects.getProject(uri);
-        Document document = project.getDocument(uri);
+        Document document = projects.getDocument(uri);
         if (document != null) {
             document.applyEdit(null, text);
         } else {
             projects.createDetachedProject(uri, text);
         }
-        // TODO: Do we need to handle canceling this?
-        sendFileDiagnostics(uri);
+
+        lifecycleManager.putTask(uri, sendFileDiagnostics(uri));
     }
 
     @Override
     public void didClose(DidCloseTextDocumentParams params) {
         LOGGER.info("DidClose");
 
         String uri = params.getTextDocument().getUri();
+        lifecycleManager.getManagedDocuments().remove(uri);
 
         if (projects.isDetached(uri)) {
             // Only cancel tasks for detached projects, since we're dropping the project
@@ -517,7 +566,11 @@ public void didSave(DidSaveTextDocumentParams params) {
             document.applyEdit(null, params.getText());
         }
 
-        triggerUpdateAndValidate(uri);
+        Project project = projects.getProject(uri);
+        CompletableFuture<Void> future = CompletableFuture
+                .runAsync(() -> project.updateAndValidateModel(uri))
+                .thenCompose(unused -> sendFileDiagnostics(uri));
+        lifecycleManager.putTask(uri, future);
     }
 
     @Override
@@ -639,20 +692,10 @@ public CompletableFuture<List<? extends TextEdit>> formatting(DocumentFormatting
         }
     }
 
-    private void triggerUpdate(String uri) {
-        Project project = projects.getProject(uri);
-        CompletableFuture<Void> future = CompletableFuture
-                .runAsync(() -> project.updateModelWithoutValidating(uri))
-                .thenComposeAsync(unused -> sendFileDiagnostics(uri));
-        lifecycleManager.putTask(uri, future);
-    }
-
-    private void triggerUpdateAndValidate(String uri) {
-        Project project = projects.getProject(uri);
-        CompletableFuture<Void> future = CompletableFuture
-                .runAsync(() -> project.updateAndValidateModel(uri))
-                .thenCompose(unused -> sendFileDiagnostics(uri));
-        lifecycleManager.putTask(uri, future);
+    private void sendFileDiagnosticsForManagedDocuments() {
+        for (String managedDocumentUri : lifecycleManager.getManagedDocuments()) {
+            lifecycleManager.putOrComposeTask(managedDocumentUri, sendFileDiagnostics(managedDocumentUri));
+        }
     }
 
     private CompletableFuture<Void> sendFileDiagnostics(String uri) {
@@ -671,6 +714,11 @@ List<Diagnostic> getFileDiagnostics(String uri) {
         }
 
         Project project = projects.getProject(uri);
+        if (project == null) {
+            client.error("Attempted to get file diagnostics for an untracked file: " + uri);
+            return Collections.emptyList();
+        }
+
         SmithyFile smithyFile = project.getSmithyFile(uri);
         String path = UriAdapter.toPath(uri);