From a70db57df946f8e037989f521d2827786e09cee6 Mon Sep 17 00:00:00 2001
From: mia <mia@mia.jetzt>
Date: Thu, 10 Oct 2024 16:45:56 -0700
Subject: [PATCH] [backend] Add prometheus metrics endpoint

---
 .pnp.cjs                                      |  40 ++++++
 ...ry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip |   3 +
 ...ntrees-npm-1.0.2-b28feeda03-071896cea5.zip |   3 +
 ...lient-npm-15.1.0-0b2231d02c-ecb6f40de7.zip |   3 +
 ...digest-npm-0.1.2-b73cfcf726-45be99fa52.zip |   3 +
 packages/backend/package.json                 |   1 +
 packages/backend/src/config/load.ts           |   3 +
 packages/backend/src/config/types.ts          |   4 +
 packages/backend/src/metrics.ts               | 125 ++++++++++++++++++
 packages/backend/src/queue/queues/deliver.ts  |   3 +
 packages/backend/src/queue/queues/inbox.ts    |   3 +
 .../src/remote/activitypub/check-fetch.ts     |   8 +-
 .../src/remote/activitypub/resolver.ts        |   7 +-
 packages/backend/src/server/activitypub.ts    |   5 +
 packages/backend/src/server/index.ts          |   5 +
 yarn.lock                                     |  34 +++++
 16 files changed, 248 insertions(+), 2 deletions(-)
 create mode 100644 .yarn/cache/@opentelemetry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip
 create mode 100644 .yarn/cache/bintrees-npm-1.0.2-b28feeda03-071896cea5.zip
 create mode 100644 .yarn/cache/prom-client-npm-15.1.0-0b2231d02c-ecb6f40de7.zip
 create mode 100644 .yarn/cache/tdigest-npm-0.1.2-b73cfcf726-45be99fa52.zip
 create mode 100644 packages/backend/src/metrics.ts

diff --git a/.pnp.cjs b/.pnp.cjs
index 68a836cb6..2edade284 100755
--- a/.pnp.cjs
+++ b/.pnp.cjs
@@ -3012,6 +3012,15 @@ const RAW_RUNTIME_STATE =
         "linkType": "HARD"\
       }]\
     ]],\
+    ["@opentelemetry/api", [\
+      ["npm:1.7.0", {\
+        "packageLocation": "./.yarn/cache/@opentelemetry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip/node_modules/@opentelemetry/api/",\
+        "packageDependencies": [\
+          ["@opentelemetry/api", "npm:1.7.0"]\
+        ],\
+        "linkType": "HARD"\
+      }]\
+    ]],\
     ["@paralleldrive/cuid2", [\
       ["npm:2.2.2", {\
         "packageLocation": "./.yarn/cache/@paralleldrive-cuid2-npm-2.2.2-e6061749b2-40ee269d6e.zip/node_modules/@paralleldrive/cuid2/",\
@@ -7299,6 +7308,7 @@ const RAW_RUNTIME_STATE =
           ["pg", "virtual:aa59773ac87791c4813d53447077fcf8a847d6de5a301d34dc31286584b1dbb26d30d3adb5b4c41c1e8aea04371e926fda05c09c6253647c432e11d872a304ba#npm:8.11.1"],\
           ["private-ip", "npm:2.3.4"],\
           ["probe-image-size", "npm:7.2.3"],\
+          ["prom-client", "npm:15.1.0"],\
           ["promise-limit", "npm:2.7.0"],\
           ["pug", "npm:3.0.2"],\
           ["punycode", "npm:2.3.0"],\
@@ -7474,6 +7484,15 @@ const RAW_RUNTIME_STATE =
         "linkType": "HARD"\
       }]\
     ]],\
+    ["bintrees", [\
+      ["npm:1.0.2", {\
+        "packageLocation": "./.yarn/cache/bintrees-npm-1.0.2-b28feeda03-071896cea5.zip/node_modules/bintrees/",\
+        "packageDependencies": [\
+          ["bintrees", "npm:1.0.2"]\
+        ],\
+        "linkType": "HARD"\
+      }]\
+    ]],\
     ["bl", [\
       ["npm:1.2.3", {\
         "packageLocation": "./.yarn/cache/bl-npm-1.2.3-49c4213ca5-11d775b09e.zip/node_modules/bl/",\
@@ -20443,6 +20462,17 @@ const RAW_RUNTIME_STATE =
         "linkType": "HARD"\
       }]\
     ]],\
+    ["prom-client", [\
+      ["npm:15.1.0", {\
+        "packageLocation": "./.yarn/cache/prom-client-npm-15.1.0-0b2231d02c-ecb6f40de7.zip/node_modules/prom-client/",\
+        "packageDependencies": [\
+          ["prom-client", "npm:15.1.0"],\
+          ["@opentelemetry/api", "npm:1.7.0"],\
+          ["tdigest", "npm:0.1.2"]\
+        ],\
+        "linkType": "HARD"\
+      }]\
+    ]],\
     ["promise", [\
       ["npm:7.3.1", {\
         "packageLocation": "./.yarn/cache/promise-npm-7.3.1-5d81d474c0-37dbe58ca7.zip/node_modules/promise/",\
@@ -23175,6 +23205,16 @@ const RAW_RUNTIME_STATE =
         "linkType": "HARD"\
       }]\
     ]],\
+    ["tdigest", [\
+      ["npm:0.1.2", {\
+        "packageLocation": "./.yarn/cache/tdigest-npm-0.1.2-b73cfcf726-45be99fa52.zip/node_modules/tdigest/",\
+        "packageDependencies": [\
+          ["tdigest", "npm:0.1.2"],\
+          ["bintrees", "npm:1.0.2"]\
+        ],\
+        "linkType": "HARD"\
+      }]\
+    ]],\
     ["terminal-link", [\
       ["npm:2.1.1", {\
         "packageLocation": "./.yarn/cache/terminal-link-npm-2.1.1-de80341758-ce3d2cd3a4.zip/node_modules/terminal-link/",\
diff --git a/.yarn/cache/@opentelemetry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip b/.yarn/cache/@opentelemetry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip
new file mode 100644
index 000000000..b54315020
--- /dev/null
+++ b/.yarn/cache/@opentelemetry-api-npm-1.7.0-6263fad98a-bcf7afa705.zip
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6b57709ed5bb5143e7607cd2eb729c80a582adaae0a482b9bf2c9ea76e693158
+size 607396
diff --git a/.yarn/cache/bintrees-npm-1.0.2-b28feeda03-071896cea5.zip b/.yarn/cache/bintrees-npm-1.0.2-b28feeda03-071896cea5.zip
new file mode 100644
index 000000000..d250f3aed
--- /dev/null
+++ b/.yarn/cache/bintrees-npm-1.0.2-b28feeda03-071896cea5.zip
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:58c1e47afedde6bb804fe2ebbdd917fabf21ef3b089093efd22d6367c170fec8
+size 1098081
diff --git a/.yarn/cache/prom-client-npm-15.1.0-0b2231d02c-ecb6f40de7.zip b/.yarn/cache/prom-client-npm-15.1.0-0b2231d02c-ecb6f40de7.zip
new file mode 100644
index 000000000..cf3b3ecb7
--- /dev/null
+++ b/.yarn/cache/prom-client-npm-15.1.0-0b2231d02c-ecb6f40de7.zip
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:060570fb9f7b9a5424b80532d4324342d3a1ddf87462af0203249fd398b8ae1a
+size 47950
diff --git a/.yarn/cache/tdigest-npm-0.1.2-b73cfcf726-45be99fa52.zip b/.yarn/cache/tdigest-npm-0.1.2-b73cfcf726-45be99fa52.zip
new file mode 100644
index 000000000..8311ca34a
--- /dev/null
+++ b/.yarn/cache/tdigest-npm-0.1.2-b73cfcf726-45be99fa52.zip
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b964b75caedfe79a4353fa7b523bbd3311a0b3da319e4b12fde84400ba721be0
+size 21869
diff --git a/packages/backend/package.json b/packages/backend/package.json
index f898ae03f..57e837b42 100644
--- a/packages/backend/package.json
+++ b/packages/backend/package.json
@@ -103,6 +103,7 @@
 		"pg": "8.11.1",
 		"private-ip": "2.3.4",
 		"probe-image-size": "7.2.3",
+		"prom-client": "^15.1.0",
 		"promise-limit": "2.7.0",
 		"punycode": "2.3.0",
 		"pureimage": "0.3.15",
diff --git a/packages/backend/src/config/load.ts b/packages/backend/src/config/load.ts
index 0c1971cdd..7120e8736 100644
--- a/packages/backend/src/config/load.ts
+++ b/packages/backend/src/config/load.ts
@@ -71,6 +71,9 @@ export default function load() {
 
 	config.searchEngine = config.searchEngine ?? 'https://duckduckgo.com/?q=';
 
+	config.metrics = config.metrics ?? {};
+	config.metrics.enable = config.metrics?.enable ?? false;
+
 	mixin.version = meta.version;
 	mixin.host = url.host;
 	mixin.hostname = url.hostname;
diff --git a/packages/backend/src/config/types.ts b/packages/backend/src/config/types.ts
index a8566cb31..9c84546e0 100644
--- a/packages/backend/src/config/types.ts
+++ b/packages/backend/src/config/types.ts
@@ -150,6 +150,10 @@ export type Source = {
 		s3ForcePathStyle?: boolean;
 	};
 	summalyProxyUrl?: string;
+	metrics?: {
+		enable?: boolean;
+		token?: string;
+	};
 };
 
 /**
diff --git a/packages/backend/src/metrics.ts b/packages/backend/src/metrics.ts
new file mode 100644
index 000000000..354e29444
--- /dev/null
+++ b/packages/backend/src/metrics.ts
@@ -0,0 +1,125 @@
+import Router from "@koa/router";
+import {
+	collectDefaultMetrics,
+	register,
+	Gauge,
+	Counter,
+	CounterConfiguration,
+} from "prom-client";
+import config from "./config/index.js";
+import { queues } from "./queue/queues.js";
+import cluster from "node:cluster";
+import Xev from "xev";
+
+const xev = new Xev();
+
+if (config.metrics?.enable) {
+	if (cluster.isPrimary) {
+		collectDefaultMetrics();
+
+		new Gauge({
+			name: "iceshrimp_queue_jobs",
+			help: "Amount of jobs in the bull queues",
+			labelNames: ["queue", "status"] as const,
+			async collect() {
+				for (const queue of queues) {
+					const counts = await queue.getJobCounts();
+					this.set({ queue: queue.name, status: "completed" }, counts.completed);
+					this.set({ queue: queue.name, status: "waiting" }, counts.waiting);
+					this.set({ queue: queue.name, status: "active" }, counts.active);
+					this.set({ queue: queue.name, status: "delayed" }, counts.delayed);
+					this.set({ queue: queue.name, status: "failed" }, counts.failed);
+				}
+			},
+		});
+	}
+}
+
+if (cluster.isPrimary) {
+	xev.on("registry-request", async () => {
+		try {
+			const metrics = await register.metrics();
+			xev.emit("registry-response", {
+				contentType: register.contentType,
+				body: metrics
+			});
+		} catch (error) {
+			xev.emit("registry-response", { error });
+		}
+	});
+}
+
+export const handleMetrics: Router.Middleware = async (ctx) => {
+	if (config.metrics?.token !== undefined) {
+		if (ctx.query.token === undefined) {
+			ctx.res.statusCode = 401;
+			ctx.body = "Missing token parameter";
+			return;
+		}
+		const correct = config.metrics.token === ctx.query.token;
+		if (!correct) {
+			ctx.res.statusCode = 403;
+			ctx.body = "Incorrect token";
+			return;
+		}
+	}
+	try {
+		if (cluster.isPrimary) {
+			ctx.set("content-type", register.contentType);
+			ctx.body = await register.metrics();
+		} else {
+			const wait = new Promise<void>((resolve, reject) => {
+				const timeout = setTimeout(
+					() => reject("Timeout while waiting for cluster master"),
+					1000 * 60
+				);
+				xev.once("registry-response", (response) => {
+					clearTimeout(timeout);
+					if (response.error) reject(response.error);
+					ctx.set("content-type", response.contentType);
+					ctx.body = response.body;
+					resolve();
+				});
+			});
+			xev.emit("registry-request");
+			await wait;
+		}
+	} catch (err) {
+		ctx.res.statusCode = 500;
+		ctx.body = err;
+	}
+};
+
+const counter = (configuration: CounterConfiguration<string>) => {
+	if (config.metrics?.enable) {
+		if (cluster.isPrimary) {
+			const counter = new Counter(configuration);
+			counter.reset(); // initialize internal hashmap
+			xev.on(`metrics-counter-${configuration.name}`, () => counter.inc());
+			return () => counter.inc();
+		} else {
+			return () => xev.emit(`metrics-counter-${configuration.name}`);
+		}
+	} else {
+		return () => { };
+	}
+};
+
+export const tickOutbox = counter({
+	name: "iceshrimp_outbox_total",
+	help: "Total AP outbox calls",
+});
+
+export const tickInbox = counter({
+	name: "iceshrimp_inbox_total",
+	help: "Total AP inbox calls",
+});
+
+export const tickFetch = counter({
+	name: "iceshrimp_fetch_total",
+	help: "Total AP fetch calls",
+});
+
+export const tickResolve = counter({
+	name: "iceshrimp_resolve_total",
+	help: "Total AP resolve calls",
diff --git a/packages/backend/src/queue/queues/deliver.ts b/packages/backend/src/queue/queues/deliver.ts
index a3544247c..7216a8afb 100644
--- a/packages/backend/src/queue/queues/deliver.ts
+++ b/packages/backend/src/queue/queues/deliver.ts
@@ -17,6 +17,7 @@ import config from "@/config/index.js";
 import { createQueue, processorTimeout } from "./index.js";
 import { ThinUser } from "../types.js";
 import { Job } from "bullmq";
+import { tickOutbox } from "@/metrics.js";
 
 export const deliverLogger = new Logger("deliver");
 
@@ -55,6 +56,8 @@ async function process(job: Job<DeliverJobData>) {
 			federationChart.deliverd(i.host, true);
 		});
 
+		tickOutbox();
+
 		return "Success";
 	} catch (res) {
 		// Update stats
diff --git a/packages/backend/src/queue/queues/inbox.ts b/packages/backend/src/queue/queues/inbox.ts
index bb0bc2438..c462f98e4 100644
--- a/packages/backend/src/queue/queues/inbox.ts
+++ b/packages/backend/src/queue/queues/inbox.ts
@@ -25,6 +25,7 @@ import { verifySignature } from "@/remote/activitypub/check-fetch.js";
 import { Job } from "bullmq";
 import { createQueue, processorTimeout } from "./index.js";
 import config from "@/config/index.js";
+import { tickInbox } from "@/metrics.js";
 
 export const inboxLogger = new Logger("inbox");
 
@@ -221,6 +222,8 @@ async function process(job: Job<InboxJobData>): Promise<string> {
 		}
 	}
 
+	tickInbox();
+
 	// アクティビティを処理
 	await perform(authUser.user, activity);
 	return "ok";
diff --git a/packages/backend/src/remote/activitypub/check-fetch.ts b/packages/backend/src/remote/activitypub/check-fetch.ts
index b7a4b4742..d9f3aeeac 100644
--- a/packages/backend/src/remote/activitypub/check-fetch.ts
+++ b/packages/backend/src/remote/activitypub/check-fetch.ts
@@ -12,6 +12,7 @@ import type { UserPublickey } from "@/models/entities/user-publickey.js";
 import { verify } from "node:crypto";
 import { toSingle } from "@/prelude/array.js";
 import { createHash } from "node:crypto";
+import { tickFetch } from "@/metrics.js";
 
 export async function hasSignature(req: IncomingMessage): Promise<string> {
 	const meta = await fetchMeta();
@@ -120,7 +121,12 @@ export async function checkFetch(req: IncomingMessage): Promise<number> {
 			return 403;
 		}
 
-		return verifySignature(signature, authUser.key) ? 200 : 401;
+		if (!verifySignature(signature, authUser.key)) {
+			return 401;
+		}
+
+		tickFetch();
+		return 200;
 	}
 	return 200;
 }
diff --git a/packages/backend/src/remote/activitypub/resolver.ts b/packages/backend/src/remote/activitypub/resolver.ts
index 17cf50e2a..c333a1e00 100644
--- a/packages/backend/src/remote/activitypub/resolver.ts
+++ b/packages/backend/src/remote/activitypub/resolver.ts
@@ -25,6 +25,7 @@ import renderFollow from "@/remote/activitypub/renderer/follow.js";
 import { shouldBlockInstance } from "@/misc/should-block-instance.js";
 import { apLogger } from "@/remote/activitypub/logger.js";
 import { In, IsNull, Not } from "typeorm";
+import { tickResolve } from "@/metrics.js";
 
 export default class Resolver {
 	private history: Set<string>;
@@ -127,7 +128,10 @@ export default class Resolver {
 		if (object.id == null) throw new Error("Object has no ID");
 		const objectId = new URL(object.id);
 		const resFinalUrl = new URL(res.finalUrl);
-		if (resFinalUrl.toString() === objectId.toString()) return object;
+		if (resFinalUrl.toString() === objectId.toString()) {
+			tickResolve();
+			return object;
+		}
 
 		if (resFinalUrl.host !== objectId.host)
 			throw new Error("Object ID host doesn't match final url host");
@@ -141,6 +145,7 @@ export default class Resolver {
 		if (finalResFinalUrl.toString() !== finalObjectId.toString())
 			throw new Error("Object ID still doesn't match final URL after second fetch attempt")
 
+		tickResolve();
 		return finalObject;
 	}
 
diff --git a/packages/backend/src/server/activitypub.ts b/packages/backend/src/server/activitypub.ts
index 9df6a0c7d..7a3d31861 100644
--- a/packages/backend/src/server/activitypub.ts
+++ b/packages/backend/src/server/activitypub.ts
@@ -35,6 +35,7 @@ import Outbox, { packActivity } from "./activitypub/outbox.js";
 import { serverLogger } from "./index.js";
 import config from "@/config/index.js";
 import Koa from "koa";
+import { tickFetch } from "@/metrics.js";
 
 // Init router
 const router = new Router();
@@ -223,6 +224,7 @@ router.get("/users/:user/collections/featured", Featured);
 router.get("/users/:user/publickey", async (ctx) => {
 	const instanceActor = await getInstanceActor();
 	if (ctx.params.user === instanceActor.id) {
+		tickFetch();
 		ctx.body = renderActivity(
 			renderKey(instanceActor, await getUserKeypair(instanceActor.id)),
 		);
@@ -287,6 +289,7 @@ router.get("/users/:user", async (ctx, next) => {
 
 	const instanceActor = await getInstanceActor();
 	if (ctx.params.user === instanceActor.id) {
+		tickFetch();
 		await userInfo(ctx, instanceActor);
 		return;
 	}
@@ -312,6 +315,7 @@ router.get("/@:user", async (ctx, next) => {
 	if (!isActivityPubReq(ctx)) return await next();
 
 	if (ctx.params.user === "instance.actor") {
+		tickFetch();
 		const instanceActor = await getInstanceActor();
 		await userInfo(ctx, instanceActor);
 		return;
@@ -333,6 +337,7 @@ router.get("/@:user", async (ctx, next) => {
 });
 
 router.get("/actor", async (ctx, next) => {
+	tickFetch();
 	const instanceActor = await getInstanceActor();
 	await userInfo(ctx, instanceActor);
 });
diff --git a/packages/backend/src/server/index.ts b/packages/backend/src/server/index.ts
index 929efd73b..c81663729 100644
--- a/packages/backend/src/server/index.ts
+++ b/packages/backend/src/server/index.ts
@@ -33,6 +33,7 @@ import removeTrailingSlash from "koa-remove-trailing-slashes";
 import { koaBody } from "koa-body";
 import { setupEndpointsAuthRoot } from "@/server/api/mastodon/endpoints/auth.js";
 import { CatchErrorsMiddleware } from "@/server/api/mastodon/middleware/catch-errors.js";
+import { handleMetrics } from "@/metrics.js";
 export const serverLogger = new Logger("server", "gray", false);
 
 // Init app
@@ -117,6 +118,10 @@ router.get("/identicon/:x", async (ctx) => {
 	}
 });
 
+if (config.metrics?.enable) {
+	router.get("/metrics", handleMetrics);
+}
+
 mastoRouter.use(
 	koaBody({
 		urlencoded: true,
diff --git a/yarn.lock b/yarn.lock
index e996c1e17..802a911c7 100644
--- a/yarn.lock
+++ b/yarn.lock
@@ -2159,6 +2159,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"@opentelemetry/api@npm:^1.4.0":
+  version: 1.7.0
+  resolution: "@opentelemetry/api@npm:1.7.0"
+  checksum: 10/bcf7afa7051dcd4583898a68f8a57fb4c85b5cedddf7b6eb3616595c0b3bcd7f5448143b8355b00935a755de004d6285489f8e132f34127efe7b1be404622a3e
+  languageName: node
+  linkType: hard
+
 "@paralleldrive/cuid2@npm:^2.2.2":
   version: 2.2.2
   resolution: "@paralleldrive/cuid2@npm:2.2.2"
@@ -5599,6 +5606,7 @@ __metadata:
     pg: "npm:8.11.1"
     private-ip: "npm:2.3.4"
     probe-image-size: "npm:7.2.3"
+    prom-client: "npm:^15.1.0"
     promise-limit: "npm:2.7.0"
     pug: "npm:3.0.2"
     punycode: "npm:2.3.0"
@@ -5759,6 +5767,13 @@ __metadata:
   languageName: node
   linkType: hard
 
+"bintrees@npm:1.0.2":
+  version: 1.0.2
+  resolution: "bintrees@npm:1.0.2"
+  checksum: 10/071896cea5ea5413316c8436e95799444c208630d5c539edd8a7089fc272fc5d3634aa4a2e4847b28350dda1796162e14a34a0eda53108cc5b3c2ff6a036c1fa
+  languageName: node
+  linkType: hard
+
 "bl@npm:^1.0.0":
   version: 1.2.3
   resolution: "bl@npm:1.2.3"
@@ -17072,6 +17087,16 @@ __metadata:
   languageName: node
   linkType: hard
 
+"prom-client@npm:^15.1.0":
+  version: 15.1.0
+  resolution: "prom-client@npm:15.1.0"
+  dependencies:
+    "@opentelemetry/api": "npm:^1.4.0"
+    tdigest: "npm:^0.1.1"
+  checksum: 10/ecb6f40de755ca9cc6dde758d195ed3e1d3b47a341d2092af8c18dbf7e6ef1079c8b8bb02496f2f430cf8bd9d391c1ea5bebbb85cdda95f67dad2dbfb90509aa
+  languageName: node
+  linkType: hard
+
 "promise-limit@npm:2.7.0":
   version: 2.7.0
   resolution: "promise-limit@npm:2.7.0"
@@ -19630,6 +19655,15 @@ __metadata:
   languageName: node
   linkType: hard
 
+"tdigest@npm:^0.1.1":
+  version: 0.1.2
+  resolution: "tdigest@npm:0.1.2"
+  dependencies:
+    bintrees: "npm:1.0.2"
+  checksum: 10/45be99fa52dab74b8edafe150e473cdc45aa1352c75ed516a39905f350a08c3175f6555598111042c3677ba042d7e3cae6b5ce4c663fe609bc634f326aabc9d6
+  languageName: node
+  linkType: hard
+
 "terminal-link@npm:^2.0.0":
   version: 2.1.1
   resolution: "terminal-link@npm:2.1.1"