import { EventEmitter } from "events" ; import { IConsumer , IQueuedMessage , IPartition , IPartitionConfig , IPartitionWithEpoch , IPartitionLambdaFactory , ILogger , LambdaCloseType , IContextErrorData , } from "@fluidframework/server-services-core" ; import { Lumberjack } from "@fluidframework/server-services-telemetry" ; import { Partition } from "./partition" ; /**
 * The PartitionManager is responsible for maintaining a list of partitions for the given Kafka topic.
 * It will route incoming messages to the appropriate partition for the messages.
 */ export class PartitionManager extends EventEmitter { private readonly partitions = new Map < number , Partition > ( ) ; private isRebalancing = true ; private stopped = false ; constructor ( private readonly factory , private readonly consumer , private readonly logger ? ) { super ( ) ; this . consumer . on ( "data" , ( message ) => { this . process ( message ) ; } ) ; this . consumer . on ( "rebalancing" , ( partitions ) => { this . rebalancing ( partitions ) ; } ) ; this . consumer . on ( "rebalanced" , ( partitions ) => { this . rebalanced ( partitions ) ; } ) ; this . consumer . on ( "error" , ( error , errorData ) => { if ( this . stopped ) { return ; } this . emit ( "error" , error , errorData ) ; } ) ; } public async stop ( ) { this . stopped = true ; this . logger ?. info ( "Stop requested" ) ; Lumberjack . info ( "Stop requested" ) ; const partitionsStoppedP = [ ] ; for ( const [ , partition ] of this . partitions ) { const stopP = partition . drain ( ) ; partitionsStoppedP . push ( stopP ) ; } await Promise . all ( partitionsStoppedP ) ; for ( const [ , partition ] of this . partitions ) { partition . close ( LambdaCloseType . Stop ) ; } this . partitions . clear ( ) ; this . removeAllListeners ( ) ; } private process ( message ) { if ( this . stopped ) { return ; } if ( this . isRebalancing ) { this . logger ?. info ( `Ignoring ${ message . topic }:${ message . partition }@${ message . offset } due to pending rebalance` ) ; Lumberjack . info ( `Ignoring ${ message . topic }:${ message . partition }@${ message . offset } due to pending rebalance` ) ; return ; } const partition = this . partitions . get ( message . partition ) ; if ( ! partition ) { this . emit ( "error" , `Received message for untracked partition ${ message . topic }:${ message . partition }@${ message . offset }` ) ; return ; } partition . process ( message ) ; } param partitions private rebalancing ( partitions ) { this . logger ?. info ( `Rebalancing partitions: ${ JSON . stringify ( partitions ) }` ) ; Lumberjack . info ( `Rebalancing partitions: ${ JSON . stringify ( partitions ) }` ) ; this . isRebalancing = true ; for ( const [ id , partition ] of this . partitions ) { this . logger ?. info ( `Closing partition ${ id } due to rebalancing` ) ; Lumberjack . info ( `Closing partition ${ id } due to rebalancing` ) ; partition . close ( LambdaCloseType . Rebalance ) ; } this . partitions . clear ( ) ; } param partitions private rebalanced ( partitions ) { if ( this . stopped ) { return ; } this . isRebalancing = false ; const partitionsMap = new Map ( partitions . map ( ( partition ) => [ partition . partition , partition ] ) ) ; const existingPartitions = Array . from ( this . partitions ) ; for ( const [ id , partition ] of existingPartitions ) { if ( ! partitionsMap . has ( id ) ) { this . logger ?. info ( `Closing partition ${ id } due to rebalancing` ) ; Lumberjack . info ( `Closing partition ${ id } due to rebalancing` ) ; partition . close ( LambdaCloseType . Rebalance ) ; this . partitions . delete ( id ) ; } } for ( const partition of partitions ) { if ( this . partitions . has ( partition . partition ) ) { continue ; } this . logger ?. info ( `Creating ${ partition . topic }: Partition ${ partition . partition }, Epoch ${ partition . leaderEpoch }, Offset ${ partition . offset } due to rebalance` ) ; Lumberjack . info ( `Creating ${ partition . topic }: Partition ${ partition . partition }, Epoch ${ partition . leaderEpoch }, Offset ${ partition . offset } due to rebalance` ) ; const newPartition = new Partition ( partition . partition , partition . leaderEpoch , this . factory , this . consumer , this . logger ) ; newPartition . on ( "error" , ( error , errorData ) => { if ( this . stopped ) { return ; } this . emit ( "error" , error , errorData ) ; } ) ; this . partitions . set ( partition . partition , newPartition ) ; } } }