@railzai/bottleneck
Version:
Distributed task scheduler and rate limiter
299 lines (254 loc) • 10.3 kB
text/coffeescript
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...) ->
options, invalid
parser.load options, , @
= new Queues NUM_PRIORITIES
= {}
= new States ["RECEIVED", "QUEUED", "RUNNING", "EXECUTING"].concat(if then ["DONE"] else [])
= null
= new Events @
= new Sync "submit",
= new Sync "register",
storeOptions = parser.load options, , {}
= if == "redis" or == "ioredis" or ?
storeInstanceOptions = parser.load options, , {}
new RedisDatastore @, storeOptions, storeInstanceOptions
else if == "local"
storeInstanceOptions = parser.load options, , {}
new LocalDatastore @, storeOptions, storeInstanceOptions
else
throw new Bottleneck::BottleneckError "Invalid datastore type: #{@datastore}"
.on "leftzero", => .heartbeat?.ref?()
.on "zero", => .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: -> .ready
clients: -> .clients
channel: -> "b_#{@id}"
channel_client: -> "b_#{@id}_#{@_store.clientId}"
publish: (message) -> .__publish__ message
disconnect: (flush=true) -> .__disconnect__ flush
chain: ( ) -> @
queued: (priority) -> .queued priority
clusterQueued: -> .__queued__()
empty: -> == 0 and .isEmpty()
running: -> .__running__()
done: -> .__done__()
jobStatus: (id) -> .jobStatus id
jobs: (status) -> .statusJobs status
counts: -> .statusCounts()
_randomIndex: -> Math.random().toString(36).slice(2)
check: (weight=1) -> .__check__ weight
_clearGlobalState: (index) ->
if [index]?
clearTimeout [index].expiration
delete [index]
true
else false
_free: (index, job, options, eventInfo) ->
try
{ running } = await .__free__ index, options.weight
.trigger "debug", "Freed #{options.id}", eventInfo
if running == 0 and then .trigger "idle"
catch e
.trigger "error", e
_run: (index, job, wait) ->
job.doRun()
clearGlobalState = .bind @, index
run = .bind @, index, job
free = .bind @, index, job
[index] =
timeout: setTimeout =>
job.doExecute , clearGlobalState, run, free
, wait
expiration: if job.options.expiration? then setTimeout ->
job.doExpire clearGlobalState, run, free
, wait + job.options.expiration
job: job
_drainOne: (capacity) ->
.schedule =>
if == 0 then return .resolve null
queue = .getFirst()
{ options, args } = next = queue.first()
if capacity? and options.weight > capacity then return .resolve null
.trigger "debug", "Draining #{options.id}", { args, options }
index =
.__register__ index, options.weight, options.expiration
.then ({ success, wait, reservoir }) =>
.trigger "debug", "Drained #{options.id}", { success, args, options }
if success
queue.shift()
empty =
if empty then .trigger "empty"
if reservoir == 0 then .trigger "depleted", empty
index, next, wait
.resolve options.weight
else
.resolve null
_drainAll: (capacity, total=0) ->
.then (drained) =>
if drained?
newCapacity = if capacity? then capacity - drained else capacity
else .resolve total
.catch (e) => .trigger "error", e
_dropAllQueued: (message) -> .shiftAll (job) -> job.doDrop { message }
stop: (options={}) ->
options = parser.load options,
waitForExecuting = (at) =>
finished = =>
counts = .counts
(counts[0] + counts[1] + counts[2] + counts[3]) == at
new (resolve, reject) =>
if finished() then resolve()
else
"done", =>
if finished()
"done"
resolve()
done = if options.dropWaitingJobs
= (index, next) -> next.doDrop { message: options.dropErrorMessage }
= => .resolve null
.schedule => .schedule =>
for k, v of
if == "RUNNING"
clearTimeout v.timeout
clearTimeout v.expiration
v.job.doDrop { message: options.dropErrorMessage }
options.dropErrorMessage
waitForExecuting(0)
else
{ priority: NUM_PRIORITIES - 1, weight: 0 }, => waitForExecuting(1)
= (job) -> job._reject new Bottleneck::BottleneckError options.enqueueErrorMessage
= => .reject new Bottleneck::BottleneckError "stop() has already been called"
done
_addToQueue: (job) =>
{ args, options } = job
try
{ reachedHWM, blocked, strategy } = await .__submit__ , options.weight
catch error
.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 .shiftLastFrom(options.priority)
else if strategy == Bottleneck::strategy.OVERFLOW_PRIORITY then .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
.push job
await
reachedHWM
_receive: (job) ->
if .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()
.schedule , job
submit: (args...) ->
if typeof args[0] == "function"
[fn, args..., cb] = args
options = parser.load {},
else
[options, fn, args..., cb] = args
options = parser.load options,
task = (args...) =>
new (resolve, reject) ->
fn args..., (args...) ->
(if args[0]? then reject else resolve) args
job = new Job task, args, options, , , , ,
job.promise
.then (args) -> cb? args...
.catch (args) -> if Array.isArray args then cb? args... else cb? args
job
schedule: (args...) ->
if typeof args[0] == "function"
[task, args...] = args
options = {}
else
[options, task, args...] = args
job = new Job task, args, options, , , , ,
job
job.promise
wrap: (fn) ->
schedule = .bind @
wrapped = (args...) -> schedule fn.bind(@), args...
wrapped.withOptions = (options, args...) -> schedule options, fn, args...
wrapped
updateSettings: (options={}) ->
await .__updateSettings__ parser.overwrite options,
parser.overwrite options, , @
@
currentReservoir: -> .__currentReservoir__()
incrementReservoir: (incr=0) -> .__incrementReservoir__ incr
module.exports = Bottleneck