From 88e971c609d112693d3c81851e31a2c44e1d27e2 Mon Sep 17 00:00:00 2001
From: Joe Previte <jjprevite@gmail.com>
Date: Wed, 4 May 2022 16:05:48 -0700
Subject: [PATCH] refactor(heart): bind class methods and make beat async
 (#5142)

* feat: set up new test for beat twice

* refactor: make Heart.beat() async

This allows us to properly await heart.beat() in our tests and remove
the HACK I added before.

* refactor: bind heart methods .beat and .alive

This allows the functions to maintain access to the Heart instance (or
`this`) even when they are passed to other functions. We do this because
we pass both `isActive` and `beat` to `heartbeatTimer`.

* feat(heart): add test to ensure no warnings called

* fixup!: revert setTimeout for heartbeatTimer

* fixup!: return promise in beat
---
 src/node/heart.ts            | 15 ++++++++++-----
 src/node/routes/index.ts     |  2 ++
 test/unit/node/heart.test.ts | 28 ++++++++++++++++------------
 3 files changed, 28 insertions(+), 17 deletions(-)

diff --git a/src/node/heart.ts b/src/node/heart.ts
index a03802035..19f9aa4ad 100644
--- a/src/node/heart.ts
+++ b/src/node/heart.ts
@@ -9,7 +9,10 @@ export class Heart {
   private heartbeatInterval = 60000
   public lastHeartbeat = 0
 
-  public constructor(private readonly heartbeatPath: string, private readonly isActive: () => Promise<boolean>) {}
+  public constructor(private readonly heartbeatPath: string, private readonly isActive: () => Promise<boolean>) {
+    this.beat = this.beat.bind(this)
+    this.alive = this.alive.bind(this)
+  }
 
   public alive(): boolean {
     const now = Date.now()
@@ -20,20 +23,22 @@ export class Heart {
    * timeout and start or reset a timer that keeps running as long as there is
    * activity. Failures are logged as warnings.
    */
-  public beat(): void {
+  public async beat(): Promise<void> {
     if (this.alive()) {
       return
     }
 
     logger.trace("heartbeat")
-    fs.writeFile(this.heartbeatPath, "").catch((error) => {
-      logger.warn(error.message)
-    })
     this.lastHeartbeat = Date.now()
     if (typeof this.heartbeatTimer !== "undefined") {
       clearTimeout(this.heartbeatTimer)
     }
     this.heartbeatTimer = setTimeout(() => heartbeatTimer(this.isActive, this.beat), this.heartbeatInterval)
+    try {
+      return await fs.writeFile(this.heartbeatPath, "")
+    } catch (error: any) {
+      logger.warn(error.message)
+    }
   }
 
   /**
diff --git a/src/node/routes/index.ts b/src/node/routes/index.ts
index 71ce557c5..13d53df86 100644
--- a/src/node/routes/index.ts
+++ b/src/node/routes/index.ts
@@ -56,6 +56,8 @@ export const register = async (app: App, args: DefaultedArgs): Promise<Disposabl
     // /healthz|/healthz/ needs to be excluded otherwise health checks will make
     // it look like code-server is always in use.
     if (!/^\/healthz\/?$/.test(req.url)) {
+      // NOTE@jsjoeio - intentionally not awaiting the .beat() call here because
+      // we don't want to slow down the request.
       heart.beat()
     }
 
diff --git a/test/unit/node/heart.test.ts b/test/unit/node/heart.test.ts
index 5360c865e..7aa6f08dc 100644
--- a/test/unit/node/heart.test.ts
+++ b/test/unit/node/heart.test.ts
@@ -23,6 +23,7 @@ describe("Heart", () => {
   })
   afterEach(() => {
     jest.resetAllMocks()
+    jest.useRealTimers()
     if (heart) {
       heart.dispose()
     }
@@ -42,11 +43,7 @@ describe("Heart", () => {
     expect(fileContents).toBe(text)
 
     heart = new Heart(pathToFile, mockIsActive(true))
-    heart.beat()
-    // HACK@jsjoeio - beat has some async logic but is not an async method
-    // Therefore, we have to create an artificial wait in order to make sure
-    // all async code has completed before asserting
-    await new Promise((r) => setTimeout(r, 100))
+    await heart.beat()
     // Check that the heart wrote to the heartbeatFilePath and overwrote our text
     const fileContentsAfterBeat = await readFile(pathToFile, { encoding: "utf8" })
     expect(fileContentsAfterBeat).not.toBe(text)
@@ -56,15 +53,11 @@ describe("Heart", () => {
   })
   it("should log a warning when given an invalid file path", async () => {
     heart = new Heart(`fakeDir/fake.txt`, mockIsActive(false))
-    heart.beat()
-    // HACK@jsjoeio - beat has some async logic but is not an async method
-    // Therefore, we have to create an artificial wait in order to make sure
-    // all async code has completed before asserting
-    await new Promise((r) => setTimeout(r, 100))
+    await heart.beat()
     expect(logger.warn).toHaveBeenCalled()
   })
-  it("should be active after calling beat", () => {
-    heart.beat()
+  it("should be active after calling beat", async () => {
+    await heart.beat()
 
     const isAlive = heart.alive()
     expect(isAlive).toBe(true)
@@ -75,6 +68,17 @@ describe("Heart", () => {
     const isAlive = heart.alive()
     expect(isAlive).toBe(false)
   })
+  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))
+    await heart.beat()
+    // we need to speed up clocks, timeouts
+    // call heartbeat again (and it won't be alive I think)
+    // then assert no warnings were called
+    jest.runAllTimers()
+    expect(logger.warn).not.toHaveBeenCalled()
+  })
 })
 
 describe("heartbeatTimer", () => {