From 93431826075103962e4c715fed5ddd457fb25f33 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Fri, 20 Jan 2023 15:02:24 +0800
Subject: [PATCH] core: add cache for domain

---
 packages/hydrooj/src/model/domain.ts | 41 ++++++++++++++++++++++------
 packages/hydrooj/src/service/bus.ts  |  1 +
 2 files changed, 34 insertions(+), 8 deletions(-)

diff --git a/packages/hydrooj/src/model/domain.ts b/packages/hydrooj/src/model/domain.ts
index 62dd9dca..87ff9174 100644
--- a/packages/hydrooj/src/model/domain.ts
+++ b/packages/hydrooj/src/model/domain.ts
@@ -1,4 +1,5 @@
 import { Dictionary } from 'lodash';
+import LRU from 'lru-cache';
 import { FilterQuery } from 'mongodb';
 import { DomainDoc } from '../interface';
 import * as bus from '../service/bus';
@@ -10,6 +11,7 @@ import UserModel, { deleteUserCache } from './user';
 
 const coll = db.collection('domain');
 const collUser = db.collection('domain.user');
+const cache = new LRU<string, DomainDoc>({ max: 1000, ttl: 300 * 1000 });
 
 interface DomainUserArg {
     _id: number,
@@ -59,19 +61,31 @@ class DomainModel {
 
     @ArgMethod
     static async get(domainId: string): Promise<DomainDoc | null> {
-        const query: FilterQuery<DomainDoc> = { lower: domainId.toLowerCase() };
+        domainId = domainId.toLowerCase();
+        const key = `id::${domainId}`;
+        if (cache.has(key)) return cache.get(key);
+        const query: FilterQuery<DomainDoc> = { lower: domainId };
         await bus.parallel('domain/before-get', query);
         const result = await coll.findOne(query);
-        if (result) await bus.parallel('domain/get', result);
+        if (result) {
+            await bus.parallel('domain/get', result);
+            cache.set(key, result);
+        }
         return result;
     }
 
     @ArgMethod
     static async getByHost(host: string): Promise<DomainDoc | null> {
+        const key = `host::${host}`;
+        // Note: cache by host might not be updated immediately
+        if (cache.has(key)) return cache.get(key);
         const query: FilterQuery<DomainDoc> = { host };
         await bus.parallel('domain/before-get', query);
         const result = await coll.findOne(query);
-        if (result) await bus.parallel('domain/get', result);
+        if (result) {
+            await bus.parallel('domain/get', result);
+            cache.set(key, result);
+        }
         return result;
     }
 
@@ -80,28 +94,32 @@ class DomainModel {
     }
 
     static async edit(domainId: string, $set: Partial<DomainDoc>) {
+        domainId = domainId.toLowerCase();
         await bus.parallel('domain/before-update', domainId, $set);
         const result = await coll.findOneAndUpdate({ _id: domainId }, { $set }, { returnDocument: 'after' });
-        if (result.value) await bus.parallel('domain/update', domainId, $set, result.value);
+        if (result.value) {
+            await bus.parallel('domain/update', domainId, $set, result.value);
+            cache.delete(`id::${domainId}`);
+        }
         return result.value;
     }
 
     @ArgMethod
     static async inc(domainId: string, field: NumberKeys<DomainDoc>, n: number): Promise<number | null> {
+        domainId = domainId.toLowerCase();
         const res = await coll.findOneAndUpdate(
             { _id: domainId },
             { $inc: { [field]: n } as any },
             { returnDocument: 'after' },
         );
+        bus.broadcast('domain/delete-cache', domainId);
         return res.value?.[field];
     }
 
     @ArgMethod
     static async getList(domainIds: string[]) {
         const r: Record<string, DomainDoc | null> = {};
-        const tasks = [];
-        for (const domainId of domainIds) tasks.push(DomainModel.get(domainId).then((ddoc) => { r[domainId] = ddoc; }));
-        await Promise.all(tasks);
+        await Promise.all(domainIds.map((domainId) => DomainModel.get(domainId).then((ddoc) => { r[domainId] = ddoc; })));
         return r;
     }
 
@@ -152,17 +170,19 @@ class DomainModel {
     }
 
     static async setRoles(domainId: string, roles: Dictionary<bigint | string>) {
-        deleteUserCache(domainId);
         const current = await DomainModel.get(domainId);
         for (const role in roles) {
             current.roles[role] = roles[role].toString();
         }
+        deleteUserCache(domainId);
+        bus.broadcast('domain/delete-cache', domainId.toLowerCase());
         return await coll.updateOne({ _id: domainId }, { $set: { roles: current.roles } });
     }
 
     static async addRole(domainId: string, name: string, permission: bigint) {
         const current = await DomainModel.get(domainId);
         current.roles[name] = permission.toString();
+        bus.broadcast('domain/delete-cache', domainId.toLowerCase());
         return await coll.updateOne({ _id: domainId }, { $set: { roles: current.roles } });
     }
 
@@ -174,6 +194,7 @@ class DomainModel {
             collUser.updateMany({ domainId, role: { $in: roles } }, { $set: { role: 'default' } }),
         ]);
         deleteUserCache(domainId);
+        bus.broadcast('domain/delete-cache', domainId.toLowerCase());
     }
 
     static async getDomainUser(domainId: string, udoc: DomainUserArg) {
@@ -254,6 +275,7 @@ class DomainModel {
         await coll.deleteOne({ _id: domainId });
         await collUser.deleteMany({ domainId });
         await bus.parallel('domain/delete', domainId);
+        bus.broadcast('domain/delete-cache', domainId.toLowerCase());
     }
 
     @ArgMethod
@@ -274,5 +296,8 @@ bus.on('ready', () => Promise.all([
         { key: { domainId: 1, rp: -1, uid: 1 }, name: 'rp', sparse: true },
     ),
 ]));
+bus.on('domain/delete-cache', (domainId: string) => {
+    cache.del(`id::${domainId}`);
+});
 export default DomainModel;
 global.Hydro.model.domain = DomainModel;
diff --git a/packages/hydrooj/src/service/bus.ts b/packages/hydrooj/src/service/bus.ts
index eb063eb9..d3335c26 100644
--- a/packages/hydrooj/src/service/bus.ts
+++ b/packages/hydrooj/src/service/bus.ts
@@ -59,6 +59,7 @@ export interface EventMap extends LifecycleEvents, HandlerEvents {
     'domain/before-update': (domainId: string, $set: Partial<DomainDoc>) => VoidReturn
     'domain/update': (domainId: string, $set: Partial<DomainDoc>, ddoc: DomainDoc) => VoidReturn
     'domain/delete': (domainId: string) => VoidReturn
+    'domain/delete-cache': (domainId: string) => VoidReturn
 
     'document/add': (doc: any) => VoidReturn
     'document/set': <T extends keyof DocType>(