Skip to content

Commit 090cb00

Browse files
author
Taois
committed
feat: ws
1 parent 3746817 commit 090cb00

File tree

2 files changed

+216
-154
lines changed

2 files changed

+216
-154
lines changed

‎controllers/websocket.js‎

Lines changed: 112 additions & 108 deletions
Original file line numberDiff line numberDiff line change
@@ -10,16 +10,24 @@ import {toBeijingTime} from "../utils/datetime-format.js";
1010
// WebSocket 客户端管理
1111
const wsClients = new Set();
1212

13+
// 需要拦截的console方法列表
14+
const CONSOLE_METHODS = [
15+
'log', 'error', 'warn', 'info', 'debug',
16+
'time', 'timeEnd', 'timeLog',
17+
'assert', 'clear', 'count', 'countReset',
18+
'dir', 'dirxml', 'group', 'groupCollapsed', 'groupEnd',
19+
'table', 'trace', 'profile', 'profileEnd'
20+
];
21+
1322
// 原始 console 方法备份
14-
const originalConsole = {
15-
log: console.log,
16-
error: console.error,
17-
warn: console.warn,
18-
info: console.info,
19-
debug: console.debug,
20-
time: console.time,
21-
timeEnd: console.timeEnd,
22-
};
23+
const originalConsole = {};
24+
25+
// 动态备份所有console方法
26+
CONSOLE_METHODS.forEach(method => {
27+
if (typeof console[method] === 'function') {
28+
originalConsole[method] = console[method];
29+
}
30+
});
2331

2432
// 广播消息到所有 WebSocket 客户端
2533
function broadcastToClients(message) {
@@ -66,24 +74,22 @@ function interceptConsole() {
6674
};
6775
};
6876

69-
console.log = createInterceptor('log', originalConsole.log);
70-
console.error = createInterceptor('error', originalConsole.error);
71-
console.warn = createInterceptor('warn', originalConsole.warn);
72-
console.info = createInterceptor('info', originalConsole.info);
73-
console.debug = createInterceptor('debug', originalConsole.debug);
74-
console.time = createInterceptor('time', originalConsole.time);
75-
console.timeEnd = createInterceptor('timeEnd', originalConsole.timeEnd);
77+
// 动态拦截所有console方法
78+
CONSOLE_METHODS.forEach(method => {
79+
if (originalConsole[method]) {
80+
console[method] = createInterceptor(method, originalConsole[method]);
81+
}
82+
});
7683
}
7784

7885
// 恢复原始 console
7986
function restoreConsole() {
80-
console.log = originalConsole.log;
81-
console.error = originalConsole.error;
82-
console.warn = originalConsole.warn;
83-
console.info = originalConsole.info;
84-
console.debug = originalConsole.debug;
85-
console.time = originalConsole.time;
86-
console.timeEnd = originalConsole.timeEnd;
87+
// 动态恢复所有console方法
88+
CONSOLE_METHODS.forEach(method => {
89+
if (originalConsole[method]) {
90+
console[method] = originalConsole[method];
91+
}
92+
});
8793
}
8894

8995
// 启动 console 拦截
@@ -96,98 +102,96 @@ interceptConsole();
96102
* @param {Function} done - 完成回调
97103
*/
98104
export default (fastify, options, done) => {
99-
// 注册WebSocket路由
100-
// fastify.register(async function (fastify) {
101-
/**
102-
* WebSocket连接路由
103-
* GET /ws - 建立WebSocket连接
104-
*/
105-
fastify.get('/ws', {websocket: true}, (socket, req) => {
106-
const clientId = Date.now() + Math.random();
107-
originalConsole.log(`WebSocket client connected: ${clientId}`);
108-
originalConsole.log('Socket type:', typeof socket);
109-
originalConsole.log('Socket has send method:', typeof socket.send);
110-
111-
// 添加到客户端集合
112-
wsClients.add(socket);
113-
114-
// 设置连接属性
115-
socket.clientId = clientId;
116-
socket.isAlive = true;
117-
socket.lastPing = Date.now();
118-
119-
// 发送欢迎消息 - 先检查send方法是否存在
120-
if (typeof socket.send === 'function') {
121-
socket.send(JSON.stringify({
122-
type: 'welcome',
123-
message: 'WebSocket connection established',
124-
clientId: clientId,
125-
timestamp: Date.now()
126-
}));
127-
} else {
128-
originalConsole.error('Socket does not have send method');
105+
/**
106+
* WebSocket连接路由
107+
* GET /ws - 建立WebSocket连接
108+
*/
109+
fastify.get('/ws', {websocket: true}, (socket, req) => {
110+
const clientId = Date.now() + Math.random();
111+
originalConsole.log(`WebSocket client connected: ${clientId}`);
112+
originalConsole.log('Socket type:', typeof socket);
113+
originalConsole.log('Socket has send method:', typeof socket.send);
114+
115+
// 添加到客户端集合
116+
wsClients.add(socket);
117+
118+
// 设置连接属性
119+
socket.clientId = clientId;
120+
socket.isAlive = true;
121+
socket.lastPing = Date.now();
122+
123+
// 发送欢迎消息 - 先检查send方法是否存在
124+
if (typeof socket.send === 'function') {
125+
socket.send(JSON.stringify({
126+
type: 'welcome',
127+
message: 'WebSocket connection established',
128+
clientId: clientId,
129+
timestamp: Date.now()
130+
}));
131+
} else {
132+
originalConsole.error('Socket does not have send method');
133+
}
134+
135+
// 设置心跳检测
136+
const heartbeatInterval = setInterval(() => {
137+
if (!socket.isAlive) {
138+
originalConsole.log(`Client ${clientId} failed heartbeat, terminating`);
139+
clearInterval(heartbeatInterval);
140+
wsClients.delete(socket); // 修复内存泄露：从客户端集合中移除
141+
socket.terminate();
142+
return;
129143
}
130144

131-
// 设置心跳检测
132-
const heartbeatInterval = setInterval(() => {
133-
if (!socket.isAlive) {
134-
originalConsole.log(`Client ${clientId} failed heartbeat, terminating`);
135-
clearInterval(heartbeatInterval);
136-
socket.terminate();
137-
return;
138-
}
145+
socket.isAlive = false;
146+
socket.ping();
147+
}, 30000); // 30秒心跳检测
148+
149+
// 处理pong响应
150+
socket.on('pong', () => {
151+
socket.isAlive = true;
152+
});
139153

140-
socket.isAlive = false;
141-
socket.ping();
142-
}, 30000); // 30秒心跳检测
143-
144-
// 处理pong响应
145-
socket.on('pong', () => {
146-
socket.isAlive = true;
147-
});
148-
149-
// 处理消息
150-
socket.on('message', (message) => {
151-
try {
152-
const data = JSON.parse(message.toString());
153-
originalConsole.log(`Received from ${clientId}:`, data);
154-
155-
// 回显消息
156-
if (socket.readyState === socket.OPEN) {
157-
socket.send(JSON.stringify({
158-
type: 'echo',
159-
originalMessage: data,
160-
timestamp: Date.now(),
161-
clientId: clientId
162-
}));
163-
}
164-
} catch (error) {
165-
originalConsole.error('Error processing message:', error);
166-
if (socket.readyState === socket.OPEN) {
167-
socket.send(JSON.stringify({
168-
type: 'error',
169-
message: 'Invalid JSON format',
170-
timestamp: Date.now()
171-
}));
172-
}
154+
// 处理消息
155+
socket.on('message', (message) => {
156+
try {
157+
const data = JSON.parse(message.toString());
158+
originalConsole.log(`Received from ${clientId}:`, data);
159+
160+
// 回显消息
161+
if (socket.readyState === socket.OPEN) {
162+
socket.send(JSON.stringify({
163+
type: 'echo',
164+
originalMessage: data,
165+
timestamp: Date.now(),
166+
clientId: clientId
167+
}));
168+
}
169+
} catch (error) {
170+
originalConsole.error('Error processing message:', error);
171+
if (socket.readyState === socket.OPEN) {
172+
socket.send(JSON.stringify({
173+
type: 'error',
174+
message: 'Invalid JSON format',
175+
timestamp: Date.now()
176+
}));
173177
}
174-
});
178+
}
179+
});
175180

176-
// 处理连接关闭
177-
socket.on('close', (code, reason) => {
178-
originalConsole.log(`Client ${clientId} disconnected: ${code} ${reason}`);
179-
wsClients.delete(socket);
180-
clearInterval(heartbeatInterval);
181-
});
181+
// 处理连接关闭
182+
socket.on('close', (code, reason) => {
183+
originalConsole.log(`Client ${clientId} disconnected: ${code} ${reason}`);
184+
wsClients.delete(socket);
185+
clearInterval(heartbeatInterval);
186+
});
182187

183-
// 处理错误
184-
socket.on('error', (error) => {
185-
originalConsole.error(`WebSocket error for client ${clientId}:`, error);
186-
wsClients.delete(socket);
187-
clearInterval(heartbeatInterval);
188-
});
188+
// 处理错误
189+
socket.on('error', (error) => {
190+
originalConsole.error(`WebSocket error for client ${clientId}:`, error);
191+
wsClients.delete(socket);
192+
clearInterval(heartbeatInterval);
189193
});
190-
// });
194+
});
191195

192196
/**
193197
* WebSocket状态查询接口

0 commit comments

Comments
 (0)