diff --git a/packages/lsp/src/virtual/document.ts b/packages/lsp/src/virtual/document.ts index cf42f021ab49..4acc1ae782c1 100644 --- a/packages/lsp/src/virtual/document.ts +++ b/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(); + for (const line of this.sourceLines.values()) { + for (const block of line.foreignDocumentsMap.values()) { + usedDocuments.add(block.virtualDocument); + } + } + + const documentIDs = new Map(); + 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(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; protected lineBlocks: Array; - protected unusedDocuments: Set; protected unusedStandaloneDocuments: DefaultMap< language, Array diff --git a/packages/lsp/test/document.spec.ts b/packages/lsp/test/document.spec.ts index f5a3daa4447f..a7d4ab7b25a2 100644 --- a/packages/lsp/test/document.spec.ts +++ b/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'](