Refactor to use event Emitter

src/node/heart.ts

@@ -1,26 +1,30 @@
 import { logger } from "@coder/logger"
 import { promises as fs } from "fs"
-import { wrapper } from "./wrapper"
+import { Emitter } from "../common/emitter"
 
 /**
  * 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
+  private readonly _onChange = new Emitter<"alive" | "idle" | "unknown">()
+  readonly onChange = this._onChange.event
+  private state: "alive" | "idle" | "unknown" = "idle"
 
   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)
+  private setState(state: typeof this.state) {
+    if (this.state !== state) {
+      this.state = state
+      this._onChange.emit(this.state)
     }
   }
 
@@ -35,6 +39,7 @@ export class Heart {
    */
   public async beat(): Promise<void> {
     if (this.alive()) {
+      this.setState("alive")
       return
     }
 
@@ -43,13 +48,22 @@ 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)
-    }
+
+    this.heartbeatTimer = setTimeout(async () => {
+      try {
+        if (await this.isActive()) {
+          this.beat()
+        } else {
+          this.setState("idle")
+        }
+      } 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) {
@@ -65,26 +79,4 @@ export class Heart {
       clearTimeout(this.heartbeatTimer)
     }
   }
-
-  private exitIfIdle(): void {
-    logger.warn(`Idle timeout of ${this.idleTimeoutSeconds} seconds exceeded`)
-    wrapper.exit(0)
-  }
-}
-
-/**
- * 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()}`)
@@ -168,6 +169,23 @@ export const runCodeServer = async (
 
   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 === "idle") {
+        startIdleShutdownTimer()
+      }
+    })
   }
 
   if (args["disable-proxy"]) {

src/node/routes/index.ts

@@ -28,8 +28,11 @@ 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"), args["idle-timeout-seconds"], async () => {
+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
       // connections.  Feels like it must be a bug?  For now add a timer to make
@@ -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,
   }
 }