fix(playgrounds): handle out of memory playground worker VSCODE-269 (#459)

Author: Anemy

playgrounds/out-of-memory-test.mongodb

@@ -0,0 +1,22 @@
+/**
+ * This playground is used to test how the language server worker
+ * handles an out of memory error in mongosh's running of a playground. VSCODE-269
+ */
+
+use('test');
+
+const mockDataArray = [];
+for(let i = 0; i < 50000; i++) {
+  mockDataArray.push(Math.random() * 10000);
+}
+
+const docs = [];
+for(let i = 0; i < 10000000; i++) {
+  docs.push({
+    mockData: [...mockDataArray],
+    a: 'test 123',
+    b: Math.ceil(Math.random() * 10000)
+  });
+}
+
+console.log('Should not show this message as the process should have run out of memory in the loop above.');

src/editors/playgroundController.ts

@@ -332,21 +332,36 @@ export default class PlaygroundController {
 
     this._statusView.showMessage('Getting results...');
 
-    // Send a request to the language server to execute scripts from a playground.
-    const result: ShellExecuteAllResult =
-      await this._languageServerController.executeAll({
-        codeToEvaluate,
-        connectionId,
-      });
+    try {
+      // Send a request to the language server to execute scripts from a playground.
+      const result: ShellExecuteAllResult =
+        await this._languageServerController.executeAll({
+          codeToEvaluate,
+          connectionId,
+        });
 
-    this._statusView.hideMessage();
-    this._telemetryService.trackPlaygroundCodeExecuted(
-      result,
-      this._isPartialRun,
-      result ? false : true
-    );
+      this._statusView.hideMessage();
+      this._telemetryService.trackPlaygroundCodeExecuted(
+        result,
+        this._isPartialRun,
+        result ? false : true
+      );
 
-    return result;
+      return result;
+    } catch (err: any) {
+      // We re-initialize the language server when we encounter an error.
+      // This happens when the language server worker runs out of memory, can't be revitalized, and restarts.
+      if (err?.code === -32097) {
+        void vscode.window.showErrorMessage(
+          'An error occurred when running the playground. This can occur when the playground runner runs out of memory.'
+        );
+
+        await this._languageServerController.startLanguageServer();
+        void this._connectToServiceProvider();
+      }
+
+      throw err;
+    }
   }
 
   _getAllText(): string {

src/language/languageServerController.ts

@@ -89,10 +89,12 @@ export default class LanguageServerController {
 
   async startLanguageServer(): Promise<void> {
     // Push the disposable client to the context's subscriptions so that the
-    // client can be deactivated on extension deactivation
-    this._context.subscriptions.push(this._client);
+    // client can be deactivated on extension deactivation.
+    if (!this._context.subscriptions.includes(this._client)) {
+      this._context.subscriptions.push(this._client);
+    }
 
-    // Subscribe on notifications from the server when the client is ready
+    // Subscribe on notifications from the server when the client is ready.
     await this._client.sendRequest(
       ServerCommands.SET_EXTENSION_PATH,
       this._context.extensionPath

src/test/suite/editors/playgroundController.test.ts

@@ -406,6 +406,48 @@ suite('Playground Controller Test Suite', function () {
         });
       });
 
+      test('it shows an error message and restarts, and connects the language server when an error occurs in executeAll (out of memory can cause this)', async () => {
+        const mockConnectionDisposedError = new Error(
+          'Pending response rejected since connection got disposed'
+        );
+        (mockConnectionDisposedError as any).code = -32097;
+        sinon
+          .stub(mockLanguageServerController, 'executeAll')
+          .rejects(mockConnectionDisposedError);
+
+        const stubStartLanguageServer = sinon
+          .stub(mockLanguageServerController, 'startLanguageServer')
+          .resolves();
+
+        const stubConnectToServiceProvider = sinon
+          .stub(testPlaygroundController, '_connectToServiceProvider')
+          .resolves();
+
+        const stubVSCodeErrorMessage = sinon
+          .stub(vscode.window, 'showErrorMessage')
+          .resolves(undefined);
+
+        try {
+          await testPlaygroundController._evaluate('console.log("test");');
+
+          // It should have thrown in the above evaluation.
+          expect(true).to.equal(false);
+        } catch (err: any) {
+          expect(err.message).to.equal(
+            'Pending response rejected since connection got disposed'
+          );
+          expect(err.code).to.equal(-32097);
+        }
+
+        expect(stubVSCodeErrorMessage.calledOnce).to.equal(true);
+        expect(stubVSCodeErrorMessage.firstCall.args[0]).to.equal(
+          'An error occurred when running the playground. This can occur when the playground runner runs out of memory.'
+        );
+
+        expect(stubStartLanguageServer.calledOnce).to.equal(true);
+        expect(stubConnectToServiceProvider.calledOnce).to.equal(true);
+      });
+
       test('playground controller loads the active editor on start', () => {
         sandbox.replaceGetter(
           vscode.window,
@@ -458,12 +500,9 @@ suite('Playground Controller Test Suite', function () {
           document: { getText: () => textFromEditor },
         } as vscode.TextEditor;
 
-        const fakeVscodeErrorMessage: any = sinon.fake();
-        sinon.replace(
-          vscode.window,
-          'showErrorMessage',
-          fakeVscodeErrorMessage
-        );
+        const fakeVscodeErrorMessage = sinon
+          .stub(vscode.window, 'showErrorMessage')
+          .resolves(undefined);
 
         playgroundControllerTest._selectedText = '{ name: qwerty }';
         playgroundControllerTest._codeActionProvider.selection = selection;
@@ -474,7 +513,9 @@ suite('Playground Controller Test Suite', function () {
 
         const expectedMessage =
           "Unable to export to csharp language: Symbol 'qwerty' is undefined";
-        expect(fakeVscodeErrorMessage.firstArg).to.be.equal(expectedMessage);
+        expect(fakeVscodeErrorMessage.firstCall.args[0]).to.equal(
+          expectedMessage
+        );
       });
     });
   });

src/test/suite/explorer/playgroundsExplorer.test.ts

@@ -54,21 +54,13 @@ suite('Playgrounds Controller Test Suite', function () {
     try {
       const children = await treeController.getPlaygrounds(rootUri);
 
-      assert.strictEqual(
-        Object.keys(children).length,
-        5,
-        `Tree playgrounds should have 5 child, found ${children.length}`
-      );
+      assert.strictEqual(Object.keys(children).length, 6);
 
       const playgrounds = Object.values(children).filter(
         (item: any) => item.label && item.label.split('.').pop() === 'mongodb'
       );
 
-      assert.strictEqual(
-        Object.keys(playgrounds).length,
-        5,
-        `Tree playgrounds should have 5 playgrounds with mongodb extension, found ${children.length}`
-      );
+      assert.strictEqual(Object.keys(playgrounds).length, 6);
     } catch (error) {
       assert(false, error as Error);
     }
@@ -92,10 +84,7 @@ suite('Playgrounds Controller Test Suite', function () {
       const treeController = new PlaygroundsTree();
       const children = await treeController.getPlaygrounds(rootUri);
 
-      assert(
-        Object.keys(children).length === 3,
-        `Tree playgrounds should have 3 child, found ${children.length}`
-      );
+      assert.strictEqual(Object.keys(children).length, 3);
     } catch (error) {
       assert(false, error as Error);
     }