shields/lib/token-pool.js

'use strict'

const PriorityQueue = require('priorityqueuejs')

// 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
  }

  static utcEpochSeconds() {
    return (Date.now() / 1000) >>> 0
  }

  get hasReset() {
    return this.constructor.utcEpochSeconds() >= 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
  }
}

// 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) {
    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 ((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
  }
}

module.exports = {
  Token,
  TokenPool,
}