1+ import { randomUUID } from 'crypto' ;
12import type { Response } from 'express' ;
2- import logger from '../logger.js' ;
3+ import logger , { requestContext } from '../logger.js' ;
34import { isRedisAvailable , getPublisher , getSubscriber } from '../lib/redis.js' ;
45
56const HEARTBEAT_INTERVAL_MS = 30_000 ;
67const MAX_WRITABLE_BUFFER = 64 * 1024 ;
78const MAX_CONNECTIONS_PER_IP = 5 ;
9+ const MAX_CONNECTIONS_PER_USER = 10 ;
810const RETRY_AFTER_SECONDS = 60 ;
911
1012interface SSEClient {
@@ -13,6 +15,7 @@ interface SSEClient {
1315 subscriptions : Set < string > ;
1416 paused : boolean ;
1517 ip : string ;
18+ userId ?: string ;
1619}
1720
1821interface SSECapacityCheckResult {
@@ -27,8 +30,17 @@ export class SSEService {
2730 private heartbeatTimer : ReturnType < typeof setInterval > | null = null ;
2831 private slowClientsDropped = 0 ;
2932 private readonly ipConnectionCounts : Map < string , number > = new Map ( ) ;
33+ private readonly userConnectionCounts : Map < string , number > = new Map ( ) ;
3034 private shuttingDown = false ;
3135 private perIpPeakConnections = 0 ;
36+ private perUserPeakConnections = 0 ;
37+
38+ /**
39+ * Stable id attached to every log line emitted by the heartbeat
40+ * setInterval callback, since it fires outside of any HTTP request and
41+ * would otherwise have no requestContext (and thus no correlation id).
42+ */
43+ private readonly heartbeatWorkerId = `sse-heartbeat:${ randomUUID ( ) } ` ;
3244
3345 private readonly maxConnections : number = ( ( ) => {
3446 const parsed = Number . parseInt ( process . env . MAX_SSE_CONNECTIONS ?? '10000' , 10 ) ;
@@ -61,7 +73,7 @@ export class SSEService {
6173 logger . info ( '[SSEService] Redis pub/sub subscription active.' ) ;
6274 }
6375
64- checkCapacity ( ip : string ) : SSECapacityCheckResult {
76+ checkCapacity ( ip : string , userId ?: string ) : SSECapacityCheckResult {
6577 if ( this . clients . size >= this . maxConnections ) {
6678 return {
6779 allowed : false ,
@@ -80,25 +92,53 @@ export class SSEService {
8092 } ;
8193 }
8294
95+ // Independent of the per-IP cap: bounds how many concurrent SSE
96+ // subscriptions a single authenticated user can hold regardless of which
97+ // IP(s) they connect from (e.g. multiple tabs/devices behind different NATs).
98+ if ( userId ) {
99+ const currentUserConnections = this . userConnectionCounts . get ( userId ) ?? 0 ;
100+ if ( currentUserConnections >= MAX_CONNECTIONS_PER_USER ) {
101+ return {
102+ allowed : false ,
103+ status : 429 ,
104+ retryAfterSeconds : RETRY_AFTER_SECONDS ,
105+ message : `Too many concurrent SSE connections for this user. Max ${ MAX_CONNECTIONS_PER_USER } .` ,
106+ } ;
107+ }
108+ }
109+
83110 return { allowed : true } ;
84111 }
85112
86- addClient ( clientId : string , res : Response , subscriptions : string [ ] = [ ] , ip = 'unknown' ) : void {
113+ addClient (
114+ clientId : string ,
115+ res : Response ,
116+ subscriptions : string [ ] = [ ] ,
117+ ip = 'unknown' ,
118+ userId ?: string ,
119+ ) : void {
87120 const nextIpCount = ( this . ipConnectionCounts . get ( ip ) ?? 0 ) + 1 ;
88121 this . ipConnectionCounts . set ( ip , nextIpCount ) ;
89122 this . perIpPeakConnections = Math . max ( this . perIpPeakConnections , nextIpCount ) ;
90123
124+ if ( userId ) {
125+ const nextUserCount = ( this . userConnectionCounts . get ( userId ) ?? 0 ) + 1 ;
126+ this . userConnectionCounts . set ( userId , nextUserCount ) ;
127+ this . perUserPeakConnections = Math . max ( this . perUserPeakConnections , nextUserCount ) ;
128+ }
129+
91130 const client : SSEClient = {
92131 id : clientId ,
93132 res,
94133 subscriptions : new Set ( subscriptions ) ,
95134 paused : false ,
96135 ip,
136+ ...( userId !== undefined && { userId } ) ,
97137 } ;
98138
99139 this . clients . set ( clientId , client ) ;
100140 logger . info (
101- `[SSEService] Connection opened: ${ clientId } , ip: ${ ip } , subscriptions: ${ subscriptions . join ( ', ' ) } `
141+ `[SSEService] Connection opened: ${ clientId } , ip: ${ ip } , userId: ${ userId ?? 'n/a' } , subscriptions: ${ subscriptions . join ( ', ' ) } `
102142 ) ;
103143
104144 res . on ( 'close' , ( ) => {
@@ -194,6 +234,18 @@ export class SSEService {
194234 return this . ipConnectionCounts . size ;
195235 }
196236
237+ getPerUserPeakConnections ( ) : number {
238+ return this . perUserPeakConnections ;
239+ }
240+
241+ getActiveUserCount ( ) : number {
242+ return this . userConnectionCounts . size ;
243+ }
244+
245+ getUserConnectionCount ( userId : string ) : number {
246+ return this . userConnectionCounts . get ( userId ) ?? 0 ;
247+ }
248+
197249 stopHeartbeat ( ) : void {
198250 if ( this . heartbeatTimer ) {
199251 clearInterval ( this . heartbeatTimer ) ;
@@ -207,7 +259,9 @@ export class SSEService {
207259 }
208260
209261 this . heartbeatTimer = setInterval ( ( ) => {
210- this . sendHeartbeat ( ) ;
262+ requestContext . run ( { requestId : this . heartbeatWorkerId } , ( ) => {
263+ this . sendHeartbeat ( ) ;
264+ } ) ;
211265 } , HEARTBEAT_INTERVAL_MS ) ;
212266 }
213267
@@ -235,6 +289,15 @@ export class SSEService {
235289 this . ipConnectionCounts . set ( client . ip , currentIpCount - 1 ) ;
236290 }
237291
292+ if ( client . userId ) {
293+ const currentUserCount = this . userConnectionCounts . get ( client . userId ) ?? 0 ;
294+ if ( currentUserCount <= 1 ) {
295+ this . userConnectionCounts . delete ( client . userId ) ;
296+ } else {
297+ this . userConnectionCounts . set ( client . userId , currentUserCount - 1 ) ;
298+ }
299+ }
300+
238301 try {
239302 if ( ! client . res . writableEnded ) {
240303 client . res . end ( ) ;
0 commit comments