Add idle timeout (#7539)

src/node/cli.ts

@@ -94,6 +94,7 @@ export interface UserProvidedArgs extends UserProvidedCodeArgs {
   "welcome-text"?: string
   "abs-proxy-base-path"?: string
   i18n?: string
+  "idle-timeout-seconds"?: number
   /* Positional arguments. */
   _?: string[]
 }
@@ -303,6 +304,10 @@ export const options: Options<Required<UserProvidedArgs>> = {
     path: true,
     description: "Path to JSON file with custom translations. Merges with default strings and supports all i18n keys.",
   },
+  "idle-timeout-seconds": {
+    type: "number",
+    description: "Timeout in seconds to wait before shutting down when idle.",
+  },
 }
 
 export const optionDescriptions = (opts: Partial<Options<Required<UserProvidedArgs>>> = options): string[] => {
@@ -396,6 +401,10 @@ export const parse = (
         throw new Error("--github-auth can only be set in the config file or passed in via $GITHUB_TOKEN")
       }
 
+      if (key === "idle-timeout-seconds" && Number(value) <= 60) {
+        throw new Error("--idle-timeout-seconds must be greater than 60 seconds.")
+      }
+
       const option = options[key]
       if (option.type === "boolean") {
         ;(args[key] as boolean) = true
@@ -611,6 +620,16 @@ export async function setDefaults(cliArgs: UserProvidedArgs, configArgs?: Config
     args["github-auth"] = process.env.GITHUB_TOKEN
   }
 
+  if (process.env.CODE_SERVER_IDLE_TIMEOUT_SECONDS) {
+    if (isNaN(Number(process.env.CODE_SERVER_IDLE_TIMEOUT_SECONDS))) {
+      logger.info("CODE_SERVER_IDLE_TIMEOUT_SECONDS must be a number")
+    }
+    if (Number(process.env.CODE_SERVER_IDLE_TIMEOUT_SECONDS) <= 60) {
+      throw new Error("--idle-timeout-seconds must be greater than 60 seconds.")
+    }
+    args["idle-timeout-seconds"] = Number(process.env.CODE_SERVER_IDLE_TIMEOUT_SECONDS)
+  }
+
   // Ensure they're not readable by child processes.
   delete process.env.PASSWORD
   delete process.env.HASHED_PASSWORD

src/node/heart.ts

@@ -1,5 +1,6 @@
 import { logger } from "@coder/logger"
 import { promises as fs } from "fs"
+import { Emitter } from "../common/emitter"
 
 /**
  * Provides a heartbeat using a local file to indicate activity.
@@ -8,6 +9,9 @@ export class Heart {
   private heartbeatTimer?: NodeJS.Timeout
   private heartbeatInterval = 60000
   public lastHeartbeat = 0
+  private readonly _onChange = new Emitter<"alive" | "expired" | "unknown">()
+  readonly onChange = this._onChange.event
+  private state: "alive" | "expired" | "unknown" = "expired"
 
   public constructor(
     private readonly heartbeatPath: string,
@@ -17,6 +21,13 @@ export class Heart {
     this.alive = this.alive.bind(this)
   }
 
+  private setState(state: typeof this.state) {
+    if (this.state !== state) {
+      this.state = state
+      this._onChange.emit(this.state)
+    }
+  }
+
   public alive(): boolean {
     const now = Date.now()
     return now - this.lastHeartbeat < this.heartbeatInterval
@@ -28,6 +39,7 @@ export class Heart {
    */
   public async beat(): Promise<void> {
     if (this.alive()) {
+      this.setState("alive")
       return
     }
 
@@ -36,7 +48,22 @@ export class Heart {
     if (typeof this.heartbeatTimer !== "undefined") {
       clearTimeout(this.heartbeatTimer)
     }
-    this.heartbeatTimer = setTimeout(() => heartbeatTimer(this.isActive, this.beat), this.heartbeatInterval)
+
+    this.heartbeatTimer = setTimeout(async () => {
+      try {
+        if (await this.isActive()) {
+          this.beat()
+        } else {
+          this.setState("expired")
+        }
+      } catch (error: unknown) {
+        logger.warn((error as Error).message)
+        this.setState("unknown")
+      }
+    }, this.heartbeatInterval)
+
+    this.setState("alive")
+
     try {
       return await fs.writeFile(this.heartbeatPath, "")
     } catch (error: any) {
@@ -53,20 +80,3 @@ export class Heart {
     }
   }
 }
-
-/**
- * Helper function for the heartbeatTimer.
- *
- * If heartbeat is active, call beat. Otherwise do nothing.
- *
- * Extracted to make it easier to test.
- */
-export async function heartbeatTimer(isActive: Heart["isActive"], beat: Heart["beat"]) {
-  try {
-    if (await isActive()) {
-      beat()
-    }
-  } catch (error: unknown) {
-    logger.warn((error as Error).message)
-  }
-}

src/node/main.ts

@@ -11,6 +11,7 @@ import { loadCustomStrings } from "./i18n"
 import { register } from "./routes"
 import { VSCodeModule } from "./routes/vscode"
 import { isDirectory, open } from "./util"
+import { wrapper } from "./wrapper"
 
 /**
  * Return true if the user passed an extension-related VS Code flag.
@@ -141,7 +142,7 @@ export const runCodeServer = async (
   const app = await createApp(args)
   const protocol = args.cert ? "https" : "http"
   const serverAddress = ensureAddress(app.server, protocol)
-  const disposeRoutes = await register(app, args)
+  const { disposeRoutes, heart } = await register(app, args)
 
   logger.info(`Using config file ${args.config}`)
   logger.info(`${protocol.toUpperCase()} server listening on ${serverAddress.toString()}`)
@@ -166,6 +167,27 @@ export const runCodeServer = async (
     logger.info("  - Not serving HTTPS")
   }
 
+  if (args["idle-timeout-seconds"]) {
+    logger.info(`  - Idle timeout set to ${args["idle-timeout-seconds"]} seconds`)
+
+    let idleShutdownTimer: NodeJS.Timeout | undefined
+    const startIdleShutdownTimer = () => {
+      idleShutdownTimer = setTimeout(() => {
+        logger.warn(`Idle timeout of ${args["idle-timeout-seconds"]} seconds exceeded`)
+        wrapper.exit(0)
+      }, args["idle-timeout-seconds"]! * 1000)
+    }
+
+    startIdleShutdownTimer()
+
+    heart.onChange((state) => {
+      clearTimeout(idleShutdownTimer)
+      if (state === "expired") {
+        startIdleShutdownTimer()
+      }
+    })
+  }
+
   if (args["disable-proxy"]) {
     logger.info("  - Proxy disabled")
   } else if (args["proxy-domain"].length > 0) {

src/node/routes/index.ts

@@ -28,7 +28,10 @@ import * as vscode from "./vscode"
 /**
  * Register all routes and middleware.
  */
-export const register = async (app: App, args: DefaultedArgs): Promise<Disposable["dispose"]> => {
+export const register = async (
+  app: App,
+  args: DefaultedArgs,
+): Promise<{ disposeRoutes: Disposable["dispose"]; heart: Heart }> => {
   const heart = new Heart(path.join(paths.data, "heartbeat"), async () => {
     return new Promise((resolve, reject) => {
       // getConnections appears to not call the callback when there are no more
@@ -173,8 +176,11 @@ export const register = async (app: App, args: DefaultedArgs): Promise<Disposabl
   app.router.use(errorHandler)
   app.wsRouter.use(wsErrorHandler)
 
-  return () => {
-    heart.dispose()
-    vscode.dispose()
+  return {
+    disposeRoutes: () => {
+      heart.dispose()
+      vscode.dispose()
+    },
+    heart,
   }
 }