From 1424232a80af95f879da30852130a859e693a32c Mon Sep 17 00:00:00 2001
From: undefined <i@undefined.moe>
Date: Tue, 28 Mar 2023 15:18:57 +0800
Subject: [PATCH] core&ui: add connection heartbeat

---
 packages/hydrooj/src/service/server.ts         | 14 +++++++++++++-
 packages/ui-default/api.ts                     |  1 +
 packages/ui-default/components/socket/index.ts | 12 ++++++++++--
 3 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/packages/hydrooj/src/service/server.ts b/packages/hydrooj/src/service/server.ts
index 783e28a8f..d0f5d9e3c 100644
--- a/packages/hydrooj/src/service/server.ts
+++ b/packages/hydrooj/src/service/server.ts
@@ -383,9 +383,21 @@ export function Connection(
             if (h.prepare) await h.prepare(args);
             // eslint-disable-next-line @typescript-eslint/no-shadow
             for (const { name, target } of h.__subscribe || []) disposables.push(bus.on(name, target.bind(h)));
-            conn.onmessage = (e) => h.message?.(JSON.parse(e.data.toString()));
+            let lastHeartbeat = Date.now();
+            const interval = setInterval(() => {
+                if (Date.now() - lastHeartbeat > 40000) conn.close(4000, 'Heartbeat timeout');
+            }, 40000);
+            conn.onmessage = (e) => {
+                if (e.data === 'ping') {
+                    conn.send('pong');
+                    lastHeartbeat = Date.now();
+                    return;
+                }
+                h.message?.(JSON.parse(e.data.toString()));
+            };
             conn.onclose = () => {
                 bus.emit('connection/close', h);
+                clearInterval(interval);
                 disposables.forEach((d) => d());
                 h.cleanup?.(args);
             };
diff --git a/packages/ui-default/api.ts b/packages/ui-default/api.ts
index 24b8e7eb9..465277b02 100644
--- a/packages/ui-default/api.ts
+++ b/packages/ui-default/api.ts
@@ -6,6 +6,7 @@ export * from './components/dialog';
 export { Context, Service, ctx } from './context';
 export { default as AnsiUp } from 'ansi_up';
 export { default as loadMonaco } from './components/monaco/loader';
+export { default as Socket } from './components/socket/index';
 export { default as uploadFiles } from './components/upload';
 export * as redux from 'react-redux';
 export * from './components/zipDownloader';
diff --git a/packages/ui-default/components/socket/index.ts b/packages/ui-default/components/socket/index.ts
index 4c66c0437..850379c04 100644
--- a/packages/ui-default/components/socket/index.ts
+++ b/packages/ui-default/components/socket/index.ts
@@ -1,23 +1,31 @@
 import ReconnectingWebSocket from 'reconnecting-websocket';
 
 export default class Sock {
-  private sock: ReconnectingWebSocket;
+  sock: ReconnectingWebSocket;
+  interval: NodeJS.Timer;
+
   constructor(public url: string) {
     const i = new URL(url, window.location.href);
     if (i.host !== window.location.host) i.searchParams.append('sid', document.cookie.split('sid=')[1].split(';')[0]);
     i.protocol = i.protocol.replace('http', 'ws');
     this.url = i.toString();
-    this.sock = new ReconnectingWebSocket(this.url);
+    this.sock = new ReconnectingWebSocket(this.url, [], {
+      maxReconnectionDelay: 10000,
+      maxRetries: 100,
+    });
     this.sock.onopen = () => {
       console.log('Connected');
       this.onopen?.(this.sock);
+      this.interval = setInterval(() => this.sock.send('ping'), 30000);
     };
     this.sock.onclose = ({ code, reason }) => {
       console.warn('Connection closed, ', code, reason);
       if (code >= 4000) this.close();
+      if (this.interval) clearInterval(this.interval);
       this.onclose?.(code, reason);
     };
     this.sock.onmessage = (message) => {
+      if (message.data === 'pong') return;
       if (process.env.NODE_ENV !== 'production') console.log('Sock.onmessage: ', message);
       const msg = JSON.parse(message.data);
       if (msg.error === 'PermissionError' || msg.error === 'PrivilegeError') this.close();