Add idle timeout (#7539)

Author: andrewbaldwin44

docs/FAQ.md

@@ -322,12 +322,8 @@ As long as there is an active browser connection, code-server touches
 `~/.local/share/code-server/heartbeat` once a minute.
 
 If you want to shutdown code-server if there hasn't been an active connection
-after a predetermined amount of time, you can do so by checking continuously for
-the last modified time on the heartbeat file. If it is older than X minutes (or
-whatever amount of time you'd like), you can kill code-server.
-
-Eventually, [#1636](https://github.com/coder/code-server/issues/1636) will make
-this process better.
+after a predetermined amount of time, you can use the --idle-timeout-seconds flag
+or set an `CODE_SERVER_IDLE_TIMEOUT_SECONDS` environment variable.
 
 ## How do I change the password?
 

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,
   }
 }

test/unit/node/heart.test.ts

@@ -1,6 +1,6 @@
 import { logger } from "@coder/logger"
 import { readFile, writeFile, stat, utimes } from "fs/promises"
-import { Heart, heartbeatTimer } from "../../../src/node/heart"
+import { Heart } from "../../../src/node/heart"
 import { clean, mockLogger, tmpdir } from "../../utils/helpers"
 
 const mockIsActive = (resolveTo: boolean) => jest.fn().mockResolvedValue(resolveTo)
@@ -82,31 +82,81 @@ describe("Heart", () => {
 })
 
 describe("heartbeatTimer", () => {
-  beforeAll(() => {
+  const testName = "heartbeatTimer"
+  let testDir = ""
+  beforeAll(async () => {
+    await clean(testName)
+    testDir = await tmpdir(testName)
     mockLogger()
   })
   afterAll(() => {
     jest.restoreAllMocks()
   })
+  beforeEach(() => {
+    jest.useFakeTimers()
+  })
   afterEach(() => {
     jest.resetAllMocks()
+    jest.clearAllTimers()
+    jest.useRealTimers()
   })
-  it("should call beat when isActive resolves to true", async () => {
+  it("should call isActive when timeout expires", async () => {
     const isActive = true
     const mockIsActive = jest.fn().mockResolvedValue(isActive)
-    const mockBeatFn = jest.fn()
-    await heartbeatTimer(mockIsActive, mockBeatFn)
+    const heart = new Heart(`${testDir}/shutdown.txt`, mockIsActive)
+    await heart.beat()
+    jest.advanceTimersByTime(60 * 1000)
     expect(mockIsActive).toHaveBeenCalled()
-    expect(mockBeatFn).toHaveBeenCalled()
   })
   it("should log a warning when isActive rejects", async () => {
     const errorMsg = "oh no"
     const error = new Error(errorMsg)
     const mockIsActive = jest.fn().mockRejectedValue(error)
-    const mockBeatFn = jest.fn()
-    await heartbeatTimer(mockIsActive, mockBeatFn)
+    const heart = new Heart(`${testDir}/shutdown.txt`, mockIsActive)
+    await heart.beat()
+    jest.advanceTimersByTime(60 * 1000)
+
     expect(mockIsActive).toHaveBeenCalled()
-    expect(mockBeatFn).not.toHaveBeenCalled()
     expect(logger.warn).toHaveBeenCalledWith(errorMsg)
   })
 })
+
+describe("stateChange", () => {
+  const testName = "stateChange"
+  let testDir = ""
+  let heart: Heart
+  beforeAll(async () => {
+    await clean(testName)
+    testDir = await tmpdir(testName)
+    mockLogger()
+  })
+  afterAll(() => {
+    jest.restoreAllMocks()
+  })
+  afterEach(() => {
+    jest.resetAllMocks()
+    if (heart) {
+      heart.dispose()
+    }
+  })
+  it("should change to alive after a beat", async () => {
+    heart = new Heart(`${testDir}/shutdown.txt`, mockIsActive(true))
+    const mockOnChange = jest.fn()
+    heart.onChange(mockOnChange)
+    await heart.beat()
+
+    expect(mockOnChange.mock.calls[0][0]).toBe("alive")
+  })
+  it.only("should change to expired when not active", async () => {
+    jest.useFakeTimers()
+    heart = new Heart(`${testDir}/shutdown.txt`, () => new Promise((resolve) => resolve(false)))
+    const mockOnChange = jest.fn()
+    heart.onChange(mockOnChange)
+    await heart.beat()
+
+    await jest.advanceTimersByTime(60 * 1000)
+    expect(mockOnChange.mock.calls[1][0]).toBe("expired")
+    jest.clearAllTimers()
+    jest.useRealTimers()
+  })
+})

test/unit/node/main.test.ts

@@ -16,6 +16,7 @@ jest.mock("@coder/logger", () => ({
     debug: jest.fn(),
     warn: jest.fn(),
     error: jest.fn(),
+    named: jest.fn(),
     level: 0,
   },
   field: jest.fn(),
@@ -94,7 +95,7 @@ describe("main", () => {
 
       // Mock routes module
       jest.doMock("../../../src/node/routes", () => ({
-        register: jest.fn().mockResolvedValue(jest.fn()),
+        register: jest.fn().mockResolvedValue({ disposeRoutes: jest.fn() }),
       }))
 
       // Mock loadCustomStrings to succeed
@@ -131,7 +132,7 @@ describe("main", () => {
 
       // Mock routes module
       jest.doMock("../../../src/node/routes", () => ({
-        register: jest.fn().mockResolvedValue(jest.fn()),
+        register: jest.fn().mockResolvedValue({ disposeRoutes: jest.fn() }),
       }))
 
       // Import runCodeServer after mocking