Add idle timeout

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 `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.IDLE_TIMEOUT_SECONDS) {
+    if (isNaN(Number(process.env.IDLE_TIMEOUT_SECONDS))) {
+      logger.info("IDLE_TIMEOUT_SECONDS must be a number")
+    }
+    if (Number(process.env.IDLE_TIMEOUT_SECONDS)) {
+      throw new Error("--idle-timeout-seconds must be greater than 60 seconds.")
+    }
+    args["idle-timeout-seconds"] = Number(process.env.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,20 +1,27 @@
 import { logger } from "@coder/logger"
 import { promises as fs } from "fs"
+import { wrapper } from "./wrapper"
 
 /**
  * Provides a heartbeat using a local file to indicate activity.
  */
 export class Heart {
   private heartbeatTimer?: NodeJS.Timeout
+  private idleShutdownTimer?: NodeJS.Timeout
   private heartbeatInterval = 60000
   public lastHeartbeat = 0
 
   public constructor(
     private readonly heartbeatPath: string,
+    private idleTimeoutSeconds: number | undefined,
     private readonly isActive: () => Promise<boolean>,
   ) {
     this.beat = this.beat.bind(this)
     this.alive = this.alive.bind(this)
+
+    if (this.idleTimeoutSeconds) {
+      this.idleShutdownTimer = setTimeout(() => this.exitIfIdle(), this.idleTimeoutSeconds * 1000)
+    }
   }
 
   public alive(): boolean {
@@ -36,7 +43,13 @@ export class Heart {
     if (typeof this.heartbeatTimer !== "undefined") {
       clearTimeout(this.heartbeatTimer)
     }
+    if (typeof this.idleShutdownTimer !== "undefined") {
+      clearInterval(this.idleShutdownTimer)
+    }
     this.heartbeatTimer = setTimeout(() => heartbeatTimer(this.isActive, this.beat), this.heartbeatInterval)
+    if (this.idleTimeoutSeconds) {
+      this.idleShutdownTimer = setTimeout(() => this.exitIfIdle(), this.idleTimeoutSeconds * 1000)
+    }
     try {
       return await fs.writeFile(this.heartbeatPath, "")
     } catch (error: any) {
@@ -52,6 +65,11 @@ export class Heart {
       clearTimeout(this.heartbeatTimer)
     }
   }
+
+  private exitIfIdle(): void {
+    logger.warn(`Idle timeout of ${this.idleTimeoutSeconds} seconds exceeded`)
+    wrapper.exit(0)
+  }
 }
 
 /**

src/node/main.ts

@@ -166,6 +166,10 @@ 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`)
+  }
+
   if (args["disable-proxy"]) {
     logger.info("  - Proxy disabled")
   } else if (args["proxy-domain"].length > 0) {

src/node/routes/index.ts

@@ -29,7 +29,7 @@ import * as vscode from "./vscode"
  * Register all routes and middleware.
  */
 export const register = async (app: App, args: DefaultedArgs): Promise<Disposable["dispose"]> => {
-  const heart = new Heart(path.join(paths.data, "heartbeat"), async () => {
+  const heart = new Heart(path.join(paths.data, "heartbeat"), args["idle-timeout-seconds"], async () => {
     return new Promise((resolve, reject) => {
       // getConnections appears to not call the callback when there are no more
       // connections.  Feels like it must be a bug?  For now add a timer to make

test/unit/node/heart.test.ts

@@ -1,10 +1,21 @@
 import { logger } from "@coder/logger"
 import { readFile, writeFile, stat, utimes } from "fs/promises"
 import { Heart, heartbeatTimer } from "../../../src/node/heart"
+import { wrapper } from "../../../src/node/wrapper"
 import { clean, mockLogger, tmpdir } from "../../utils/helpers"
 
 const mockIsActive = (resolveTo: boolean) => jest.fn().mockResolvedValue(resolveTo)
 
+jest.mock("../../../src/node/wrapper", () => {
+  const original = jest.requireActual("../../../src/node/wrapper")
+  return {
+    ...original,
+    wrapper: {
+      exit: jest.fn(),
+    },
+  }
+})
+
 describe("Heart", () => {
   const testName = "heartTests"
   let testDir = ""
@@ -16,7 +27,7 @@ describe("Heart", () => {
     testDir = await tmpdir(testName)
   })
   beforeEach(() => {
-    heart = new Heart(`${testDir}/shutdown.txt`, mockIsActive(true))
+    heart = new Heart(`${testDir}/shutdown.txt`, undefined, mockIsActive(true))
   })
   afterAll(() => {
     jest.restoreAllMocks()
@@ -42,7 +53,7 @@ describe("Heart", () => {
 
     expect(fileContents).toBe(text)
 
-    heart = new Heart(pathToFile, mockIsActive(true))
+    heart = new Heart(pathToFile, undefined, mockIsActive(true))
     await heart.beat()
     // Check that the heart wrote to the heartbeatFilePath and overwrote our text
     const fileContentsAfterBeat = await readFile(pathToFile, { encoding: "utf8" })
@@ -52,7 +63,7 @@ describe("Heart", () => {
     expect(fileStatusAfterEdit.mtimeMs).toBeGreaterThan(0)
   })
   it("should log a warning when given an invalid file path", async () => {
-    heart = new Heart(`fakeDir/fake.txt`, mockIsActive(false))
+    heart = new Heart(`fakeDir/fake.txt`, undefined, mockIsActive(false))
     await heart.beat()
     expect(logger.warn).toHaveBeenCalled()
   })
@@ -71,7 +82,7 @@ describe("Heart", () => {
   it("should beat twice without warnings", async () => {
     // Use fake timers so we can speed up setTimeout
     jest.useFakeTimers()
-    heart = new Heart(`${testDir}/hello.txt`, mockIsActive(true))
+    heart = new Heart(`${testDir}/hello.txt`, undefined, mockIsActive(true))
     await heart.beat()
     // we need to speed up clocks, timeouts
     // call heartbeat again (and it won't be alive I think)
@@ -110,3 +121,32 @@ describe("heartbeatTimer", () => {
     expect(logger.warn).toHaveBeenCalledWith(errorMsg)
   })
 })
+
+describe("idleTimeout", () => {
+  const testName = "idleHeartTests"
+  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 call beat when isActive resolves to true", async () => {
+    jest.useFakeTimers()
+    heart = new Heart(`${testDir}/shutdown.txt`, 60, mockIsActive(true))
+
+    jest.advanceTimersByTime(60 * 1000)
+    expect(wrapper.exit).toHaveBeenCalled()
+    jest.clearAllTimers()
+    jest.useRealTimers()
+  })
+})