Merge pull request #847 from krassowski/diagnostic-and-highlight-performance

krassowski · web-flow · commit 1a9922643ecd · 2022-09-03T12:43:55.000+01:00
Improve performance of diagnostics and highlights
diff --git a/packages/jupyterlab-lsp/src/editor_integration/codemirror.ts b/packages/jupyterlab-lsp/src/editor_integration/codemirror.ts
@@ -245,6 +245,9 @@ export abstract class CodeMirrorIntegration
     }
   }
 
+  /**
+   * @deprecated
+   */
   protected highlight_range(
     range: IEditorRange,
     class_name: string
diff --git a/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.ts b/packages/jupyterlab-lsp/src/features/diagnostics/diagnostics.ts
@@ -306,6 +306,18 @@ export const diagnostics_databases = new WeakMap<
   DiagnosticsDatabase
 >();
 
+interface IMarkerDefinition {
+  options: CodeMirror.TextMarkerOptions;
+  start: IEditorPosition;
+  end: IEditorPosition;
+  hash: string;
+}
+
+interface IMarkedDiagnostic {
+  editor: CodeMirror.Editor;
+  marker: CodeMirror.TextMarker;
+}
+
 export class DiagnosticsCM extends CodeMirrorIntegration {
   private last_response: lsProtocol.PublishDiagnosticsParams;
 
@@ -339,7 +351,7 @@ export class DiagnosticsCM extends CodeMirrorIntegration {
   }
 
   private unique_editor_ids: DefaultMap<CodeMirror.Editor, number>;
-  private marked_diagnostics: Map<string, CodeMirror.TextMarker> = new Map();
+  private marked_diagnostics: Map<string, IMarkedDiagnostic> = new Map();
   /**
    * Allows access to the most recent diagnostics in context of the editor.
    *
@@ -477,6 +489,11 @@ export class DiagnosticsCM extends CodeMirrorIntegration {
       this.filterDiagnostics(response.diagnostics)
     );
 
+    const markerOptionsByEditor = new Map<
+      CodeMirror.Editor,
+      IMarkerDefinition[]
+    >();
+
     diagnostics_by_range.forEach(
       (diagnostics: lsProtocol.Diagnostic[], range: lsProtocol.Range) => {
         const start = PositionConverter.lsp_to_cm(
@@ -641,23 +658,70 @@ export class DiagnosticsCM extends CodeMirrorIntegration {
               .join('\n'),
             className: classNames.join(' ')
           };
+
+          let optionsList = markerOptionsByEditor.get(cm_editor);
+          if (!optionsList) {
+            optionsList = [];
+            markerOptionsByEditor.set(cm_editor, optionsList);
+          }
+          optionsList.push({
+            options,
+            start: start_in_editor,
+            end: end_in_editor,
+            hash: diagnostic_hash
+          });
+        }
+      }
+    );
+
+    for (const [
+      cmEditor,
+      markerDefinitions
+    ] of markerOptionsByEditor.entries()) {
+      // note: could possibly be wrapped in `requestAnimationFrame()`
+      // at a risk of sometimes having an inconsistent state in database.
+      // note: using `operation()` significantly improves performance.
+      // test cases:
+      //   - one cell with 1000 `math.pi` and `import math`; comment out import,
+      //     wait for 1000 diagnostics, then uncomment import, wait for removal:
+      //     - before:
+      //        - diagnostics show up in: 13.6s (blocking!)
+      //        - diagnostics removal in: 13.2s (blocking!)
+      //     - after:
+      //        - diagnostics show up in: 254ms
+      //        - diagnostics removal in: 160.4ms
+      //   - first open of a notebook with 10 cells, each with 88 diagnostics:
+      //     - before: 2.75s (each time)
+      //     - after 208.75ms (order of magnitude faster!)
+      //   - first open of a notebook with 100 cells, each with 1 diagnostic
+      //       this scenario is expected to have no gain (measures overhead)
+      //     - before 280.34ms, 377ms, 399ms
+      //     - after 385.29ms, 301.97ms, 309.4ms
+      cmEditor.operation(() => {
+        const doc = cmEditor.getDoc();
+        for (const definition of markerDefinitions) {
           let marker;
           try {
-            marker = cm_editor
-              .getDoc()
-              .markText(start_in_editor, end_in_editor, options);
+            marker = doc.markText(
+              definition.start,
+              definition.end,
+              definition.options
+            );
           } catch (e) {
             this.console.warn(
               'Marking inspection (diagnostic text) failed:',
-              diagnostics,
+              definition,
               e
             );
             return;
           }
-          this.marked_diagnostics.set(diagnostic_hash, marker);
+          this.marked_diagnostics.set(definition.hash, {
+            marker,
+            editor: cmEditor
+          });
         }
-      }
-    );
+      });
+    }
 
     // remove the markers which were not included in the new message
     this.removeUnusedDiagnosticMarkers(markers_to_retain);
@@ -696,14 +760,36 @@ export class DiagnosticsCM extends CodeMirrorIntegration {
   }
 
   protected removeUnusedDiagnosticMarkers(to_retain: Set<string>) {
-    this.marked_diagnostics.forEach(
-      (marker: CodeMirror.TextMarker, diagnostic_hash: string) => {
-        if (!to_retain.has(diagnostic_hash)) {
-          this.marked_diagnostics.delete(diagnostic_hash);
-          marker.clear();
+    const toRemoveByEditor = new Map<
+      CodeMirror.Editor,
+      { marker: CodeMirror.TextMarker; hash: string }[]
+    >();
+
+    for (const [
+      diagnosticHash,
+      markedDiagnostic
+    ] of this.marked_diagnostics.entries()) {
+      if (!to_retain.has(diagnosticHash)) {
+        let diagnosticsList = toRemoveByEditor.get(markedDiagnostic.editor);
+        if (!diagnosticsList) {
+          diagnosticsList = [];
+          toRemoveByEditor.set(markedDiagnostic.editor, diagnosticsList);
         }
+        diagnosticsList.push({
+          marker: markedDiagnostic.marker,
+          hash: diagnosticHash
+        });
       }
-    );
+    }
+
+    for (const [cmEditor, markers] of toRemoveByEditor.entries()) {
+      cmEditor.operation(() => {
+        for (const markerData of markers) {
+          markerData.marker.clear();
+          this.marked_diagnostics.delete(markerData.hash);
+        }
+      });
+    }
   }
 
   remove(): void {
diff --git a/packages/jupyterlab-lsp/src/features/highlights.ts b/packages/jupyterlab-lsp/src/features/highlights.ts
@@ -14,7 +14,11 @@ import { CodeHighlights as LSPHighlightsSettings } from '../_highlights';
 import { CodeMirrorIntegration } from '../editor_integration/codemirror';
 import { FeatureSettings } from '../feature';
 import { DocumentHighlightKind } from '../lsp';
-import { IRootPosition, IVirtualPosition } from '../positioning';
+import {
+  IEditorPosition,
+  IRootPosition,
+  IVirtualPosition
+} from '../positioning';
 import { ILSPFeatureManager, PLUGIN_ID } from '../tokens';
 import { VirtualDocument } from '../virtual/document';
 
@@ -23,8 +27,19 @@ export const highlightIcon = new LabIcon({
   svgstr: highlightSvg
 });
 
+interface IHighlightDefinition {
+  options: CodeMirror.TextMarkerOptions;
+  start: IEditorPosition;
+  end: IEditorPosition;
+}
+
 export class HighlightsCM extends CodeMirrorIntegration {
-  protected highlight_markers: CodeMirror.TextMarker[] = [];
+  protected highlightMarkers: Map<CodeMirror.Editor, CodeMirror.TextMarker[]> =
+    new Map();
+  /*
+   * @deprecated
+   */
+  protected highlight_markers: CodeMirror.TextMarker[];
   private debounced_get_highlight: Debouncer<
     lsProtocol.DocumentHighlight[] | undefined
   >;
@@ -63,9 +78,14 @@ export class HighlightsCM extends CodeMirrorIntegration {
   }
 
   protected clear_markers() {
-    for (let marker of this.highlight_markers) {
-      marker.clear();
+    for (const [cmEditor, markers] of this.highlightMarkers.entries()) {
+      cmEditor.operation(() => {
+        for (const marker of markers) {
+          marker.clear();
+        }
+      });
     }
+    this.highlightMarkers = new Map();
     this.highlight_markers = [];
   }
 
@@ -78,16 +98,75 @@ export class HighlightsCM extends CodeMirrorIntegration {
       return;
     }
 
+    const highlightOptionsByEditor = new Map<
+      CodeMirror.Editor,
+      IHighlightDefinition[]
+    >();
+
     for (let item of items) {
       let range = this.range_to_editor_range(item.range);
       let kind_class = item.kind
         ? 'cm-lsp-highlight-' + DocumentHighlightKind[item.kind]
         : '';
-      let marker = this.highlight_range(
-        range,
-        'cm-lsp-highlight ' + kind_class
-      );
-      this.highlight_markers.push(marker);
+
+      let optionsList = highlightOptionsByEditor.get(range.editor);
+
+      if (!optionsList) {
+        optionsList = [];
+        highlightOptionsByEditor.set(range.editor, optionsList);
+      }
+      optionsList.push({
+        options: {
+          className: 'cm-lsp-highlight ' + kind_class
+        },
+        start: range.start,
+        end: range.end
+      });
+    }
+
+    for (const [
+      cmEditor,
+      markerDefinitions
+    ] of highlightOptionsByEditor.entries()) {
+      // note: using `operation()` significantly improves performance.
+      // test cases:
+      //   - one cell with 1000 `math.pi` and `import math`; move cursor to `math`,
+      //     wait for 1000 highlights, then move to `pi`:
+      //     - before:
+      //        - highlight `math`: 13.1s
+      //        - then highlight `pi`: 16.6s
+      //     - after:
+      //        - highlight `math`: 160ms
+      //        - then highlight `pi`: 227ms
+      //   - 100 cells with `math.pi` and one with `import math`; move cursor to `math`,
+      //     wait for 1000 highlights, then move to `pi` (this is overhead control,
+      //     no gains expected):
+      //     - before:
+      //        - highlight `math`: 385ms
+      //        - then highlight `pi`: 683 ms
+      //     - after:
+      //        - highlight `math`: 390ms
+      //        - then highlight `pi`: 870ms
+      cmEditor.operation(() => {
+        const doc = cmEditor.getDoc();
+        const markersList: CodeMirror.TextMarker[] = [];
+        for (const definition of markerDefinitions) {
+          let marker;
+          try {
+            marker = doc.markText(
+              definition.start,
+              definition.end,
+              definition.options
+            );
+          } catch (e) {
+            this.console.warn('Marking highlight failed:', definition, e);
+            return;
+          }
+          markersList.push(marker);
+          this.highlight_markers.push(marker);
+        }
+        this.highlightMarkers.set(cmEditor, markersList);
+      });
     }
   };
 

Original file line number	Diff line number	Diff line change
`@@ -245,6 +245,9 @@ export abstract class CodeMirrorIntegration`
`245`	`245`	`}`
`246`	`246`	`}`
`247`	`247`
	`248`	`+ /**`
	`249`	`+ * @deprecated`
	`250`	`+ */`
`248`	`251`	`protected highlight_range(`
`249`	`252`	`range: IEditorRange,`
`250`	`253`	`class_name: string`