import { map , debounce , forEach } from 'lodash' ; import PQueue from 'p-queue' ; import { Redis } from 'ioredis' ; import { setIntervalAsync , clearIntervalAsync , SetIntervalAsyncTimer , } from 'set-interval-async/dynamic' ; import debugF from 'debug' ; import { RateLimiterRedis } from 'rate-limiter-flexible' ; import { sleep , createWorkerId } from '../utils/general' ; import { getRetryDelayType , TaskSuccessCb , TaskErrorCb , TaskFailedCb , Handler , } from './handle-task' ; import { ensureConnected , tryIgnore , ensureDisconnected , publish , set , createClientAndLoadLuaScripts , RedisConfig , } from '../utils/redis' ; import { takeTaskBlocking } from './take-task-blocking' ; import { processTask , OnAfterTaskSuccess , OnAfterTaskError , OnAfterTaskFail , } from './process-task' ; import { getWorkerStartedChannel , getWorkerPausedChannel , getWorkerKey , } from '../utils/keys' ; import { serializeEvent } from '../domain/events/serialize-event' ; import { EventType } from '../domain/events/event-type' ; import { WorkerInstance } from '../domain/worker/worker-instance' ; import { serializeWorker } from '../domain/worker/serialize-worker' ; import { Task } from '../domain/tasks/task' ; import { Worker } from '../domain/worker/worker' ; import { markTaskProcessing } from './mark-task-processing' ; import { getQueueRateLimitConfig , QueueRateLimitConfig , } from './get-queue-rate-limit-config' ; import { createListener } from './create-listener' ; const debug = debugF ( 'conveyor-mq:worker' ) ; export type OnBeforeTaskProcessing = ( { taskId } : { taskId : string ; } ) => any ; export type OnAfterTaskProcessing = ( { task } : { task : Task ; } ) => any ; export interface WorkerInput { queue : string ; redisConfig : RedisConfig ; redisClient ? : Redis ; handler : Handler ; concurrency ? : number ; defaultStallTimeout ? : number ; defaultTaskAcknowledgementInterval ? : number ; getRetryDelay ? : getRetryDelayType ; onTaskSuccess ? : TaskSuccessCb ; onTaskError ? : TaskErrorCb ; onTaskFailed ? : TaskFailedCb ; onHandlerError ? : ( error : any ) => any ; onIdle ? : ( ) => any ; idleTimeout ? : number ; onReady ? : ( ) => any ; autoStart ? : boolean ; removeOnSuccess ? : boolean ; removeOnFailed ? : boolean ; hooks ? : { onBeforeTaskProcessing ? : OnBeforeTaskProcessing ; onAfterTaskProcessing ? : OnAfterTaskProcessing ; onAfterTaskSuccess ? : OnAfterTaskSuccess ; onAfterTaskError ? : OnAfterTaskError ; onAfterTaskFail ? : OnAfterTaskFail ; } ; } param queue param redisConfig param redisClient param handler param concurrency param defaultStallTimeout param defaultTaskAcknowledgementInterval param getRetryDelay param onTaskSuccess param onTaskError param onTaskFailed param onHandlerError param onIdle param idleTimeout param onReady param autoStart param removeOnSuccess param removeOnFailed returns export const createWorker = ( { queue , redisConfig , redisClient , handler , concurrency = 10 , defaultStallTimeout = 1000 , defaultTaskAcknowledgementInterval , getRetryDelay , onTaskSuccess , onTaskError , onTaskFailed , onHandlerError , onIdle , idleTimeout = 250 , onReady , autoStart = true , removeOnSuccess = false , removeOnFailed = false , hooks , } ) : Worker => { debug ( 'Starting' ) ; let isPausing = false ; let isPaused = true ; let isShuttingDown = false ; let isShutdown = false ; let upsertInterval ; let rateLimitConfig ; let rateLimiter ; const worker = { id : createWorkerId ( ) , createdAt : new Date ( ) , } ; debug ( 'Starting promise queues' ) ; const takerQueue = new PQueue ( { concurrency , autoStart } ) ; const workerQueue = new PQueue ( { concurrency } ) ; if ( onIdle ) workerQueue . on ( 'idle' , debounce ( onIdle , idleTimeout ) ) ; debug ( 'Creating clients' ) ; const takerClient = createClientAndLoadLuaScripts ( { ... redisConfig , enableReadyCheck : false , } ) ; const workerClient = redisClient || createClientAndLoadLuaScripts ( redisConfig ) ; const isActive = ( ) => ! isPausing && ! isPaused && ! isShuttingDown && ! isShutdown ; const takeAndProcessTask = async ( t ? ) => { debug ( t ? `Processing pre-fetched task ${ t . id }` : `Starting to check for tasks to process` ) ; try { if ( rateLimiter ) { try { await rateLimiter . consume ( queue ) ; } catch ( e ) { if ( e . msBeforeNext ) { await sleep ( e . msBeforeNext ) ; takerQueue . add ( takeAndProcessTask ) ; return ; } } } const task = t || ( await tryIgnore ( async ( ) => { const taskId = await takeTaskBlocking ( { queue , client : takerClient , } ) ; if ( ! taskId ) return null ; if ( hooks ?. onBeforeTaskProcessing ) { await hooks . onBeforeTaskProcessing ( { taskId } ) ; } const processingTask = await markTaskProcessing ( { taskId , queue , stallTimeout : defaultStallTimeout , client : workerClient , } ) ; if ( hooks ?. onAfterTaskProcessing ) { await hooks . onAfterTaskProcessing ( { task : processingTask } ) ; } return processingTask ; } , ( ) => isActive ( ) ) ) ; if ( task ) { debug ( `Adding task to worker queue ${ task . id }` ) ; const nextTask = await workerQueue . add ( async ( ) => tryIgnore ( async ( ) => { debug ( `Processing task ${ task . id }` ) ; const nextTaskToHandle = await processTask ( { task , queue , client : workerClient , handler , stallTimeout : task . stallTimeout || defaultStallTimeout , taskAcknowledgementInterval : task . taskAcknowledgementInterval || defaultTaskAcknowledgementInterval || defaultStallTimeout / 2 , onAcknowledgeTask : ( { task : ackTask } ) => { debug ( `Acknowledging task ${ ackTask }` ) ; } , onAcknowledgedTask : ( { task : ackTask } ) => { debug ( `Acknowledged task ${ ackTask }` ) ; } , getRetryDelay , onTaskSuccess , onTaskError , onTaskFailed , removeOnSuccess : task . removeOnSuccess !== undefined ? task . removeOnSuccess : removeOnSuccess , removeOnFailed : task . removeOnFailed !== undefined ? task . removeOnFailed : removeOnFailed , hooks : { onAfterTaskSuccess : hooks ?. onAfterTaskSuccess , onAfterTaskError : hooks ?. onAfterTaskError , onAfterTaskFail : hooks ?. onAfterTaskFail , } , } ) ; debug ( `Processed task ${ task . id }` ) ; return nextTaskToHandle ; } , ( ) => isActive ( ) ) ) ; if ( isActive ( ) ) { debug ( 'calling takeAndProcessTask with a next task' ) ; takerQueue . add ( ( ) => takeAndProcessTask ( nextTask ) ) ; } } else if ( isActive ( ) ) { debug ( 'calling takeAndProcessTask without a next task' ) ; takerQueue . add ( takeAndProcessTask ) ; } } catch ( e ) { debug ( 'takeAndProcessTask error' ) ; if ( onHandlerError ) onHandlerError ( e ) ; await sleep ( 1000 ) ; if ( isActive ( ) ) { takerQueue . add ( takeAndProcessTask ) ; } } } ; const upsertWorker = async ( ) => { debug ( 'Upserting' ) ; await set ( { key : getWorkerKey ( { workerId : worker . id , queue , } ) , client : workerClient , value : serializeWorker ( worker ) , ttl : 30000 , } ) ; } ; const setQueueRateLimit = ( { points , duration , } ) => { rateLimitConfig = { points , duration } ; rateLimiter = new RateLimiterRedis ( { storeClient : workerClient , points , duration , keyPrefix : queue , inmemoryBlockOnConsumed : points + 1 , } ) ; } ; const clearQueueRateLimit = ( ) => { rateLimitConfig = undefined ; rateLimiter = undefined ; } ; const setupListener = async ( ) => { const listener = createListener ( { queue , redisConfig , events : [ EventType . QueueRateLimitUpdated ] , } ) ; await listener . onReady ( ) ; listener . on ( EventType . QueueRateLimitUpdated , ( { event } ) => { if ( event . data ?. rateLimitConfig ) { setQueueRateLimit ( event . data . rateLimitConfig ) ; } else { clearQueueRateLimit ( ) ; } } ) ; } ; const pause = async ( { killProcessingTasks , } ) => { debug ( 'Pausing' ) ; if ( isShutdown || isShuttingDown ) { throw new Error ( 'Cannot pause a shutdown worker.' ) ; } isPausing = true ; await clearIntervalAsync ( upsertInterval ) ; forEach ( killProcessingTasks ? [ takerQueue , workerQueue ] : [ takerQueue ] , ( q ) => { q . pause ( ) ; q . clear ( ) ; } ) ; await publish ( { channel : getWorkerPausedChannel ( { queue } ) , message : serializeEvent ( { createdAt : new Date ( ) , type : EventType . WorkerPaused , worker , } ) , client : workerClient , } ) ; await Promise . all ( [ ... map ( killProcessingTasks && ! redisClient ? [ takerClient , workerClient ] : [ takerClient ] , ( client ) => ensureDisconnected ( { client } ) ) , ... map ( [ workerQueue , takerQueue ] , ( q ) => q . onIdle ( ) ) , ] ) ; isPaused = true ; isPausing = false ; debug ( 'Paused' ) ; } ; const start = async ( ) => { debug ( 'Starting' ) ; if ( isShuttingDown || isShutdown ) { throw new Error ( 'Cannot start a shutdown worker.' ) ; } if ( isPausing ) { throw new Error ( 'Cannot start a pausing worker.' ) ; } if ( isPaused ) { await Promise . all ( map ( [ takerClient , workerClient ] , ( client ) => ensureConnected ( { client } ) ) ) ; const queueRateLimitConfig = await getQueueRateLimitConfig ( { queue , client : workerClient , } ) ; if ( queueRateLimitConfig && queueRateLimitConfig . duration && queueRateLimitConfig . points ) { setQueueRateLimit ( { points : queueRateLimitConfig . points , duration : queueRateLimitConfig . duration , } ) ; } forEach ( [ workerQueue , takerQueue ] , ( q ) => q . start ( ) ) ; await upsertWorker ( ) ; upsertInterval = setIntervalAsync ( upsertWorker , 15000 ) ; takerQueue . addAll ( map ( Array . from ( { length : concurrency } ) , ( ) => takeAndProcessTask ) ) ; isPaused = false ; await publish ( { channel : getWorkerStartedChannel ( { queue } ) , message : serializeEvent ( { createdAt : new Date ( ) , type : EventType . WorkerStarted , worker , } ) , client : workerClient , } ) ; debug ( 'Started' ) ; } } ; const shutdown = async ( { killProcessingTasks , } ) => { debug ( 'Shutting down' ) ; if ( isShuttingDown || isShutdown ) { throw new Error ( 'Cannot shutdown an already shutdown worker.' ) ; } if ( isPausing ) { throw new Error ( 'Cannot shutdown a pausing worker.' ) ; } isShuttingDown = true ; await clearIntervalAsync ( upsertInterval ) ; forEach ( killProcessingTasks ? [ takerQueue , workerQueue ] : [ takerQueue ] , ( q ) => { q . pause ( ) ; q . clear ( ) ; } ) ; await publish ( { channel : getWorkerStartedChannel ( { queue } ) , message : serializeEvent ( { createdAt : new Date ( ) , type : EventType . WorkerShutdown , worker , } ) , client : workerClient , } ) ; await Promise . all ( [ ... map ( killProcessingTasks && ! redisClient ? [ takerClient , workerClient ] : [ takerClient ] , ( client ) => ensureDisconnected ( { client } ) ) , ... map ( [ workerQueue , takerQueue ] , ( q ) => q . onIdle ( ) ) , ] ) ; isShutdown = true ; isShuttingDown = false ; debug ( 'Shutdown' ) ; } ; const ready = async ( ) => { await setupListener ( ) ; if ( autoStart ) await start ( ) ; if ( onReady ) onReady ( ) ; debug ( 'Ready' ) ; } ; const readyPromise = ready ( ) ; return { id : worker . id , createdAt : worker . createdAt , onReady : async ( ) => { debug ( 'onReady' ) ; await readyPromise ; } , pause : async ( killProcessingTasks ? ) => { await readyPromise ; return pause ( { killProcessingTasks } ) ; } , start : ( ) => start ( ) , getQueueRateLimitConfig : async ( ) => rateLimitConfig , shutdown : async ( killProcessingTasks ? ) => { await readyPromise ; return shutdown ( { killProcessingTasks } ) ; } , onIdle : async ( ) => { debug ( 'onIdle' ) ; await Promise . all ( map ( [ takerQueue , workerQueue ] , ( q ) => q . onIdle ( ) ) ) ; } , } ; } ;