From 75c8fdeed2b1dc0c69993d2e95a11df9db349827 Mon Sep 17 00:00:00 2001
From: Jacob Goldman <jacobmgoldman@icloud.com>
Date: Mon, 31 Aug 2020 11:29:12 -0400
Subject: [PATCH] Added /healthz JSON response for heartbeat data. #1940
 (#1984)

---
 src/node/app/health.ts | 32 ++++++++++++++++++++++++++++++++
 src/node/entry.ts      |  4 +++-
 src/node/http.ts       | 17 +++++++++++------
 3 files changed, 46 insertions(+), 7 deletions(-)
 create mode 100644 src/node/app/health.ts

diff --git a/src/node/app/health.ts b/src/node/app/health.ts
new file mode 100644
index 00000000..6a3aae94
--- /dev/null
+++ b/src/node/app/health.ts
@@ -0,0 +1,32 @@
+import * as http from "http"
+import { HttpCode, HttpError } from "../../common/http"
+import { HttpProvider, HttpResponse, Route, Heart, HttpProviderOptions } from "../http"
+
+/**
+ * Check the heartbeat.
+ */
+export class HealthHttpProvider extends HttpProvider {
+  public constructor(options: HttpProviderOptions, private readonly heart: Heart) {
+    super(options)
+  }
+
+  public async handleRequest(route: Route, request: http.IncomingMessage): Promise<HttpResponse> {
+    if (!this.authenticated(request)) {
+      if (this.isRoot(route)) {
+        return { redirect: "/login", query: { to: route.fullPath } }
+      }
+      throw new HttpError("Unauthorized", HttpCode.Unauthorized)
+    }
+
+    const result = {
+      cache: false,
+      mime: "application/json",
+      content: {
+        status: this.heart.alive() ? "alive" : "expired",
+        lastHeartbeat: this.heart.lastHeartbeat,
+      },
+    }
+
+    return result
+  }
+}
diff --git a/src/node/entry.ts b/src/node/entry.ts
index 61f417f5..739476e1 100644
--- a/src/node/entry.ts
+++ b/src/node/entry.ts
@@ -3,8 +3,9 @@ import * as cp from "child_process"
 import { promises as fs } from "fs"
 import http from "http"
 import * as path from "path"
-import { CliMessage, OpenCommandPipeArgs } from "../../lib/vscode/src/vs/server/ipc"
+import { CliMessage } from "../../lib/vscode/src/vs/server/ipc"
 import { plural } from "../common/util"
+import { HealthHttpProvider } from "./app/health"
 import { LoginHttpProvider } from "./app/login"
 import { ProxyHttpProvider } from "./app/proxy"
 import { StaticHttpProvider } from "./app/static"
@@ -80,6 +81,7 @@ const main = async (args: Args, cliArgs: Args, configArgs: Args): Promise<void>
   httpServer.registerHttpProvider("/proxy", ProxyHttpProvider)
   httpServer.registerHttpProvider("/login", LoginHttpProvider, args.config!, envPassword)
   httpServer.registerHttpProvider("/static", StaticHttpProvider)
+  httpServer.registerHttpProvider("/healthz", HealthHttpProvider, httpServer.heart)
 
   await loadPlugins(httpServer, args)
 
diff --git a/src/node/http.ts b/src/node/http.ts
index 5c8346f7..37bbcfbd 100644
--- a/src/node/http.ts
+++ b/src/node/http.ts
@@ -396,23 +396,26 @@ export abstract class HttpProvider {
 export class Heart {
   private heartbeatTimer?: NodeJS.Timeout
   private heartbeatInterval = 60000
-  private lastHeartbeat = 0
+  public lastHeartbeat = 0
 
   public constructor(private readonly heartbeatPath: string, private readonly isActive: () => Promise<boolean>) {}
 
+  public alive(): boolean {
+    const now = Date.now()
+    return now - this.lastHeartbeat < this.heartbeatInterval
+  }
   /**
    * Write to the heartbeat file if we haven't already done so within the
    * timeout and start or reset a timer that keeps running as long as there is
    * activity. Failures are logged as warnings.
    */
   public beat(): void {
-    const now = Date.now()
-    if (now - this.lastHeartbeat >= this.heartbeatInterval) {
+    if (!this.alive()) {
       logger.trace("heartbeat")
       fs.outputFile(this.heartbeatPath, "").catch((error) => {
         logger.warn(error.message)
       })
-      this.lastHeartbeat = now
+      this.lastHeartbeat = Date.now()
       if (typeof this.heartbeatTimer !== "undefined") {
         clearTimeout(this.heartbeatTimer)
       }
@@ -457,7 +460,7 @@ export class HttpServer {
   private listenPromise: Promise<string | null> | undefined
   public readonly protocol: "http" | "https"
   private readonly providers = new Map<string, HttpProvider>()
-  private readonly heart: Heart
+  public readonly heart: Heart
   private readonly socketProvider = new SocketProxyProvider()
 
   /**
@@ -602,8 +605,10 @@ export class HttpServer {
   }
 
   private onRequest = async (request: http.IncomingMessage, response: http.ServerResponse): Promise<void> => {
-    this.heart.beat()
     const route = this.parseUrl(request)
+    if (route.providerBase !== "/healthz") {
+      this.heart.beat()
+    }
     const write = (payload: HttpResponse): void => {
       response.writeHead(payload.redirect ? HttpCode.Redirect : payload.code || HttpCode.Ok, {
         "Content-Type": payload.mime || getMediaMime(payload.filePath),