prom: export metrics

packages/hydrooj/src/service/server.ts

@@ -36,7 +36,7 @@ export interface HydroRequest {
     host: string;
     hostname: string;
     ip: string;
-    headers: any;
+    headers: Koa.Request['headers'];
     cookies: any;
     body: any;
     files: Record<string, import('formidable').File>;

packages/prom-client/handler.ts

@@ -0,0 +1,52 @@
+import { AggregatorRegistry, metric } from 'prom-client';
+import * as system from 'hydrooj/src/model/system';
+import * as bus from 'hydrooj/src/service/bus';
+import { Handler, Route } from 'hydrooj/src/service/server';
+import { registry } from './metrics';
+
+declare module 'hydrooj/src/service/bus' {
+    interface EventMap {
+        metrics: (id: string, metrics: any) => void;
+    }
+}
+declare module 'hydrooj/src/interface' {
+    interface SystemKeys {
+        'prom-client.name': string;
+        'prom-client.password': string;
+        'prom-client.collect_rate': number;
+    }
+}
+
+const instances: Record<string, metric[]> = {};
+
+class MetricsHandler extends Handler {
+    noCheckPermView = true;
+    notUsage = true;
+
+    async get() {
+        if (!this.request.headers.authorization) {
+            this.response.status = 401;
+            this.response.body = {};
+            this.response.addHeader('WWW-Authenticate', 'Basic');
+            return;
+        }
+        const [name, password] = system.getMany(['prom-client.name', 'prom-client.password']);
+        const key = this.request.headers.authorization?.split('Basic ')?.[1];
+        if (!key || key !== Buffer.from(`${name}:${password}`).toString('base64')) {
+            this.response.status = 403;
+            this.response.body = {};
+            return;
+        }
+        this.response.body = await AggregatorRegistry.aggregate(Object.values(instances)).metrics();
+        this.response.type = 'text/plain';
+    }
+}
+
+bus.on('metrics', (id, metrics) => { instances[id] = metrics; });
+setInterval(async () => {
+    bus.broadcast('metrics', process.env.NODE_APP_INSTANCE!, await registry.getMetricsAsJSON());
+}, 5000 * (+system.get('prom-client.collect_rate') || 1));
+
+global.Hydro.handler.prom = () => {
+    Route('metrics', '/metrics', MetricsHandler);
+};

packages/prom-client/metrics.ts

@@ -0,0 +1,51 @@
+import {
+    collectDefaultMetrics, Counter, Gauge, Metric, Registry,
+} from 'prom-client';
+import * as bus from 'hydrooj/src/service/bus';
+import * as db from 'hydrooj/src/service/db';
+
+export const registry = new Registry();
+registry.setDefaultLabels({ instanceid: process.env.NODE_APP_INSTANCE });
+function createMetric<Q extends string, T extends (new (a: any) => Metric<Q>)>(
+    C: T, name: string, help: string, extra?: T extends new (a: infer R) => any ? Partial<R> : never,
+): T extends new (a) => Counter<Q> ? Counter<Q> : T extends new (a) => Gauge<Q> ? Gauge<Q> : Metric<Q> {
+    const metric = new C({ name, help, ...(extra || {}) });
+    registry.registerMetric(metric);
+    return metric as any;
+}
+
+const reqCounter = createMetric(Counter, 'hydro_reqcount', 'reqcount', {
+    labelNames: ['domainId'],
+});
+bus.on('handler/create', (h) => reqCounter.inc({ domainId: h.args.domainId }));
+
+const judgeCounter = createMetric(Counter, 'hydro_judgecount', 'judgecount');
+bus.on('record/judge', () => judgeCounter.inc());
+
+createMetric(Gauge, 'hydro_regcount', 'regcount', {
+    async collect() {
+        this.set({}, await db.collection('user').countDocuments());
+    },
+});
+
+const submissionCounter = createMetric(Counter, 'hydro_submission', 'submissioncount', {
+    labelNames: ['lang', 'domainId'],
+});
+bus.on('handler/after/ProblemSubmit', (that) => {
+    submissionCounter.inc({ lang: that.args.lang, domainId: that.args.domainId });
+});
+
+const taskColl = db.collection('task');
+createMetric(Gauge, 'hydro_task', 'taskcount', {
+    async collect() {
+        const data = await taskColl.aggregate([
+            { $group: { _id: '$type', count: { $sum: 1 } } },
+        ]).toArray();
+        for (const line of data) {
+            this.set({ type: line._id as unknown as string }, line.count);
+        }
+    },
+    labelNames: ['type'],
+});
+
+collectDefaultMetrics({ register: registry });

packages/prom-client/package.json

@@ -0,0 +1,7 @@
+{
+    "name": "@hydrooj/prom-client",
+    "version": "0.0.1",
+    "dependencies": {
+        "prom-client": "^14.0.1"
+    }
+}

packages/prom-client/setting.yaml

@@ -0,0 +1,15 @@
+name:
+  type: text
+  name: Auth Username
+  desc: basic auth username while requesting metrics
+  default: admin
+password:
+  type: password
+  name: Auth Password
+  desc: basic auth password while requesting metrics
+  default: admin
+collect_rate:
+  type: number
+  name: Collect Rate
+  desc: The collect rate scale (default:1)
+  default: 1