modules/measurement-rate-control/hyper-p2p-bucket-rate-limit/index.js

require('bare-process/global')
const EventEmitter = require('bare-events')
const { setInterval, clearInterval } = require('bare-timers')
const b4a = require('b4a')
const { initModuleSwarm, gossipSend } = require('../../_shared/p2p-bare.js')

const PROTOCOL = 'bucket-rate-limit/v1'

class HyperP2PBucketRateLimit extends EventEmitter {
  constructor (opts = {}) {
    super()
    this._stats = { ops: 0, errors: 0 }

    this.topic = opts.topic || null
    this.keyPair = opts.keyPair || require('hypercore-crypto').keyPair()
    this._buckets = new Map()
    this.rate = opts.rate ?? 10
    this.burst = opts.burst ?? 20
    this.syncIntervalMs = opts.syncIntervalMs ?? 10000
    this.enableBackgroundTimers = opts.enableBackgroundTimers === true
    this._syncTimer = null
    this.swarm = null
    this._peerMsgs = null
  }

  configure ({ rate, burst }) {
    if (rate != null) this.rate = rate
    if (burst != null) this.burst = burst
    const snap = this.snapshot()
    this.emit('configure', snap)
    this._gossip({ type: 'configure', ...snap })
    return snap
  }

  snapshot () {
    return { rate: this.rate, burst: this.burst, buckets: this._bucketSnapshot() }
  }

  _bucketSnapshot () {
    const out = {}
    for (const [peerId, b] of this._buckets) {
      out[peerId] = { tokens: b.tokens, updatedAt: b.updatedAt }
    }
    return out
  }

  _gossip (data) {
    if (this._peerMsgs) gossipSend(this, data)
  }

  _bucket (peerId) {
    const key = typeof peerId === 'string' ? peerId : b4a.toString(peerId, 'hex')
    let b = this._buckets.get(key)
    if (!b) {
      b = { tokens: this.burst, updatedAt: Date.now() }
      this._buckets.set(key, b)
    }
    const now = Date.now()
    const elapsed = (now - b.updatedAt) / 1000
    b.tokens = Math.min(this.burst, b.tokens + elapsed * this.rate)
    b.updatedAt = now
    return b
  }

  getBucket (peerId) {
    const key = typeof peerId === 'string' ? peerId : b4a.toString(peerId, 'hex')
    const b = this._bucket(key)
    return { peerId: key, tokens: b.tokens, rate: this.rate, burst: this.burst, updatedAt: b.updatedAt }
  }

  tryConsume (peerId, cost = 1) {
    const b = this._bucket(peerId)
    if (b.tokens < cost) {
      this.emit('reject', { peerId, cost })
      return false
    }
    b.tokens -= cost
    this.emit('consume', { peerId, cost, tokens: b.tokens })
    this._gossip({ type: 'bucket', peerId: typeof peerId === 'string' ? peerId : b4a.toString(peerId, 'hex'), tokens: b.tokens, updatedAt: b.updatedAt })
    return true
  }

  _applyRemote (data) {
    if (!data) return
    if (data.type === 'configure') {
      if (data.rate != null) this.rate = data.rate
      if (data.burst != null) this.burst = data.burst
      this.emit('remote-configure', { rate: this.rate, burst: this.burst })
    } else if (data.type === 'sync' && data.buckets) {
      for (const [peerId, remote] of Object.entries(data.buckets)) {
        const cur = this._buckets.get(peerId)
        if (!cur || remote.updatedAt > cur.updatedAt) {
          this._buckets.set(peerId, { tokens: remote.tokens, updatedAt: remote.updatedAt })
        }
      }
      this.emit('sync', { peers: Object.keys(data.buckets).length })
    } else if (data.type === 'bucket' && data.peerId) {
      const cur = this._buckets.get(data.peerId)
      if (!cur || data.updatedAt >= cur.updatedAt) {
        this._buckets.set(data.peerId, { tokens: data.tokens, updatedAt: data.updatedAt })
      }
    }
  }

  sync () {
    const snap = { type: 'sync', ...this.snapshot() }
    this._gossip(snap)
    this.emit('sync', snap)
    return snap
  }

  _ensureSyncTimer () {
    if (!this.enableBackgroundTimers || this._syncTimer) return
    this._syncTimer = setInterval(() => this.sync(), this.syncIntervalMs)
  }

  async ready () {
    if (this.swarm || !this.topic) return this
    await initModuleSwarm(this, {
      keyPair: this.keyPair,
      topic: this.topic,
      protocol: PROTOCOL,
      onmessage: (data) => this._applyRemote(data)
    })
    this._ensureSyncTimer()
    return this
  }


  getStats () {
    return { ...this._stats }
  }

  async close () {
    if (this._syncTimer) {
      clearInterval(this._syncTimer)
      this._syncTimer = null
    }
    if (this.swarm) await this.swarm.destroy().catch(() => {})
    this.swarm = null
    this.emit('closed')
  }
}

module.exports = { HyperP2PBucketRateLimit, PROTOCOL }