Refactor vscode router to load async.

Author: GirlBossRush

src/node/routes/index.ts

@@ -10,7 +10,7 @@ import { HttpCode, HttpError } from "../../common/http"
 import { plural } from "../../common/util"
 import { App } from "../app"
 import { AuthType, DefaultedArgs } from "../cli"
-import { commit, isDevMode, rootPath } from "../constants"
+import { commit, rootPath } from "../constants"
 import { Heart } from "../heart"
 import { ensureAuthenticated, redirect } from "../http"
 import { PluginAPI } from "../plugin"
@@ -23,7 +23,7 @@ import * as login from "./login"
 import * as logout from "./logout"
 import * as pathProxy from "./pathProxy"
 import * as update from "./update"
-import { createVSServerRouter, VSServerResult } from "./vscode"
+import { CodeServerRouteWrapper } from "./vscode"
 
 /**
  * Register all routes and middleware.
@@ -138,20 +138,12 @@ export const register = async (app: App, args: DefaultedArgs): Promise<Disposabl
 
   app.router.use("/update", update.router)
 
-  let vscode: VSServerResult
-  try {
-    vscode = await createVSServerRouter(args)
-    app.router.use("/", vscode.router)
-    app.wsRouter.use("/", vscode.wsRouter.router)
-    app.router.use("/vscode", vscode.router)
-    app.wsRouter.use("/vscode", vscode.wsRouter.router)
-  } catch (error: any) {
-    if (isDevMode) {
-      logger.warn(error)
-      logger.warn("VS Server router may still be compiling.")
-    } else {
-      throw error
-    }
+  const vsServerRouteHandler = new CodeServerRouteWrapper()
+
+  // Note that the root route is replaced in Coder Enterprise by the plugin API.
+  for (const routePrefix of ["/", "/vscode"]) {
+    app.router.use(routePrefix, vsServerRouteHandler.router)
+    app.wsRouter.use(routePrefix, vsServerRouteHandler.wsRouter)
   }
 
   app.router.use(() => {
@@ -164,6 +156,6 @@ export const register = async (app: App, args: DefaultedArgs): Promise<Disposabl
   return () => {
     heart.dispose()
     pluginApi?.dispose()
-    vscode?.codeServerMain.dispose()
+    vsServerRouteHandler.dispose()
   }
 }

src/node/routes/vscode.ts

@@ -1,65 +1,105 @@
+import { logger } from "@coder/logger"
 import * as express from "express"
-import { DefaultedArgs } from "../cli"
+import { WebsocketRequest } from "../../../typings/pluginapi"
+import { logError } from "../../common/util"
+import { isDevMode } from "../constants"
 import { ensureAuthenticated, authenticated, redirect } from "../http"
 import { loadAMDModule } from "../util"
-import { Router as WsRouter, WebsocketRouter } from "../wsRouter"
+import { Router as WsRouter } from "../wsRouter"
 import { errorHandler } from "./errors"
 
-export interface VSServerResult {
-  router: express.Router
-  wsRouter: WebsocketRouter
-  codeServerMain: CodeServerLib.IServerAPI
-}
-
-export const createVSServerRouter = async (args: DefaultedArgs): Promise<VSServerResult> => {
-  // See ../../../vendor/modules/code-oss-dev/src/vs/server/main.js.
-  const createVSServer = await loadAMDModule<CodeServerLib.CreateServer>(
-    "vs/server/remoteExtensionHostAgent",
-    "createServer",
-  )
+export class CodeServerRouteWrapper {
+  /** Assigned in `ensureCodeServerLoaded` */
+  private _codeServerMain!: CodeServerLib.IServerAPI
+  private _wsRouterWrapper = WsRouter()
+  public router = express.Router()
 
-  const codeServerMain = await createVSServer(null, {
-    connectionToken: "0000",
-    ...args,
-    // For some reason VS Code takes the port as a string.
-    port: typeof args.port !== "undefined" ? args.port.toString() : undefined,
-  })
+  public get wsRouter() {
+    return this._wsRouterWrapper.router
+  }
 
-  const router = express.Router()
-  const wsRouter = WsRouter()
+  //#region Route Handlers
 
-  router.get("/", async (req, res, next) => {
+  private $root: express.Handler = async (req, res, next) => {
     const isAuthenticated = await authenticated(req)
+
     if (!isAuthenticated) {
       return redirect(req, res, "login", {
         // req.baseUrl can be blank if already at the root.
         to: req.baseUrl && req.baseUrl !== "/" ? req.baseUrl : undefined,
       })
     }
+
     next()
-  })
+  }
 
-  router.all("*", ensureAuthenticated, (req, res, next) => {
-    req.on("error", (error: any) => {
+  private $proxyRequest: express.Handler = async (req, res, next) => {
+    // We allow certain errors to propagate so that other routers may handle requests
+    // outside VS Code
+    const requestErrorHandler = (error: any) => {
       if (error instanceof Error && ["EntryNotFound", "FileNotFound", "HttpError"].includes(error.message)) {
         next()
       }
 
       errorHandler(error, req, res, next)
-    })
+    }
+
+    req.once("error", requestErrorHandler)
 
-    codeServerMain.handleRequest(req, res)
-  })
+    this._codeServerMain.handleRequest(req, res)
+  }
 
-  wsRouter.ws("/", ensureAuthenticated, (req) => {
-    codeServerMain.handleUpgrade(req, req.socket)
+  private $proxyWebsocket = async (req: WebsocketRequest) => {
+    this._codeServerMain.handleUpgrade(req, req.socket)
 
     req.socket.resume()
-  })
+  }
+
+  //#endregion
+
+  /**
+   * Fetches a code server instance asynchronously to avoid an initial memory overhead.
+   */
+  private ensureCodeServerLoaded: express.Handler = async (req, _res, next) => {
+    if (this._codeServerMain) {
+      return next()
+    }
+
+    const { args } = req
+
+    /**
+     * @file ../../../vendor/modules/code-oss-dev/src/vs/server/main.js
+     */
+    const createVSServer = await loadAMDModule<CodeServerLib.CreateServer>(
+      "vs/server/remoteExtensionHostAgent",
+      "createServer",
+    )
+
+    try {
+      this._codeServerMain = await createVSServer(null, {
+        connectionToken: "0000",
+        ...args,
+        // For some reason VS Code takes the port as a string.
+        port: args.port?.toString(),
+      })
+    } catch (createServerError) {
+      logError(logger, "CodeServerRouteWrapper", createServerError)
+
+      const loggedError = isDevMode ? new Error("VS Code may still be compiling...") : createServerError
+
+      return next(loggedError)
+    }
+
+    return next()
+  }
+
+  constructor() {
+    this.router.get("/", this.ensureCodeServerLoaded, this.$root)
+    this.router.all("*", ensureAuthenticated, this.ensureCodeServerLoaded, this.$proxyRequest)
+    this._wsRouterWrapper.ws("/", ensureAuthenticated, this.ensureCodeServerLoaded, this.$proxyWebsocket)
+  }
 
-  return {
-    router,
-    wsRouter,
-    codeServerMain,
+  dispose() {
+    this._codeServerMain?.dispose()
   }
 }