jormungandr-bite/packages/backend/src/metrics.ts

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();
const metricsMaster = cluster.worker?.id === 1;

if (config.enableMetrics) {
	if (metricsMaster) {
		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 (metricsMaster) {
	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) => {
	try {
		if (metricsMaster) {
			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.enableMetrics) {
		if (metricsMaster) {
			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",
});

export const tickObliterate = counter({
		name: "iceshrimp_obliterate_total",
		help: "Total obliterate jobs processed",
	});
	

export const tickBiteIncoming = counter({
	name: "iceshrimp_bite_remote_incoming_total",
	help: "Total bites received from remote",
});

export const tickBiteOutgoing = counter({
	name: "iceshrimp_bite_remote_outgoing_total",
	help: "Total bites sent to remote",
});

export const tickBiteLocal = counter ({
	name: "iceshrimp_bite_local_total",
	help: "Total local bites"
});