299 lines
8.0 KiB
JavaScript
299 lines
8.0 KiB
JavaScript
'use strict'
|
|
|
|
const crypto = require('crypto')
|
|
const PriorityQueue = require('priorityqueuejs')
|
|
|
|
// Compute a one-way hash of the input string.
|
|
function sanitizeToken(id) {
|
|
return crypto
|
|
.createHash('sha256')
|
|
.update(id, 'utf-8')
|
|
.digest('hex')
|
|
}
|
|
|
|
function getUtcEpochSeconds() {
|
|
return (Date.now() / 1000) >>> 0
|
|
}
|
|
|
|
// Encapsulate a rate-limited token, with a user-provided ID and user-provided data.
|
|
//
|
|
// Each token has a notion of the number of uses remaining until exhausted,
|
|
// and the next reset time, when it can be used again even if it's exhausted.
|
|
class Token {
|
|
constructor(id, data, usesRemaining, nextReset) {
|
|
// Use underscores to avoid conflict with property accessors.
|
|
Object.assign(this, {
|
|
_id: id,
|
|
_data: data,
|
|
_usesRemaining: usesRemaining,
|
|
_nextReset: nextReset,
|
|
_isValid: true,
|
|
_isFrozen: false,
|
|
})
|
|
}
|
|
|
|
get id() {
|
|
return this._id
|
|
}
|
|
get data() {
|
|
return this._data
|
|
}
|
|
get usesRemaining() {
|
|
return this._usesRemaining
|
|
}
|
|
get nextReset() {
|
|
return this._nextReset
|
|
}
|
|
get isValid() {
|
|
return this._isValid
|
|
}
|
|
get isFrozen() {
|
|
return this._isFrozen
|
|
}
|
|
|
|
get hasReset() {
|
|
return getUtcEpochSeconds() >= this.nextReset
|
|
}
|
|
|
|
get isExhausted() {
|
|
return this.usesRemaining <= 0 && !this.hasReset
|
|
}
|
|
|
|
// Update the uses remaining and next reset time for a token.
|
|
update(usesRemaining, nextReset) {
|
|
if (!Number.isInteger(usesRemaining)) {
|
|
throw Error('usesRemaining must be an integer')
|
|
}
|
|
if (!Number.isInteger(nextReset)) {
|
|
throw Error('nextReset must be an integer')
|
|
}
|
|
|
|
if (this._isFrozen) {
|
|
return
|
|
}
|
|
|
|
// Since the token pool will typically return the same token for many uses
|
|
// before moving on to another, `update()` may be called many times. Since
|
|
// the sequence of responses may be indeterminate, keep the "worst case"
|
|
// value for uses remaining.
|
|
if (
|
|
this._nextReset === this.constructor.nextResetNever ||
|
|
nextReset > this._nextReset
|
|
) {
|
|
this._nextReset = nextReset
|
|
this._usesRemaining = usesRemaining
|
|
} else if (nextReset === this._nextReset) {
|
|
this._usesRemaining = Math.min(this._usesRemaining, usesRemaining)
|
|
} else {
|
|
// Discard the new update; it's older than the values we have.
|
|
}
|
|
}
|
|
|
|
// Indicate that the token should no longer be used.
|
|
invalidate() {
|
|
this._isValid = false
|
|
}
|
|
|
|
// Freeze the uses remaining and next reset values. Helpful for keeping
|
|
// stable ordering for a valid priority queue.
|
|
freeze() {
|
|
this._isFrozen = true
|
|
}
|
|
|
|
// Unfreeze the uses remaining and next reset values.
|
|
unfreeze() {
|
|
this._isFrozen = false
|
|
}
|
|
|
|
getDebugInfo({ sanitize = true } = {}) {
|
|
const { id, data, usesRemaining, nextReset, isValid, isFrozen } = this
|
|
|
|
if (sanitize) {
|
|
return {
|
|
id: sanitizeToken(id),
|
|
data: '[redacted]',
|
|
usesRemaining,
|
|
nextReset,
|
|
isValid,
|
|
isFrozen,
|
|
}
|
|
} else {
|
|
return { id, data, usesRemaining, nextReset, isValid, isFrozen }
|
|
}
|
|
}
|
|
}
|
|
|
|
// Large sentinel value which means "never reset".
|
|
Token.nextResetNever = Number.MAX_SAFE_INTEGER
|
|
|
|
// Encapsulate a collection of rate-limited tokens and choose the next
|
|
// available token when one is needed.
|
|
//
|
|
// Designed for the Github API, though may be also useful with other rate-
|
|
// limited APIs.
|
|
class TokenPool {
|
|
// batchSize: The maximum number of times we use each token before moving
|
|
// on to the next one.
|
|
constructor({ batchSize = 1 } = {}) {
|
|
this.batchSize = batchSize
|
|
|
|
this.currentBatch = { currentToken: null, remaining: 0 }
|
|
|
|
// A set of IDs used for deduplication.
|
|
this.tokenIds = new Set()
|
|
|
|
// See discussion on the FIFO and priority queues in `next()`.
|
|
this.fifoQueue = []
|
|
this.priorityQueue = new PriorityQueue(this.constructor.compareTokens)
|
|
}
|
|
|
|
// Use the token whose current rate allotment is expiring soonest.
|
|
static compareTokens(first, second) {
|
|
return second.nextReset - first.nextReset
|
|
}
|
|
|
|
// Add a token with user-provided ID and data.
|
|
//
|
|
// The ID can be a primitive value or an object reference, and is used (with
|
|
// `Set`) for deduplication. If a token already exists with a given id, it
|
|
// will be ignored.
|
|
add(id, data, usesRemaining, nextReset) {
|
|
if (this.tokenIds.has(id)) {
|
|
return false
|
|
}
|
|
this.tokenIds.add(id)
|
|
|
|
usesRemaining = usesRemaining === undefined ? this.batchSize : usesRemaining
|
|
nextReset = nextReset === undefined ? Token.nextResetNever : nextReset
|
|
|
|
const token = new Token(id, data, usesRemaining, nextReset)
|
|
this.fifoQueue.push(token)
|
|
|
|
return true
|
|
}
|
|
|
|
// Prepare to start a new batch by obtaining and returning the next usable
|
|
// token.
|
|
_nextBatch() {
|
|
let next
|
|
|
|
while ((next = this.fifoQueue.shift())) {
|
|
if (!next.isValid) {
|
|
// Discard, and
|
|
continue
|
|
} else if (next.isExhausted) {
|
|
next.freeze()
|
|
this.priorityQueue.enq(next)
|
|
} else {
|
|
return next
|
|
}
|
|
}
|
|
|
|
while (
|
|
!this.priorityQueue.isEmpty() &&
|
|
(next = this.priorityQueue.peek())
|
|
) {
|
|
if (!next.isValid) {
|
|
// Discard, and
|
|
continue
|
|
} else if (next.isExhausted) {
|
|
// No need to check any more tokens, since they all reset after this
|
|
// one.
|
|
break
|
|
} else {
|
|
this.priorityQueue.deq() // deq next
|
|
next.unfreeze()
|
|
return next
|
|
}
|
|
}
|
|
|
|
throw Error('Token pool is exhausted')
|
|
}
|
|
|
|
// Obtain the next available token, returning `null` if no tokens are
|
|
// available.
|
|
//
|
|
// Tokens are initially pulled from a FIFO queue. The first token is used
|
|
// for a batch of requests, then returned to the queue to give those
|
|
// requests the opportunity to complete. The next token is used for the next
|
|
// batch of requests.
|
|
//
|
|
// This strategy allows a token to be used for concurrent requests, not just
|
|
// sequential request, and simplifies token recovery after errored and timed
|
|
// out requests.
|
|
//
|
|
// By the time the original token re-emerges, its requests should have long
|
|
// since completed. Even if a couple them are still running, they can
|
|
// reasonably be ignored. The uses remaining won't be 100% correct, but
|
|
// that's fine, because Shields uses only 75%
|
|
//
|
|
// The process continues until an exhausted token is pulled from the FIFO
|
|
// queue. At that time it's placed in the priority queue based on its
|
|
// scheduled reset time. To ensure the priority queue works as intended,
|
|
// the scheduled reset time is frozen then.
|
|
//
|
|
// After obtaining a token using `next()`, invoke `update()` on it to set a
|
|
// new use-remaining count and next-reset time. Invoke `invalidate()` to
|
|
// indicate it should not be reused.
|
|
next() {
|
|
let token = this.currentBatch.token
|
|
const remaining = this.currentBatch.remaining
|
|
|
|
if (remaining <= 0 || !token.isValid || token.isExhausted) {
|
|
token = this._nextBatch()
|
|
this.currentBatch = {
|
|
token,
|
|
remaining: token.hasReset
|
|
? this.batchSize
|
|
: Math.min(this.batchSize, token.usesRemaining),
|
|
}
|
|
this.fifoQueue.push(token)
|
|
}
|
|
|
|
this.currentBatch.remaining -= 1
|
|
|
|
return token
|
|
}
|
|
|
|
// Iterate over all valid tokens.
|
|
forEach(callback) {
|
|
function visit(item) {
|
|
if (item.isValid) {
|
|
callback(item)
|
|
}
|
|
}
|
|
|
|
this.fifoQueue.forEach(visit)
|
|
this.priorityQueue.forEach(visit)
|
|
}
|
|
|
|
allValidTokenIds() {
|
|
const result = []
|
|
this.forEach(({ id }) => result.push(id))
|
|
return result
|
|
}
|
|
|
|
serializeDebugInfo({ sanitize = true } = {}) {
|
|
const maybeSanitize = sanitize ? id => sanitizeToken(id) : id => id
|
|
|
|
const priorityQueue = []
|
|
this.priorityQueue.forEach(t =>
|
|
priorityQueue.push(t.getDebugInfo({ sanitize }))
|
|
)
|
|
|
|
return {
|
|
utcEpochSeconds: getUtcEpochSeconds(),
|
|
allValidTokenIds: this.allValidTokenIds().map(maybeSanitize),
|
|
fifoQueue: this.fifoQueue.map(t => t.getDebugInfo({ sanitize })),
|
|
priorityQueue,
|
|
sanitized: sanitize,
|
|
}
|
|
}
|
|
}
|
|
|
|
module.exports = {
|
|
sanitizeToken,
|
|
Token,
|
|
TokenPool,
|
|
}
|