import { Injectable , OnModuleDestroy } from '@nestjs/common' ; import { each } from 'async' ; import { serialize } from 'class-transformer' ; import { Response } from 'express' ; import { RedisService } from 'nestjs-redis' ; import { ERROR_MESSAGES } from '@koh/common' ; import * as Sentry from '@sentry/node' ; /**
 * A connection to a particular frontend client
 */ interface Connection { res : Response ; cleanup : ( ) => Promise < void > ; } interface RedisClientInfo < T > { clientId : number ; metadata : T ; } /**
 * T is metadata associated with each Client
 *
 * Low level abstraction for sending SSE to "rooms" of clients.
 * Probably don't use this directly, and wrap it in a service specific to that event source
 *
 * This handles when there's multiple backend instances by assigning unique client ids to each connection.
 * When one instance wants to send to a client, it publishes to a Redis channel for the client.
 * All instances listen to Redis, and if they are the one managing that client, they send the msg.
 *
 * Rooms with client metadata are also maintained in Redis key/value store.
 */ @ Injectable ( ) export class SSEService < T > implements OnModuleDestroy { private directConnnections : Record < string , Connection > = { } ; constructor ( private readonly redisService ) { const redisSub = this . redisService . getClient ( 'sub' ) ; if ( ! redisSub ) { Sentry . captureException ( ERROR_MESSAGES . sseService . getSubClient ) ; throw new Error ( ERROR_MESSAGES . sseService . getSubClient ) ; } redisSub . on ( 'message' , ( channel , message ) => { const id = /sse::client-(\d+)/ . exec ( channel ) ; if ( id && id [ 1 ] in this . directConnnections ) { this . directConnnections [ id [ 1 ] ] . res . write ( `data: ${ message }\n\n` ) ; } } ) ; } async onModuleDestroy ( ) { await each ( Object . values ( this . directConnnections ) , async ( conn ) => { await conn . cleanup ( ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . cleanupConnection ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; } ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . moduleDestroy ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; } /**
     * Get redis channel name from client id
     */ private idToChannel ( clientId ) { return `sse::client-${ clientId }` ; } /** Add a client to a room */ async subscribeClient ( room , res , metadata ) { const redisSub = this . redisService . getClient ( 'sub' ) ; const redis = this . redisService . getClient ( 'db' ) ; if ( ! redisSub ) { Sentry . captureException ( ERROR_MESSAGES . sseService . getSubClient ) ; throw new Error ( ERROR_MESSAGES . sseService . getSubClient ) ; } if ( ! redis ) { Sentry . captureException ( ERROR_MESSAGES . sseService . getDBClient ) ; throw new Error ( ERROR_MESSAGES . sseService . getDBClient ) ; } const clientId = await redis . incr ( 'sse::client::id' ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . clientIdSubscribe ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; if ( ! clientId ) { Sentry . captureException ( ERROR_MESSAGES . sseService . clientIdNotFound ) ; throw new Error ( ERROR_MESSAGES . sseService . clientIdNotFound ) ; } await redisSub . subscribe ( this . idToChannel ( clientId ) ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . subscribe ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; const clientInfo = JSON . stringify ( { clientId , metadata : metadata , } as RedisClientInfo < T > ) ; await redis . sadd ( room , clientInfo ) . catch ( ( err ) => { console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; this . directConnnections [ clientId ] = { res , cleanup : async ( ) => { await redis . srem ( room , clientInfo ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . removeFromRoom ) ; console . error ( err ) ; } ) ; await redisSub . unsubscribe ( this . idToChannel ( clientId ) ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . unsubscribe ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; res . end ( ) ; } , } ; res . write ( '\n' ) ; res . socket . on ( 'end' , async ( ) => { await this . directConnnections [ clientId ] . cleanup ( ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . directConnections ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; delete this . directConnnections [ clientId ] ; } ) ; } /** Send some data to everyone in a room */ async sendEvent < D > ( room , payload ) { const redisPub = this . redisService . getClient ( 'pub' ) ; const redis = this . redisService . getClient ( 'db' ) ; if ( ! redisPub ) { Sentry . captureException ( ERROR_MESSAGES . sseService . getPubClient ) ; throw new Error ( ERROR_MESSAGES . sseService . getPubClient ) ; } if ( ! redis ) { Sentry . captureException ( ERROR_MESSAGES . sseService . getDBClient ) ; throw new Error ( ERROR_MESSAGES . sseService . getDBClient ) ; } const roomInfo = await redis . smembers ( room ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . roomMembers ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; if ( room && roomInfo ) { const clients = roomInfo . map ( ( s ) => JSON . parse ( s ) ) ; console . log ( `sending sse to ${ clients . length } clients in ${ room }` ) ; console . time ( `sending sse time: ` ) ; await each ( clients , async ( { clientId , metadata } ) => { const toSend = serialize ( await payload ( metadata ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . serialize ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ) ; await redisPub . publish ( this . idToChannel ( clientId ) , toSend ) . catch ( ( err ) => { console . error ( ERROR_MESSAGES . sseService . publish ) ; console . error ( err ) ; Sentry . captureException ( err ) ; } ) ; } ) ; console . timeEnd ( `sending sse time: ` ) ; } } }