UNPKG

@railzai/bottleneck

Version:

Distributed task scheduler and rate limiter

299 lines (254 loc) 10.3 kB
NUM_PRIORITIES = 10 DEFAULT_PRIORITY = 5 parser = require "./parser" Queues = require "./Queues" Job = require "./Job" LocalDatastore = require "./LocalDatastore" RedisDatastore = require "./RedisDatastore" Events = require "./Events" States = require "./States" Sync = require "./Sync" class Bottleneck Bottleneck.default = Bottleneck Bottleneck.Events = Events Bottleneck.version = Bottleneck::version = require("./version.json").version Bottleneck.strategy = Bottleneck::strategy = { LEAK:1, OVERFLOW:2, OVERFLOW_PRIORITY:4, BLOCK:3 } Bottleneck.BottleneckError = Bottleneck::BottleneckError = require "./BottleneckError" Bottleneck.Group = Bottleneck::Group = require "./Group" Bottleneck.RedisConnection = Bottleneck::RedisConnection = require "./RedisConnection" Bottleneck.IORedisConnection = Bottleneck::IORedisConnection = require "./IORedisConnection" Bottleneck.Batcher = Bottleneck::Batcher = require "./Batcher" jobDefaults: priority: DEFAULT_PRIORITY weight: 1 expiration: null id: "<no-id>" storeDefaults: maxConcurrent: null minTime: 0 highWater: null strategy: Bottleneck::strategy.LEAK penalty: null reservoir: null reservoirRefreshInterval: null reservoirRefreshAmount: null reservoirIncreaseInterval: null reservoirIncreaseAmount: null reservoirIncreaseMaximum: null localStoreDefaults: Promise: Promise timeout: null heartbeatInterval: 250 redisStoreDefaults: Promise: Promise timeout: null heartbeatInterval: 5000 clientTimeout: 10000 Redis: null clientOptions: {} clusterNodes: null clearDatastore: false connection: null instanceDefaults: datastore: "local" connection: null id: "<no-id>" rejectOnDrop: true trackDoneStatus: false Promise: Promise stopDefaults: enqueueErrorMessage: "This limiter has been stopped and cannot accept new jobs." dropWaitingJobs: true dropErrorMessage: "This limiter has been stopped." constructor: (options={}, invalid...) -> @_validateOptions options, invalid parser.load options, @instanceDefaults, @ @_queues = new Queues NUM_PRIORITIES @_scheduled = {} @_states = new States ["RECEIVED", "QUEUED", "RUNNING", "EXECUTING"].concat(if @trackDoneStatus then ["DONE"] else []) @_limiter = null @Events = new Events @ @_submitLock = new Sync "submit", @Promise @_registerLock = new Sync "register", @Promise storeOptions = parser.load options, @storeDefaults, {} @_store = if @datastore == "redis" or @datastore == "ioredis" or @connection? storeInstanceOptions = parser.load options, @redisStoreDefaults, {} new RedisDatastore @, storeOptions, storeInstanceOptions else if @datastore == "local" storeInstanceOptions = parser.load options, @localStoreDefaults, {} new LocalDatastore @, storeOptions, storeInstanceOptions else throw new Bottleneck::BottleneckError "Invalid datastore type: #{@datastore}" @_queues.on "leftzero", => @_store.heartbeat?.ref?() @_queues.on "zero", => @_store.heartbeat?.unref?() _validateOptions: (options, invalid) -> unless options? and typeof options == "object" and invalid.length == 0 throw new Bottleneck::BottleneckError "Bottleneck v2 takes a single object argument. Refer to https://github.com/SGrondin/bottleneck#upgrading-to-v2 if you're upgrading from Bottleneck v1." ready: -> @_store.ready clients: -> @_store.clients channel: -> "b_#{@id}" channel_client: -> "b_#{@id}_#{@_store.clientId}" publish: (message) -> @_store.__publish__ message disconnect: (flush=true) -> @_store.__disconnect__ flush chain: (@_limiter) -> @ queued: (priority) -> @_queues.queued priority clusterQueued: -> @_store.__queued__() empty: -> @queued() == 0 and @_submitLock.isEmpty() running: -> @_store.__running__() done: -> @_store.__done__() jobStatus: (id) -> @_states.jobStatus id jobs: (status) -> @_states.statusJobs status counts: -> @_states.statusCounts() _randomIndex: -> Math.random().toString(36).slice(2) check: (weight=1) -> @_store.__check__ weight _clearGlobalState: (index) -> if @_scheduled[index]? clearTimeout @_scheduled[index].expiration delete @_scheduled[index] true else false _free: (index, job, options, eventInfo) -> try { running } = await @_store.__free__ index, options.weight @Events.trigger "debug", "Freed #{options.id}", eventInfo if running == 0 and @empty() then @Events.trigger "idle" catch e @Events.trigger "error", e _run: (index, job, wait) -> job.doRun() clearGlobalState = @_clearGlobalState.bind @, index run = @_run.bind @, index, job free = @_free.bind @, index, job @_scheduled[index] = timeout: setTimeout => job.doExecute @_limiter, clearGlobalState, run, free , wait expiration: if job.options.expiration? then setTimeout -> job.doExpire clearGlobalState, run, free , wait + job.options.expiration job: job _drainOne: (capacity) -> @_registerLock.schedule => if @queued() == 0 then return @Promise.resolve null queue = @_queues.getFirst() { options, args } = next = queue.first() if capacity? and options.weight > capacity then return @Promise.resolve null @Events.trigger "debug", "Draining #{options.id}", { args, options } index = @_randomIndex() @_store.__register__ index, options.weight, options.expiration .then ({ success, wait, reservoir }) => @Events.trigger "debug", "Drained #{options.id}", { success, args, options } if success queue.shift() empty = @empty() if empty then @Events.trigger "empty" if reservoir == 0 then @Events.trigger "depleted", empty @_run index, next, wait @Promise.resolve options.weight else @Promise.resolve null _drainAll: (capacity, total=0) -> @_drainOne(capacity) .then (drained) => if drained? newCapacity = if capacity? then capacity - drained else capacity @_drainAll(newCapacity, total + drained) else @Promise.resolve total .catch (e) => @Events.trigger "error", e _dropAllQueued: (message) -> @_queues.shiftAll (job) -> job.doDrop { message } stop: (options={}) -> options = parser.load options, @stopDefaults waitForExecuting = (at) => finished = => counts = @_states.counts (counts[0] + counts[1] + counts[2] + counts[3]) == at new @Promise (resolve, reject) => if finished() then resolve() else @on "done", => if finished() @removeAllListeners "done" resolve() done = if options.dropWaitingJobs @_run = (index, next) -> next.doDrop { message: options.dropErrorMessage } @_drainOne = => @Promise.resolve null @_registerLock.schedule => @_submitLock.schedule => for k, v of @_scheduled if @jobStatus(v.job.options.id) == "RUNNING" clearTimeout v.timeout clearTimeout v.expiration v.job.doDrop { message: options.dropErrorMessage } @_dropAllQueued options.dropErrorMessage waitForExecuting(0) else @schedule { priority: NUM_PRIORITIES - 1, weight: 0 }, => waitForExecuting(1) @_receive = (job) -> job._reject new Bottleneck::BottleneckError options.enqueueErrorMessage @stop = => @Promise.reject new Bottleneck::BottleneckError "stop() has already been called" done _addToQueue: (job) => { args, options } = job try { reachedHWM, blocked, strategy } = await @_store.__submit__ @queued(), options.weight catch error @Events.trigger "debug", "Could not queue #{options.id}", { args, options, error } job.doDrop { error } return false if blocked job.doDrop() return true else if reachedHWM shifted = if strategy == Bottleneck::strategy.LEAK then @_queues.shiftLastFrom(options.priority) else if strategy == Bottleneck::strategy.OVERFLOW_PRIORITY then @_queues.shiftLastFrom(options.priority + 1) else if strategy == Bottleneck::strategy.OVERFLOW then job if shifted? then shifted.doDrop() if not shifted? or strategy == Bottleneck::strategy.OVERFLOW if not shifted? then job.doDrop() return reachedHWM job.doQueue reachedHWM, blocked @_queues.push job await @_drainAll() reachedHWM _receive: (job) -> if @_states.jobStatus(job.options.id)? job._reject new Bottleneck::BottleneckError "A job with the same id already exists (id=#{job.options.id})" false else job.doReceive() @_submitLock.schedule @_addToQueue, job submit: (args...) -> if typeof args[0] == "function" [fn, args..., cb] = args options = parser.load {}, @jobDefaults else [options, fn, args..., cb] = args options = parser.load options, @jobDefaults task = (args...) => new @Promise (resolve, reject) -> fn args..., (args...) -> (if args[0]? then reject else resolve) args job = new Job task, args, options, @jobDefaults, @rejectOnDrop, @Events, @_states, @Promise job.promise .then (args) -> cb? args... .catch (args) -> if Array.isArray args then cb? args... else cb? args @_receive job schedule: (args...) -> if typeof args[0] == "function" [task, args...] = args options = {} else [options, task, args...] = args job = new Job task, args, options, @jobDefaults, @rejectOnDrop, @Events, @_states, @Promise @_receive job job.promise wrap: (fn) -> schedule = @schedule.bind @ wrapped = (args...) -> schedule fn.bind(@), args... wrapped.withOptions = (options, args...) -> schedule options, fn, args... wrapped updateSettings: (options={}) -> await @_store.__updateSettings__ parser.overwrite options, @storeDefaults parser.overwrite options, @instanceDefaults, @ @ currentReservoir: -> @_store.__currentReservoir__() incrementReservoir: (incr=0) -> @_store.__incrementReservoir__ incr module.exports = Bottleneck