feat: allow disabling metrics server, consolidate exit handlers (#1201)

* feat: allow disabling metrics server, consolidate exit handlers

* Make server off by default, enable in prods, update docs accordingly

---------

Co-authored-by: Nathan Sarrazin <[email protected]>

.env

@@ -159,6 +159,7 @@ ALLOWED_USER_EMAILS=`[]` # if it's defined, only these emails will be allowed to
 USAGE_LIMITS=`{}`
 
 ALLOW_INSECURE_COOKIES=false # recommended to keep this to false but set to true if you need to run over http without tls
-METRICS_PORT=
+METRICS_ENABLED=false
+METRICS_PORT=5565
 LOG_LEVEL=info
 BODY_SIZE_LIMIT=15728640

chart/env/prod.yaml

@@ -35,6 +35,7 @@ envVars:
   EXPOSE_API: "true"
   METRICS_PORT: 5565
   LOG_LEVEL: "debug"
+  METRICS_ENABLED: "true"
   MODELS: >
     [
       {

docs/source/_toctree.yml

@@ -20,6 +20,8 @@
       title: OpenID
     - local: configuration/web-search
       title: Web Search
+    - local: configuration/metrics
+      title: Metrics
     - local: configuration/embeddings
       title: Text Embedding Models
     - title: Models

docs/source/configuration/metrics.md

@@ -0,0 +1,9 @@
+# Metrics
+
+The server can expose prometheus metrics on port `5565` but is off by default. You may enable the metrics server with `METRICS_ENABLED=true` and change the port with `METRICS_PORT=1234`.
+
+<Tip>
+
+In development with `npm run dev`, the metrics server does not shutdown gracefully due to Sveltekit not providing hooks for restart. It's recommended to disable the metrics server in this case.
+
+</Tip>

src/hooks.server.ts

@@ -13,10 +13,13 @@ import { refreshAssistantsCounts } from "$lib/assistantStats/refresh-assistants-
 import { logger } from "$lib/server/logger";
 import { AbortedGenerations } from "$lib/server/abortedGenerations";
 import { MetricsServer } from "$lib/server/metrics";
+import { initExitHandler } from "$lib/server/exitHandler";
 import { ObjectId } from "mongodb";
 
 // TODO: move this code on a started server hook, instead of using a "building" flag
 if (!building) {
+	initExitHandler();
+
 	await checkAndRunMigrations();
 	if (env.ENABLE_ASSISTANTS) {
 		refreshAssistantsCounts();

src/lib/server/abortedGenerations.ts

@@ -2,6 +2,7 @@
 
 import { logger } from "$lib/server/logger";
 import { collections } from "$lib/server/database";
+import { onExit } from "./exitHandler";
 
 export class AbortedGenerations {
 	private static instance: AbortedGenerations;
@@ -10,10 +11,7 @@ export class AbortedGenerations {
 
 	private constructor() {
 		const interval = setInterval(this.updateList, 1000);
-
-		process.on("SIGINT", () => {
-			clearInterval(interval);
-		});
+		onExit(() => clearInterval(interval));
 	}
 
 	public static getInstance(): AbortedGenerations {

src/lib/server/database.ts

@@ -15,6 +15,7 @@ import type { Semaphore } from "$lib/types/Semaphore";
 import type { AssistantStats } from "$lib/types/AssistantStats";
 import { logger } from "$lib/server/logger";
 import { building } from "$app/environment";
+import { onExit } from "./exitHandler";
 
 export const CONVERSATION_STATS_COLLECTION = "conversations.stats";
 
@@ -41,15 +42,8 @@ export class Database {
 		this.client.db(env.MONGODB_DB_NAME + (import.meta.env.MODE === "test" ? "-test" : ""));
 		this.client.on("open", () => this.initDatabase());
 
-		// Disconnect DB on process kill
-		process.on("SIGINT", async () => {
-			await this.client.close(true);
-
-			// https://github.com/sveltejs/kit/issues/9540
-			setTimeout(() => {
-				process.exit(0);
-			}, 100);
-		});
+		// Disconnect DB on exit
+		onExit(() => this.client.close(true));
 	}
 
 	public static getInstance(): Database {

src/lib/server/exitHandler.ts

@@ -0,0 +1,41 @@
+import { randomUUID } from "$lib/utils/randomUuid";
+import { timeout } from "$lib/utils/timeout";
+import { logger } from "./logger";
+
+type ExitHandler = () => void | Promise<void>;
+type ExitHandlerUnsubscribe = () => void;
+
+const listeners = new Map<string, ExitHandler>();
+
+export function onExit(cb: ExitHandler): ExitHandlerUnsubscribe {
+	const uuid = randomUUID();
+	listeners.set(uuid, cb);
+	return () => {
+		listeners.delete(uuid);
+	};
+}
+
+async function runExitHandler(handler: ExitHandler): Promise<void> {
+	return timeout(Promise.resolve().then(handler), 30_000).catch((err) => {
+		logger.error("Exit handler failed to run", err);
+	});
+}
+
+export function initExitHandler() {
+	let signalCount = 0;
+	const exitHandler = async () => {
+		signalCount++;
+		if (signalCount === 1) {
+			logger.info("Received signal... Exiting");
+			await Promise.all(Array.from(listeners.values()).map(runExitHandler));
+			logger.info("All exit handlers ran... Waiting for svelte server to exit");
+		}
+		if (signalCount === 3) {
+			logger.warn("Received 3 signals... Exiting immediately");
+			process.exit(1);
+		}
+	};
+
+	process.on("SIGINT", exitHandler);
+	process.on("SIGTERM", exitHandler);
+}

src/lib/server/metrics.ts

@@ -4,6 +4,8 @@ import { logger } from "$lib/server/logger";
 import { env } from "$env/dynamic/private";
 import type { Model } from "$lib/types/Model";
 import type { Tool } from "$lib/types/Tool";
+import { onExit } from "./exitHandler";
+import { promisify } from "util";
 
 interface Metrics {
 	model: {
@@ -37,11 +39,26 @@ export class MetricsServer {
 
 	private constructor() {
 		const app = express();
-		const port = env.METRICS_PORT || "5565";
 
-		const server = app.listen(port, () => {
-			logger.info(`Metrics server listening on port ${port}`);
-		});
+		const port = Number(env.METRICS_PORT || "5565");
+		if (isNaN(port) || port < 0 || port > 65535) {
+			logger.warn(`Invalid value for METRICS_PORT: ${env.METRICS_PORT}`);
+		}
+
+		if (env.METRICS_ENABLED !== "false" && env.METRICS_ENABLED !== "true") {
+			logger.warn(`Invalid value for METRICS_ENABLED: ${env.METRICS_ENABLED}`);
+		}
+		if (env.METRICS_ENABLED === "true") {
+			const server = app.listen(port, () => {
+				logger.info(`Metrics server listening on port ${port}`);
+			});
+			const closeServer = promisify(server.close);
+			onExit(async () => {
+				logger.info("Disconnecting metrics server ...");
+				await closeServer();
+				logger.info("Server stopped ...");
+			});
+		}
 
 		const register = new Registry();
 		collectDefaultMetrics({ register });
@@ -160,14 +177,6 @@ export class MetricsServer {
 				res.send(metrics);
 			});
 		});
-
-		process.on("SIGINT", async () => {
-			logger.info("Sigint received, disconnect metrics server ...");
-			server.close(() => {
-				logger.info("Server stopped ...");
-			});
-			process.exit();
-		});
 	}
 
 	public static getInstance(): MetricsServer {

src/lib/server/websearch/scrape/playwright.ts

@@ -9,6 +9,7 @@ import {
 import { PlaywrightBlocker } from "@cliqz/adblocker-playwright";
 import { env } from "$env/dynamic/private";
 import { logger } from "$lib/server/logger";
+import { onExit } from "$lib/server/exitHandler";
 
 const blocker = await PlaywrightBlocker.fromPrebuiltAdsAndTracking(fetch)
 	.then((blker) => {
@@ -24,7 +25,7 @@ const blocker = await PlaywrightBlocker.fromPrebuiltAdsAndTracking(fetch)
 let browserSingleton: Promise<Browser> | undefined;
 async function getBrowser() {
 	const browser = await chromium.launch({ headless: true });
-	process.on("SIGINT", () => browser.close());
+	onExit(() => browser.close());
 	browser.on("disconnected", () => {
 		logger.warn("Browser closed");
 		browserSingleton = undefined;