import { App } from './../app' ; import { ConsumptionResponse } from './rate-limiter-interface' ; import { LocalRateLimiter } from './local-rate-limiter' ; import { RateLimiterAbstract , RateLimiterClusterMasterPM2 } from 'rate-limiter-flexible' ; import { Server } from '../server' ; const cluster = require ( 'cluster' ) ; const pm2 = require ( 'pm2' ) ; export interface ConsumptionMessage { app : App ; eventKey : string ; points : number ; maxPoints : number ; } export class ClusterRateLimiter extends LocalRateLimiter { /**
     * Initialize the local rate limiter driver.
     */ constructor ( protected server ) { super ( server ) ; if ( cluster . isPrimary || typeof cluster . isPrimary === 'undefined' ) { if ( server . pm2 ) { new RateLimiterClusterMasterPM2 ( pm2 ) ; } else { server . discover . join ( 'rate_limiter:limiters' , ( rateLimiters ) => { this . rateLimiters = Object . fromEntries ( Object . entries ( rateLimiters ) . map ( ( [ key , rateLimiterObject ] ) => { return [ key , this . createNewRateLimiter ( key . split ( ':' ) [ 0 ] , rateLimiterObject . _points ) , ] ; } ) ) ; } ) ; server . discover . join ( 'rate_limiter:consume' , ( { app , eventKey , points , maxPoints } ) => { super . consume ( app , eventKey , points , maxPoints ) ; } ) ; server . discover . on ( 'added' , ( ) => { if ( server . nodes . get ( 'self' ) . isMaster ) { this . sendRateLimiters ( ) ; } } ) ; } } } /**
     * Consume points for a given key, then
     * return a response object with headers and the success indicator.
     */ protected consume ( app , eventKey , points , maxPoints ) { return super . consume ( app , eventKey , points , maxPoints ) . then ( ( response ) => { if ( response . canContinue ) { this . server . discover . send ( 'rate_limiter:consume' , { app , eventKey , points , maxPoints , } ) ; } return response ; } ) ; } /**
     * Clear the rate limiter or active connections.
     */ clear ( closeConnections = false ) { return super . clear ( closeConnections ) . then ( ( ) => { if ( closeConnections && this . server . nodes . get ( 'self' ) . isMaster ) { this . server . discover . demote ( ) ; this . server . discover . send ( 'rate_limiter:limiters' , this . rateLimiters ) ; } } ) ; } /**
     * Send the stored rate limiters this instance currently have.
     */ protected sendRateLimiters ( ) { this . server . discover . send ( 'rate_limiter:limiters' , this . rateLimiters ) ; } }