refactor: simplify websocket code (#425)

client/src/engine/jobs.ts

@@ -47,8 +47,6 @@ export enum Request {
   lspWorkspaceSymbols = 'lsp/workspaceSymbols',
   lspShowAst = 'lsp/showAst',
 
-  internalRestart = 'ext/restart', // Internal Extension Request
-  internalDownloadLatest = 'ext/downloadLatest', // Internal Extension Request
   internalReady = 'ext/ready', // Internal Extension Request
   internalMessage = 'ext/message', // Internal Extension Request
   internalError = 'ext/error', // Internal Extension Request

client/src/extension.ts

@@ -270,8 +270,6 @@ async function startSession(
 
   client.onNotification(jobs.Request.internalDiagnostics, handlePrintDiagnostics)
 
-  client.onNotification(jobs.Request.internalRestart, makeHandleRestartClient(context))
-
   client.onNotification(jobs.Request.internalMessage, vscode.window.showInformationMessage)
 
   client.onNotification(jobs.Request.internalError, handleError)

server/package.json

@@ -13,9 +13,13 @@
   },
   "dependencies": {
     "portfinder": "^1.0.32",
+    "reconnecting-websocket": "^4.4.0",
     "vscode-languageserver": "^8.0.2",
     "vscode-languageserver-textdocument": "^1.0.8",
     "ws": "^8.17.1"
   },
+  "devDependencies": {
+    "@types/ws": "^8.5.10"
+  },
   "scripts": {}
 }

server/src/engine/jobs.ts

@@ -50,8 +50,6 @@ export enum Request {
   lspSemanticTokens = 'lsp/semanticTokens',
   lspShowAst = 'lsp/showAst',
 
-  internalRestart = 'ext/restart', // Internal Extension Request
-  internalDownloadLatest = 'ext/downloadLatest', // Internal Extension Request
   internalReady = 'ext/ready', // Internal Extension Request
   internalMessage = 'ext/message', // Internal Extension Request
   internalError = 'ext/error', // Internal Extension Request

server/src/engine/socket.ts

@@ -19,13 +19,12 @@ import * as queue from './queue'
 import { clearDiagnostics, sendNotification } from '../server'
 import { EventEmitter } from 'events'
 import { handleCrash, lspCheckResponseHandler } from '../handlers'
-import { getPort } from 'portfinder'
 import { USER_MESSAGE } from '../util/userMessages'
 import { StatusCode } from '../util/statusCodes'
+import ReconnectingWebSocket from 'reconnecting-websocket'
+import WebSocket from 'ws'
 
-const WebSocket = require('ws')
-
-let webSocket: any
+let webSocket: ReconnectingWebSocket
 let webSocketOpen = false
 
 // event emitter to handle communication between socket handlers and connection handlers
@@ -85,60 +84,41 @@ export function isClosed() {
 let lastManualStopTimestamp: number = 0
 
 export function initialiseSocket({ uri, onOpen, onClose }: InitialiseSocketInput) {
-  if (!uri) {
-    throw 'Must be called with an uri'
-  }
-  webSocket = new WebSocket(uri)
+  webSocket = new ReconnectingWebSocket(uri, [], {
+    WebSocket,
+  })
+
+  webSocket.addEventListener('open', function handleOpen() {
+    // The 'open' event is emitted every time the connection is established,
+    // even if it was just a temprorary interruption.
+    // This handler should only be called once, so remove it after the first call.
+    webSocket.removeEventListener('open', handleOpen)
 
-  webSocket.on('open', () => {
     webSocketOpen = true
-    onOpen && setTimeout(onOpen!, 0)
+
+    onOpen?.()
   })
 
-  webSocket.on('close', () => {
+  webSocket.addEventListener('close', () => {
     webSocketOpen = false
+
     if (lastManualStopTimestamp + 15000 < Date.now()) {
-      // This happends when the connections breaks unintentionally
+      // This happens when the connections breaks unintentionally
       console.log(USER_MESSAGE.CONNECTION_LOST())
-      tryToConnect({ uri, onOpen, onClose }, 5).then(connected => {
-        if (!connected) {
-          console.log(USER_MESSAGE.CONNECTION_LOST_RESTARTING())
-          sendNotification(jobs.Request.internalRestart)
-        }
-      })
       return
     }
-    onClose && setTimeout(onClose!, 0)
+
+    onClose?.()
   })
 
-  webSocket.on('message', (data: string) => {
-    const flixResponse: FlixResponse = JSON.parse(data)
-    const job: jobs.EnqueuedJob = jobs.getJob(flixResponse.id)
+  webSocket.addEventListener('message', message => {
+    const flixResponse: FlixResponse = JSON.parse(message.data)
+    const job = jobs.getJob(flixResponse.id)
 
     handleResponse(flixResponse, job)
   })
 }
 
-async function tryToConnect({ uri, onOpen, onClose }: InitialiseSocketInput, times: number) {
-  const uriPort = parseInt(uri.slice(-4))
-  getPort({ port: uriPort }, (err, freePort) => {
-    if (uriPort === freePort) {
-      // This happens if the previously used port is now free
-      sendNotification(jobs.Request.internalRestart)
-      return
-    }
-  })
-  let retries = times
-  while (retries-- > 0) {
-    initialiseSocket({ uri, onOpen, onClose })
-    await sleep(1000)
-    if (webSocketOpen) {
-      return true
-    }
-  }
-  return false
-}
-
 function clearTimer(id: string) {
   clearTimeout(sentMessagesMap[id])
   delete sentMessagesMap[id]
@@ -171,22 +151,7 @@ export async function closeSocket() {
   }
 }
 
-export function sendMessage(job: jobs.EnqueuedJob, expectResponse = true, retries = 0) {
-  if (isClosed()) {
-    if (retries > 2) {
-      const errorMessage = USER_MESSAGE.REQUEST_TIMEOUT(retries)
-      sendNotification(jobs.Request.internalError, {
-        message: errorMessage,
-        actions: [],
-      })
-      return
-    }
-    setTimeout(() => {
-      sendMessage(job, expectResponse, retries + 1)
-    }, 1000)
-    return
-  }
-
+export function sendMessage(job: jobs.EnqueuedJob, expectResponse = true) {
   if (expectResponse) {
     // register a timer to handle timeouts
     sentMessagesMap[job.id] = setTimeout(() => {

server/src/handlers/handlers.ts

@@ -90,7 +90,7 @@ export function handleReplaceConfiguration(userConfiguration: engine.UserConfigu
 }
 
 /**
- * Simulates the language server disconnecting.
+ * Simulates the compiler disconnecting.
  * Used for testing.
  */
 export function handleDisconnect() {

server/src/util/userMessages.ts

@@ -8,10 +8,6 @@ export class USER_MESSAGE {
     return 'Connection to the flix server was lost, trying to reconnect...'
   }
 
-  static CONNECTION_LOST_RESTARTING() {
-    return 'Failed to connect to the flix server, restarting the compiler...'
-  }
-
   static FAILED_TO_START() {
     return 'Failed starting Flix'
   }

server/tsconfig.json

@@ -8,7 +8,9 @@
     "strict": true,
     "outDir": "out",
     "rootDir": "src",
-    "allowJs": true
+    "allowJs": true,
+    "esModuleInterop": true,
+    "skipLibCheck": true
   },
   "include": ["src"],
   "exclude": ["node_modules", ".vscode-test"]

test/src/disconnect.test.ts

@@ -24,16 +24,18 @@ suite('Server disconnect', () => {
   })
 
   test('Should reconnect automatically when server is disconnected', async () => {
-    await vscode.commands.executeCommand('flix.simulateDisconnect')
+    for (let i = 0; i < 10; i++) {
+      await vscode.commands.executeCommand('flix.simulateDisconnect')
 
-    // Wait for the server to disconnect, otherwise the next command will hang
-    await sleep(1000)
+      // Wait for the server to disconnect, otherwise the next command will hang
+      await sleep(1000)
 
-    // Ensure that the server is reconnected
-    const docUri = getTestDocUri('src/Main.flix')
-    const position = new vscode.Position(9, 12)
-    const r = await vscode.commands.executeCommand<vscode.Hover[]>('vscode.executeHoverProvider', docUri, position)
-    const contents = r[0].contents[0] as vscode.MarkdownString
-    assert.strictEqual(contents.value.includes('Type'), true)
+      // Ensure that the server is reconnected
+      const docUri = getTestDocUri('src/Main.flix')
+      const position = new vscode.Position(9, 12)
+      const r = await vscode.commands.executeCommand<vscode.Hover[]>('vscode.executeHoverProvider', docUri, position)
+      const contents = r[0].contents[0] as vscode.MarkdownString
+      assert.strictEqual(contents.value.includes('Type'), true)
+    }
   })
 })