Backport PR jupyterlab#15105: Remove unusedDocuments, fix culling o…

…f foreign documents

packages/lsp/src/virtual/document.ts

@@ -241,7 +241,6 @@ export class VirtualDocument implements IDisposable {
     );
     this._remainingLifetime = 6;
 
-    this.unusedDocuments = new Set();
     this.documentInfo = new VirtualDocumentInfo(this);
     this.updateManager = new UpdateManager(this);
     this.updateManager.updateBegan.connect(this._updateBeganSlot, this);
@@ -446,7 +445,6 @@ export class VirtualDocument implements IDisposable {
 
     this.foreignDocuments.clear();
     this.sourceLines.clear();
-    this.unusedDocuments.clear();
     this.unusedStandaloneDocuments.clear();
     this.virtualLines.clear();
 
@@ -470,7 +468,6 @@ export class VirtualDocument implements IDisposable {
     // TODO - deep clear (assure that there is no memory leak)
     this.unusedStandaloneDocuments.clear();
 
-    this.unusedDocuments = new Set();
     this.virtualLines.clear();
     this.sourceLines.clear();
     this.lastVirtualLine = 0;
@@ -806,10 +803,34 @@ export class VirtualDocument implements IDisposable {
    * Close all expired documents.
    */
   closeExpiredDocuments(): void {
-    for (let document of this.unusedDocuments.values()) {
+    const usedDocuments = new Set<VirtualDocument>();
+    for (const line of this.sourceLines.values()) {
+      for (const block of line.foreignDocumentsMap.values()) {
+        usedDocuments.add(block.virtualDocument);
+      }
+    }
+
+    const documentIDs = new Map<VirtualDocument, string[]>();
+    for (const [id, document] of this.foreignDocuments.entries()) {
+      const ids = documentIDs.get(document);
+      if (typeof ids !== 'undefined') {
+        documentIDs.set(document, [...ids, id]);
+      }
+      documentIDs.set(document, [id]);
+    }
+    const allDocuments = new Set<VirtualDocument>(documentIDs.keys());
+    const unusedVirtualDocuments = new Set(
+      [...allDocuments].filter(x => !usedDocuments.has(x))
+    );
+
+    for (let document of unusedVirtualDocuments.values()) {
       document.remainingLifetime -= 1;
       if (document.remainingLifetime <= 0) {
         document.dispose();
+        const ids = documentIDs.get(document)!;
+        for (const id of ids) {
+          this.foreignDocuments.delete(id);
+        }
       }
     }
   }
@@ -901,7 +922,7 @@ export class VirtualDocument implements IDisposable {
 
   /**
    * When this counter goes down to 0, the document will be destroyed and the associated connection will be closed;
-   * This is meant to reduce the number of open connections when a a foreign code snippet was removed from the document.
+   * This is meant to reduce the number of open connections when a foreign code snippet was removed from the document.
    *
    * Note: top level virtual documents are currently immortal (unless killed by other means); it might be worth
    * implementing culling of unused documents, but if and only if JupyterLab will also implement culling of
@@ -927,7 +948,6 @@ export class VirtualDocument implements IDisposable {
   protected sourceLines: Map<number, ISourceLine>;
   protected lineBlocks: Array<string>;
 
-  protected unusedDocuments: Set<VirtualDocument>;
   protected unusedStandaloneDocuments: DefaultMap<
     language,
     Array<VirtualDocument>

packages/lsp/test/document.spec.ts

@@ -137,7 +137,6 @@ describe('@jupyterlab/lsp', () => {
         document.dispose();
         expect(document.foreignDocuments.size).toEqual(0);
         expect(document['sourceLines'].size).toEqual(0);
-        expect(document['unusedDocuments'].size).toEqual(0);
         expect(document['unusedStandaloneDocuments'].size).toEqual(0);
         expect(document['virtualLines'].size).toEqual(0);
       });
@@ -239,6 +238,34 @@ describe('@jupyterlab/lsp', () => {
         expect(md.uri).toBe('test.ipynb.python-text.txt');
       });
     });
+    describe('#closeExpiredDocuments', () => {
+      it('should close expired foreign documents', async () => {
+        // We start with a notebook having a code cell, a raw cell and a markdown cell
+        // this means we have two foreign documents which can expire: one for markdown
+        // cell and one for raw cell.
+        expect(new Set(document.foreignDocuments.keys())).toEqual(
+          new Set(['text', 'markdown'])
+        );
+        const newNotebookState = [
+          {
+            value: 'test line in markdown 1\ntest line in markdown 2',
+            ceEditor: {} as Document.IEditor,
+            type: 'markdown'
+          }
+        ];
+        // If user just removed the last raw cell we do expire the virtual document
+        // associated immediately (because if they are just cutting and pasting
+        // cells we may be able to recycle connections). However if user keeps modifying
+        // the notebook, we ultimately discard the virtual document as expired.
+        for (let i = 0; i < 10; i++) {
+          await document.updateManager.updateDocuments(newNotebookState);
+          document.closeExpiredDocuments();
+        }
+        expect(new Set(document.foreignDocuments.keys())).toEqual(
+          new Set(['markdown'])
+        );
+      });
+    });
     describe('#openForeign', () => {
       it('should create a new foreign document', () => {
         const doc: VirtualDocument = document['openForeign'](