From 3c97ff74a975ff6efbf264fa32290e1d4d47a6de Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Sun, 6 Nov 2022 07:17:11 +0800
Subject: [PATCH] prom-client: add category for connection handler

---
 packages/hydrooj/src/handler/home.ts  | 2 ++
 packages/hydrooj/src/handler/judge.ts | 1 +
 packages/prom-client/metrics.ts       | 4 ++--
 3 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/packages/hydrooj/src/handler/home.ts b/packages/hydrooj/src/handler/home.ts
index 52cad05d..f82344bf 100644
--- a/packages/hydrooj/src/handler/home.ts
+++ b/packages/hydrooj/src/handler/home.ts
@@ -363,6 +363,8 @@ class HomeDomainCreateHandler extends Handler {
 }
 
 class HomeMessagesHandler extends Handler {
+    category = '#message';
+
     async get() {
         // TODO(iceboy): projection, pagination.
         const messages = await message.getByUser(this.user._id);
diff --git a/packages/hydrooj/src/handler/judge.ts b/packages/hydrooj/src/handler/judge.ts
index fa80a3a1..29e9eacd 100644
--- a/packages/hydrooj/src/handler/judge.ts
+++ b/packages/hydrooj/src/handler/judge.ts
@@ -172,6 +172,7 @@ export class SubmissionDataDownloadHandler extends Handler {
 }
 
 class JudgeConnectionHandler extends ConnectionHandler {
+    category = '#judge';
     processing: any = null;
     closed = false;
     query: any = { type: 'judge' };
diff --git a/packages/prom-client/metrics.ts b/packages/prom-client/metrics.ts
index 286e07f4..7d8f18e0 100644
--- a/packages/prom-client/metrics.ts
+++ b/packages/prom-client/metrics.ts
@@ -40,10 +40,10 @@ export function createRegistry(ctx: Context) {
         labelNames: ['domainId'],
     });
     ctx.on('connection/active', (h) => {
-        connectionGauge.inc({ domainId: h.args.domainId });
+        connectionGauge.inc({ domainId: (h as any).category || h.args.domainId });
     });
     ctx.on('connection/close', (h) => {
-        connectionGauge.dec({ domainId: h.args.domainId });
+        connectionGauge.dec({ domainId: (h as any).category || h.args.domainId });
     });
 
     const taskColl = db.collection('task');