import * as AsyncLock from 'async-lock' ; import * as EventEmitter from 'events' ; import { Collection , MongoClient } from 'mongodb' ; import * as assert from 'assert' ; import * as _ from 'lodash' ; export interface TxConfig < Message , Payload , State = any > { dbHost : string ; dbPort : number ; dbUri ? : string ; dbName : string ; reinitDbOnStartup : boolean ; sendCommitToCohort : ( msg : Message ) => Promise < void > ; sendAbortToCohort : ( msg : Message ) => Promise < void > ; sendVoteToCohort : ( msg : Message ) => Promise < Vote < Payload > > ; /**
     * `causedAbort` is `true` if this is the client that caused the abort.
     * `payload` is `undefined` if the client never got to the point of voting.
     */ sendAbortToClient : ( causedAbort : boolean , msg ? : Message , payload ? : Payload , state ? : State , index ? : number ) => void ; sendToClient : ( payload : Payload , state ? : State , index ? : number ) => void ; onError : ( error : Error , msg : Message , state ? : State , index ? : number ) => void ; } export interface Vote < Payload > { result : 'yes' | 'no' ; payload : Payload ; } type TxState = ( 'voting' | 'committing' | 'committed' | 'aborting' | 'aborted' ) ; type CohortState = TxState | 'waitingForCompletion' ; interface Cohort < Message , Payload > { id : string ; msg ? : Message ; vote ? : Vote < Payload > | undefined ; state ? : CohortState ; } interface TxDoc < Message , Payload > { id : string ; state : TxState ; cohorts : Cohort < Message , Payload > [ ] ; startedOn : Date ; } interface Transition { newTxState ? : 'committing' | 'aborting' ; newCohortState ? : 'waitingForCompletion' ; } const TX_TIMEOUT_SECONDS = 10 ; const MS_IN_S = 1000 ; function txMsg ( msg , txId ) { return `[txId: ${ txId }] ${ msg }` ; } function cohortMsg ( msg , cohortId ) { return `(cohortId: ${ cohortId }) ${ msg }` ; } function txCohortMsg ( msg , txId , cohortId ) { return txMsg ( cohortMsg ( msg , cohortId ) , txId ) ; } function log ( msg , txId , cohortId ? ) { const msgWithCohort = cohortId === undefined ? msg : cohortMsg ( msg , cohortId ) ; console . log ( txMsg ( msgWithCohort , txId ) ) ; } /**
 *  Coordinates between a set of participants on whether to commit or abort
 *  a transaction.
 */ export class TxCoordinator < Message , Payload , State = any > { private lock = new AsyncLock ( ) ; private completed = new EventEmitter ( ) ; private txs : Collection < TxDoc < Message , Payload > > | undefined ; constructor ( private config ) { } async start ( ) { this . txs = await this . getTxsCollection ( ) ; setInterval ( this . timeoutAbort . bind ( this ) , TX_TIMEOUT_SECONDS * MS_IN_S ) ; } async processMessage ( txId , cohortId , cohorts , msg , state ? , index ? ) { return this . doProcessMessage ( txId , cohortId , cohorts , msg , state , index ) . catch ( ( e ) => { this . config . onError ( e , msg , state ) ; } ) ; } private async doProcessMessage ( txId , cohortId , cohorts , msg , state ? , index ? ) { if ( ! this . txs ) { throw new Error ( 'TxCoordinator hasn\'t been started yet: call start()' ) ; } log ( `Processing msg ${ JSON . stringify ( msg ) }` , txId , cohortId ) ; const tx = await this . getTx ( txId , cohorts ) ; const cohortIds = _ . map ( tx . cohorts , 'id' ) ; if ( ! _ . includes ( cohortIds , cohortId ) ) { throw new Error ( txMsg ( `${ cohortId } is not part of this tx. Cohorts are ` + JSON . stringify ( cohortIds ) , txId ) ) ; } if ( this . shouldAbort ( tx ) ) { this . config . sendAbortToClient ( false , msg , undefined , state , index ) ; return ; } const update = await this . txs . updateOne ( { id : txId , cohorts : { $elemMatch : { id : cohortId , msg : { $exists : false } } } } , { $set : { 'cohorts.$.msg' : msg , 'cohorts.$.vote' : undefined , 'cohorts.$.state' : 'voting' } } ) ; if ( update . matchedCount === 0 ) { throw new Error ( txMsg ( `Duplicate message from ${ cohortId }` , txId ) ) ; } let vote ; try { log ( `Sending vote to cohort` , txId , cohortId ) ; vote = await this . config . sendVoteToCohort ( msg ) ; } catch ( e ) { throw new Error ( txCohortMsg ( `Sending vote to cohort failed: ${ e }` , txId , cohortId ) ) ; } log ( `Received vote ${ JSON . stringify ( vote ) }` , txId , cohortId ) ; this . saveVote ( txId , cohortId , vote ) ; const transition = await this . lock . acquire ( txId , ( ) => { return this . processVote ( txId , cohortId , vote , ( ) => { log ( 'Not waiting anymore (tx committed). ' + `Send payload to client of cohort ${ cohortId }` , txId ) ; this . config . sendToClient ( vote . payload , state , index ) ; } , ( ) => { log ( 'Not waiting anymore (tx aborted). ' + `Send payload to client of cohort ${ cohortId }` , txId ) ; this . config . sendAbortToClient ( false , msg , vote . payload , state , index ) ; } ) ; } ) ; log ( `Processing tx transition caused by cohort ${ cohortId }` , txId ) ; let ret ; if ( _ . isEmpty ( transition ) ) { log ( `Tx aborted. Send abort to client of cohort ${ cohortId }` , txId ) ; this . config . sendAbortToClient ( false , msg , vote . payload , state , index ) ; ret = this . completeMessage ( txId , cohortId , msg , false ) ; } else if ( transition . newTxState === 'committing' && ! transition . newCohortState ) { log ( `Tx committed. Send payload to client of cohort ${ cohortId }` , txId ) ; ret = Promise . all ( [ this . completeTx ( txId , true ) , this . completeMessage ( txId , cohortId , msg , true ) ] ) . then ( ( ) => { this . config . sendToClient ( vote . payload , state , index ) ; this . completed . emit ( txId + '-commit' ) ; } ) ; } else if ( ! transition . newTxState && transition . newCohortState === 'waitingForCompletion' ) { log ( `Tx pending. Cohort ${ cohortId } is waiting for completion` , txId ) ; } else if ( transition . newTxState === 'aborting' ) { log ( `Tx is aborting. Send abort to client of cohort ${ cohortId }` , txId ) ; this . config . sendAbortToClient ( true , msg , vote . payload , state , index ) ; this . completed . emit ( txId + '-abort' ) ; ret = Promise . all ( [ this . completeTx ( txId , false ) , this . completeMessage ( txId , cohortId , msg , false ) ] ) ; } else { throw new Error ( 'Invalid transition' ) ; } return ret ; } private async getTx ( txId , cohorts ) { const cohortObjects = _ . map ( cohorts , ( id ) => ( { id } ) ) ; try { return ( await this . txs ! . findOneAndUpdate ( { id : txId } , { $setOnInsert : { state : 'voting' , startedOn : new Date ( ) , cohorts : cohortObjects } } , { returnOriginal : false , upsert : true } ) ) . value ! ; } catch ( e ) { log ( `We got ${ e }, this is not an error (it is expected to ` + 'happen in some cases)' , txId ) ; return ( await this . txs ! . findOne ( { id : txId } ) ) ! ; } } private async processVote ( txId , cohortId , vote , onCommit , onAbort ) { const tx = ( await this . txs ! . findOne ( { id : txId } ) ) ! ; const ret = { } ; if ( this . shouldAbort ( tx ) ) { return ret ; } if ( vote . result === 'yes' ) { if ( this . allVotedButOne ( tx ) ) { await this . updateTxState ( txId , 'committing' ) ; ret . newTxState = 'committing' ; } else { this . completed . on ( txId + '-commit' , onCommit ) ; this . completed . on ( txId + '-abort' , onAbort ) ; await this . updateCohortState ( txId , cohortId , 'waitingForCompletion' ) ; ret . newCohortState = 'waitingForCompletion' ; } } else { if ( vote . result !== 'no' ) { log ( `Got a vote back that was not 'yes'/'no'. The concept ` + 'doesn\'t correctly implement voting. We got back ' + JSON . stringify ( vote ) , txId , cohortId ) ; } await this . updateTxState ( txId , 'aborting' ) ; ret . newTxState = 'aborting' ; } return ret ; } private shouldAbort ( tx ) { return tx . state === 'aborting' || tx . state === 'aborted' ; } private allVotedButOne ( tx ) { return _ . chain ( tx . cohorts ) . map ( 'state' ) . filter ( ( s ) => s === 'waitingForCompletion' ) . value ( ) . length === tx . cohorts . length - 1 ; } private saveVote ( txId , cohortId , vote ) { return this . txs ! . updateOne ( { id : txId , 'cohorts.id' : cohortId } , { $set : { [ `cohorts.$.vote` ] : vote } } ) ; } private async updateCohortState ( txId , cohortId , newState ) { const update = await this . txs ! . updateOne ( { id : txId , 'cohorts.id' : cohortId } , { $set : { [ `cohorts.$.state` ] : newState } } ) ; if ( update . modifiedCount === 0 ) { throw new Error ( txCohortMsg ( `Couldn't set state to ${ newState } ` + `(matched: ${ update . matchedCount })` , txId , cohortId ) ) ; } } private updateTxState ( id , newState ) { return this . txs ! . updateOne ( { id : id } , { $set : { state : newState } } ) ; } /**
     *  Every time this function runs, it aborts active txs whose
     *  `startedOn + TX_TIMEOUT_SECONDS` time is less than the current time.
     */ private async timeoutAbort ( ) { const threshold = new Date ( ) ; threshold . setSeconds ( threshold . getSeconds ( ) - TX_TIMEOUT_SECONDS ) ; const txsToAbort = await this . txs ! . find < { id : string ; } > ( { state : 'voting' , startedOn : { $lt : threshold } } , { projection : { id : 1 } } ) . toArray ( ) ; console . log ( `Found ${ txsToAbort . length } transactions to abort due to timeout` ) ; await Promise . all ( _ . map ( txsToAbort , async ( txToAbort ) => { return this . lock . acquire ( txToAbort . id , async ( ) : Promise < boolean > => { const tx = ( await this . txs ! . findOne ( { id : txToAbort . id } , { projection : { state : 1 } } ) ) ! ; if ( tx . state !== 'voting' ) { return false ; } log ( 'Aborting due to timeout' , txToAbort . id ) ; await this . updateTxState ( txToAbort . id , 'aborting' ) ; return true ; } ) . then ( ( shouldCompleteTx ) : Promise < any > | undefined => { if ( shouldCompleteTx ) { log ( 'Emitting abort event' , txToAbort . id ) ; this . completed . emit ( txToAbort . id + '-abort' ) ; return this . completeTx ( txToAbort . id , false ) ; } return undefined ; } ) ; } ) ) ; } /**
     * Send a complete ('commit' or 'abort') message to a cohort server.
     *
     * After an ACK is received, the cohort state is updated.
     *
     * Function returns after an ACK is received.
     */ private async completeMessage ( txId , cohortId , msg , success ) { const completedState = success ? 'committed' : 'aborted' ; const completingState = success ? 'committing' : 'aborting' ; const send = success ? this . config . sendCommitToCohort : this . config . sendAbortToCohort ; log ( `Sending complete message to cohort, success: ${ success }` , txId , cohortId ) ; await this . updateCohortState ( txId , cohortId , completingState ) ; try { await send ( msg ) ; } catch ( e ) { log ( e , txId ) ; throw new Error ( txCohortMsg ( `Sending commit/abort failed` , txId , cohortId ) ) ; } await this . updateCohortState ( txId , cohortId , completedState ) ; } private async completeTx ( txId , success ) { const completedState = success ? 'committed' : 'aborted' ; log ( `Completing tx (status: ${ completedState })` , txId ) ; const tx = ( await this . txs ! . findOne ( { id : txId } , { projection : { 'cohorts.msg' : 1 , 'cohorts.id' : 1 , 'cohorts.state' : 1 } } ) ) ! ; await Promise . all ( _ . map ( tx . cohorts , ( cohort ) => { if ( cohort . state !== 'waitingForCompletion' ) { log ( `cohort ${ cohort . id } is not waiting for completion ` + `(it's ${ cohort . state }), not sending complete message` , txId ) ; return undefined ; } log ( `On complete tx, sending complete message to ${ cohort . id } ` + `(state is: ${ cohort . state })` , txId ) ; assert . ok ( ! _ . isNil ( cohort . msg ) , txMsg ( `Expected msg for ${ cohort . id } but got undefined` , txId ) ) ; return this . completeMessage ( txId , cohort . id , cohort . msg ! , success ) ; } ) ) ; return this . updateTxState ( txId , completedState ) ; } private async getTxsCollection ( ) { const mongoUri = _ . isEmpty ( this . config . dbUri ) ? `mongodb://${ this . config . dbHost }:${ this . config . dbPort }` : this . config . dbUri ; const client = await MongoClient . connect ( mongoUri ) ; const db = client . db ( this . config . dbName ) ; if ( this . config . reinitDbOnStartup ) { await db . dropDatabase ( ) ; console . log ( `Reinitialized db ${ this . config . dbName }` ) ; } const txs = db . collection ( 'txs' ) ; await Promise . all ( [ txs . createIndex ( { id : 1 } , { unique : true } ) , txs . createIndex ( { 'cohorts.id' : 1 } ) ] ) ; return txs ; } }