From 3ad9d5e614228fcbaec1db406ffd9d273fe16813 Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Sat, 6 Apr 2024 05:37:52 +0800
Subject: [PATCH] core: support random picking cdn and ws endpoints

---
 packages/hydrooj/src/service/layers/base.ts | 9 +++++++--
 packages/utils/lib/common.ts                | 4 ++++
 2 files changed, 11 insertions(+), 2 deletions(-)

diff --git a/packages/hydrooj/src/service/layers/base.ts b/packages/hydrooj/src/service/layers/base.ts
index 8faaab745..08d27bc67 100644
--- a/packages/hydrooj/src/service/layers/base.ts
+++ b/packages/hydrooj/src/service/layers/base.ts
@@ -1,6 +1,7 @@
 import { PassThrough } from 'stream';
 import type { Next } from 'koa';
 import { cloneDeep, omit, pick } from 'lodash';
+import { randomPick } from '@hydrooj/utils';
 import { PERM } from '../../model/builtin';
 import * as system from '../../model/system';
 import token from '../../model/token';
@@ -59,8 +60,12 @@ export default async (ctx: KoaContext, next: Next) => {
         domainId, ...ctx.params, ...ctx.query, ...ctx.request.body, __start: Date.now(),
     };
     const UiContext: any = cloneDeep(UiContextBase);
-    if (!process.env.DEV) UiContext.cdn_prefix = system.get('server.cdn');
-    if (!process.env.DEV) UiContext.ws_prefix = system.get('server.ws');
+    if (!process.env.DEV) {
+        UiContext.cdn_prefix = system.get('server.cdn');
+        if (UiContext.cdn_prefix.includes(',')) UiContext.cdn_prefix = randomPick(UiContext.cdn_prefix.split(','));
+        UiContext.ws_prefix = system.get('server.ws');
+        if (UiContext.ws_prefix.includes(',')) UiContext.ws_prefix = randomPick(UiContext.ws_prefix.split(','));
+    }
     UiContext.domainId = domainId;
     UiContext.domain = domainInfo;
     ctx.HydroContext = {
diff --git a/packages/utils/lib/common.ts b/packages/utils/lib/common.ts
index a254ba670..be43a3521 100644
--- a/packages/utils/lib/common.ts
+++ b/packages/utils/lib/common.ts
@@ -199,6 +199,10 @@ export function size(s: number, base = 1) {
     return `${Math.round(s * unit)} ${unitNames[unitNames.length - 1]}`;
 }
 
+export function randomPick<T>(arr: T[]): T {
+    return arr[Math.floor(Math.random() * arr.length)];
+}
+
 export type StringKeys<O> = {
     [K in keyof O]: string extends O[K] ? K : never
 }[keyof O];