misty.cafe/mistykey
1
0
Fork 0
Code Issues Pull requests 1 Projects Releases Packages Wiki Activity Actions

merge: Enable NestJS shutdown hooks for clean exit (!1145)

Browse source 
View MR for information: https://activitypub.software/TransFem-org/Sharkey/-/merge_requests/1145

Approved-by: dakkar <dakkar@thenautilus.net>
Approved-by: Marie <github@yuugi.dev>
This commit is contained in:
Hazelnoot 2025-06-29 19:23:55 +00:00
commit be1063238f
9 changed files with 173 additions and 52 deletions
Download patch file Download diff file Expand all files Collapse all files

48
packages/backend/src/server/api/StreamingApiServerService.ts
View file

@ -4,7 +4,7 @@
*/
import { EventEmitter } from 'events';
import { Inject, Injectable } from '@nestjs/common';
import { Inject, Injectable, OnApplicationShutdown } from '@nestjs/common';
import * as Redis from 'ioredis';
import * as WebSocket from 'ws';
import proxyAddr from 'proxy-addr';
@ -32,11 +32,12 @@ import type * as http from 'node:http';
const MAX_CONNECTIONS_PER_CLIENT = 32;
@Injectable()
export class StreamingApiServerService {
export class StreamingApiServerService implements OnApplicationShutdown {
#wss: WebSocket.WebSocketServer;
#connections = new Map<WebSocket.WebSocket, number>();
#connectionsByClient = new Map<string, Set<WebSocket.WebSocket>>(); // key: IP / user ID -> value: connection
#cleanConnectionsIntervalId: NodeJS.Timeout | null = null;
readonly #globalEv = new EventEmitter();
constructor(
@Inject(DI.redisForSub)
@ -67,6 +68,14 @@ export class StreamingApiServerService {
@Inject(DI.config)
private config: Config,
) {
this.redisForSub.on('message', this.onRedis);
}
@bindThis
onApplicationShutdown() {
this.redisForSub.off('message', this.onRedis);
this.#globalEv.removeAllListeners();
// Other shutdown logic is handled by detach(), which gets called by ServerServer's own shutdown handler.
}
@bindThis
@ -79,6 +88,12 @@ export class StreamingApiServerService {
return rateLimit.blocked;
}
@bindThis
private onRedis(_: string, data: string) {
const parsed = JSON.parse(data);
this.#globalEv.emit('message', parsed);
}
@bindThis
public attach(server: http.Server): void {
this.#wss = new WebSocket.WebSocketServer({
@ -213,13 +228,6 @@ export class StreamingApiServerService {
});
});
const globalEv = new EventEmitter();
this.redisForSub.on('message', (_: string, data: string) => {
const parsed = JSON.parse(data);
globalEv.emit('message', parsed);
});
this.#wss.on('connection', async (connection: WebSocket.WebSocket, request: http.IncomingMessage, ctx: {
stream: MainStreamConnection,
user: MiLocalUser | null;
@ -233,12 +241,13 @@ export class StreamingApiServerService {
ev.emit(data.channel, data.message);
}
globalEv.on('message', onRedisMessage);
this.#globalEv.on('message', onRedisMessage);
await stream.listen(ev, connection);
this.#connections.set(connection, Date.now());
// TODO use collapsed queue
const userUpdateIntervalId = user ? setInterval(() => {
this.usersService.updateLastActiveDate(user);
}, 1000 * 60 * 5) : null;
@ -249,7 +258,7 @@ export class StreamingApiServerService {
connection.once('close', () => {
ev.removeAllListeners();
stream.dispose();
globalEv.off('message', onRedisMessage);
this.#globalEv.off('message', onRedisMessage);
this.#connections.delete(connection);
if (userUpdateIntervalId) clearInterval(userUpdateIntervalId);
});
@ -274,13 +283,24 @@ export class StreamingApiServerService {
}
@bindThis
public detach(): Promise<void> {
public async detach(): Promise<void> {
if (this.#cleanConnectionsIntervalId) {
clearInterval(this.#cleanConnectionsIntervalId);
this.#cleanConnectionsIntervalId = null;
}
return new Promise((resolve) => {
this.#wss.close(() => resolve());
for (const connection of this.#connections.keys()) {
connection.close();
}
this.#connections.clear();
this.#connectionsByClient.clear();
await new Promise<void>((resolve, reject) => {
this.#wss.close(err => {
if (err) reject(err);
else resolve();
});
});
}
}