fix(crashMonitoring): handle sleep/wake appropriately

Problem:

When a user does a sleep then wake, our heartbeats are not up to date since
they cannot send when the user's computer is asleep.

Due to this there is a race condition between the heartbeats sent and when we check the
heartbeats to determine if they are stale (crash).

Solution:

Keep a TimeLag class that helps to assert that there is no lag in our time.
It works by updating a state every second. And if we determine that the next
update to that state took longer than a second, we determine that there was a lag.
Then we simply skip the next crash check, allowing a fresh heartbeat to be sent.

Signed-off-by: nkomonen-amazon <[email protected]>

Author: nkomonen-amazon

packages/core/src/shared/crashMonitoring.ts

@@ -61,6 +61,7 @@ const className = 'CrashMonitoring'
  */
 export class CrashMonitoring {
     protected heartbeat: Heartbeat | undefined
+    protected timeLag: TimeLag | undefined
     protected crashChecker: CrashChecker | undefined
 
     constructor(
@@ -111,16 +112,27 @@ export class CrashMonitoring {
         }
 
         try {
+            this.timeLag = new TimeLag()
+
             this.heartbeat = new Heartbeat(this.state, this.checkInterval, this.isDevMode)
-            this.crashChecker = new CrashChecker(this.state, this.checkInterval, this.isDevMode, this.devLogger)
+            this.heartbeat.onFailure(() => this.cleanup())
+
+            this.crashChecker = new CrashChecker(
+                this.state,
+                this.checkInterval,
+                this.isDevMode,
+                this.devLogger,
+                this.timeLag
+            )
+            this.crashChecker.onFailure(() => this.cleanup())
 
+            this.timeLag.start()
             await this.heartbeat.start()
             await this.crashChecker.start()
         } catch (error) {
             emitFailure({ functionName: 'start', error })
             try {
-                this.crashChecker?.cleanup()
-                await this.heartbeat?.cleanup()
+                await this.cleanup()
             } catch {}
 
             // Surface errors during development, otherwise it can be missed.
@@ -135,6 +147,7 @@ export class CrashMonitoring {
         try {
             this.crashChecker?.cleanup()
             await this.heartbeat?.shutdown()
+            this.timeLag?.cleanup()
         } catch (error) {
             try {
                 // This probably wont emit in time before shutdown, but may be written to the logs
@@ -146,6 +159,12 @@ export class CrashMonitoring {
             }
         }
     }
+
+    public async cleanup() {
+        this.crashChecker?.cleanup()
+        await this.heartbeat?.cleanup()
+        this.timeLag?.cleanup()
+    }
 }
 
 /**
@@ -154,15 +173,19 @@ export class CrashMonitoring {
  */
 class Heartbeat {
     private intervalRef: NodeJS.Timer | undefined
+    private _onFailure = new vscode.EventEmitter<void>()
+    public onFailure: vscode.Event<void> = this._onFailure.event
+    private readonly heartbeatInterval: number
+
     constructor(
         private readonly state: FileSystemState,
-        private readonly checkInterval: number,
+        checkInterval: number,
         private readonly isDevMode: boolean
-    ) {}
+    ) {
+        this.heartbeatInterval = checkInterval / 2
+    }
 
     public async start() {
-        const heartbeatInterval = this.checkInterval / 2
-
         // Send an initial heartbeat immediately
         await withFailCtx('initialSendHeartbeat', () => this.state.sendHeartbeat())
 
@@ -179,14 +202,15 @@ class Heartbeat {
                 if (this.isDevMode) {
                     throw e
                 }
+                this._onFailure.fire()
             }
-        }, heartbeatInterval)
+        }, this.heartbeatInterval)
     }
 
     /** Stops everything, signifying a graceful shutdown */
     public async shutdown() {
         globals.clock.clearInterval(this.intervalRef)
-        return this.state.indicateGracefulShutdown()
+        await this.state.indicateGracefulShutdown()
     }
 
     /**
@@ -217,12 +241,15 @@ class Heartbeat {
  */
 class CrashChecker {
     private intervalRef: NodeJS.Timer | undefined
+    private _onFailure = new vscode.EventEmitter<void>()
+    public onFailure = this._onFailure.event
 
     constructor(
         private readonly state: FileSystemState,
         private readonly checkInterval: number,
         private readonly isDevMode: boolean,
-        private readonly devLogger: Logger | undefined
+        private readonly devLogger: Logger | undefined,
+        private readonly timeLag: TimeLag
     ) {}
 
     public async start() {
@@ -231,20 +258,21 @@ class CrashChecker {
 
             // do an initial check
             await withFailCtx('initialCrashCheck', () =>
-                tryCheckCrash(this.state, this.checkInterval, this.isDevMode, this.devLogger)
+                tryCheckCrash(this.state, this.checkInterval, this.isDevMode, this.devLogger, this.timeLag)
             )
 
             // check on an interval
             this.intervalRef = globals.clock.setInterval(async () => {
                 try {
-                    await tryCheckCrash(this.state, this.checkInterval, this.isDevMode, this.devLogger)
+                    await tryCheckCrash(this.state, this.checkInterval, this.isDevMode, this.devLogger, this.timeLag)
                 } catch (e) {
                     emitFailure({ functionName: 'checkCrashInterval', error: e })
-                    this.cleanup()
 
                     if (this.isDevMode) {
                         throw e
                     }
+
+                    this._onFailure.fire()
                 }
             }, this.checkInterval)
         }
@@ -255,8 +283,15 @@ class CrashChecker {
             state: FileSystemState,
             checkInterval: number,
             isDevMode: boolean,
-            devLogger: Logger | undefined
+            devLogger: Logger | undefined,
+            timeLag: TimeLag
         ) {
+            if (await timeLag.didLag()) {
+                timeLag.reset()
+                devLogger?.warn('crashMonitoring: SKIPPED check crash due to time lag')
+                return
+            }
+
             // Iterate all known extensions and for each check if they have crashed
             const knownExts = await state.getAllExts()
             const runningExts: ExtInstanceHeartbeat[] = []
@@ -324,6 +359,84 @@ class CrashChecker {
 
     /** Mimics a crash, only for testing */
     public testCrash() {
+        this.cleanup()
+    }
+}
+
+/**
+ * This class is required for the following edge case:
+ * 1. Heartbeat is sent
+ * 2. Computer goes to sleep for X minutes
+ * 3. Wake up computer. But before a new heartbeat can be sent, a crash checker (can be from another ext instance) runs
+ *   and sees a stale heartbeat. It assumes a crash.
+ *
+ * Why? Intervals do not run while the computer is asleep, so the latest heartbeat has a "lag" since it wasn't able to send
+ *      a new heartbeat.
+ *      Then on wake, there is a racecondition for the next heartbeat to be sent before the next crash check. If the crash checker
+ *      runs first it will incorrectly conclude a crash.
+ *
+ * Solution: Keep track of the lag, and then skip the next crash check if there was a lag. This will give time for the
+ *           next heartbeat to be sent.
+ */
+class TimeLag {
+    private intervalRef: NodeJS.Timer | undefined
+    private isCompleted: Promise<void> | undefined
+    /** Resolves {@link isCompleted} when the next interval of lag checking is completed */
+    private setCompleted: (() => void) | undefined
+
+    /** The last timestamp we remember. If this is more than 1 {@link lagCheckInterval} we probably did a sleep+wake */
+    private latestTimestamp: number = 0
+    /** The accumulation of lag before the next crash checker interval, since the user can sleep+wake multiple times in between */
+    private totalLag: number = 0
+    private readonly lagCheckInterval = 1000
+
+    public start() {
+        this.reset() // initialize
+
+        // Every interval calculate the lag
+        this.intervalRef = globals.clock.setInterval(() => {
+            const expectedNow = this.latestTimestamp + this.lagCheckInterval
+            const actualNow = globals.clock.Date.now()
+            const lag = actualNow - expectedNow
+
+            // interval callback execution is not exact, so this buffer avoids micro lag from being
+            // considered actual lag
+            if (lag > 5000) {
+                this.totalLag += lag
+            }
+
+            this.latestTimestamp = Date.now()
+
+            // race condition between crash checker and lag checker on computer wake. So we have
+            // the crash checker block until this is completed.
+            this.setCompleted!()
+            this.setInProgress()
+        }, this.lagCheckInterval)
+    }
+
+    /** True if there is a time lag */
+    public async didLag() {
+        await this.isCompleted!
+        return this.totalLag > 0
+    }
+
+    /** Indicates the next time lag interval check is not completed */
+    private setInProgress() {
+        this.isCompleted = new Promise((resolve) => {
+            this.setCompleted = resolve
+        })
+    }
+
+    /**
+     * Call this once the user of this instance has handled the lag. If not done, {@link didLag} can
+     * be permanently true.
+     */
+    public reset() {
+        this.totalLag = 0
+        this.latestTimestamp = globals.clock.Date.now()
+        this.setInProgress()
+    }
+    public cleanup() {
         globals.clock.clearInterval(this.intervalRef)
     }
 }
@@ -617,7 +730,10 @@ export type ExtInstance = {
     isDebug?: boolean
 }
 
-type ExtInstanceHeartbeat = ExtInstance & { lastHeartbeat: number }
+type ExtInstanceHeartbeat = ExtInstance & {
+    /** Timestamp of the last heartbeat in milliseconds */
+    lastHeartbeat: number
+}
 
 function isExtHeartbeat(ext: unknown): ext is ExtInstanceHeartbeat {
     return typeof ext === 'object' && ext !== null && 'lastHeartbeat' in ext && ext.lastHeartbeat !== undefined