From 0726dba4921ac0f6ff5f3bc7a8c9c8d76298cba1 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Mon, 30 Jan 2023 11:40:19 +0800
Subject: [PATCH] core: optimize database operation

---
 packages/hydrooj/src/handler/status.ts |  2 +-
 packages/hydrooj/src/model/task.ts     |  2 +-
 packages/hydrooj/src/service/db.ts     | 16 ++++++++++++++++
 3 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/packages/hydrooj/src/handler/status.ts b/packages/hydrooj/src/handler/status.ts
index dd9e4a2de..ab729f1e1 100644
--- a/packages/hydrooj/src/handler/status.ts
+++ b/packages/hydrooj/src/handler/status.ts
@@ -74,5 +74,5 @@ export async function apply(ctx: Context) {
     ctx.Route('status', '/status', StatusHandler);
     ctx.Route('status_admin', '/.status', AdminStatusHandler);
     ctx.Route('status_update', '/status/update', StatusUpdateHandler);
-    await db.ensureIndexes(coll, { key: { updateAt: 1 }, expireAfterSeconds: 24 * 2600 });
+    await db.ensureIndexes(coll, { name: 'expire', key: { updateAt: 1 }, expireAfterSeconds: 24 * 2600 });
 }
diff --git a/packages/hydrooj/src/model/task.ts b/packages/hydrooj/src/model/task.ts
index eaa6a9b6a..822655472 100644
--- a/packages/hydrooj/src/model/task.ts
+++ b/packages/hydrooj/src/model/task.ts
@@ -130,7 +130,7 @@ export async function apply(ctx: Context) {
         while (true) {
             // eslint-disable-next-line no-await-in-loop
             const res = await collEvent.findOneAndUpdate(
-                { ack: { $nin: [id] } },
+                { expire: { $gt: new Date() }, ack: { $nin: [id] } },
                 { $push: { ack: id } },
             );
             // eslint-disable-next-line no-await-in-loop
diff --git a/packages/hydrooj/src/service/db.ts b/packages/hydrooj/src/service/db.ts
index adb275574..0cef5aacf 100644
--- a/packages/hydrooj/src/service/db.ts
+++ b/packages/hydrooj/src/service/db.ts
@@ -2,6 +2,7 @@
 import {
     Collection, Db, IndexSpecification, MongoClient,
 } from 'mongodb';
+import { Time } from '@hydrooj/utils';
 import { Logger } from '../logger';
 import options from '../options';
 import * as bus from './bus';
@@ -51,6 +52,7 @@ class MongoService {
         });
         this.db = this.client.db(opts.name || 'hydro');
         await bus.parallel('database/connect', this.db);
+        setInterval(() => this.fixExpireAfter(), Time.hour);
     }
 
     public collection<K extends keyof Collections>(c: K): Collection<Collections[K]> {
@@ -58,6 +60,20 @@ class MongoService {
         return this.db.collection(c);
     }
 
+    public async fixExpireAfter() {
+        // Sometimes mongo's expireAfterSeconds is not working in non-replica set mode;
+        const collections = await this.db.listCollections().toArray();
+        for (const c of collections) {
+            const coll = this.db.collection(c.name);
+            const indexes = await coll.listIndexes().toArray();
+            for (const i of indexes) {
+                if (!i.expireAfterSeconds) continue;
+                const key = Object.keys(i.key)[0];
+                await coll.deleteMany({ [key]: { $lt: new Date(Date.now() - i.expireAfterSeconds * 1000) } });
+            }
+        }
+    }
+
     public async ensureIndexes<T>(coll: Collection<T>, ...args: IndexSpecification[]) {
         if (process.env.NODE_APP_INSTANCE !== '0') return;
         let existed: any[];