SolidScribe/server/helpers/HostQueue.js

//
// Per-host politeness scheduler. Knows nothing about scraping.
//
// This is the piece that stops the scraper looking like a crawler. The old code
// did foundUrls.forEach(fetch) - a note with 40 links opened 40 simultaneous
// connections, which is the most bot-like traffic shape available. Here, each
// host gets one request at a time with jittered spacing between them, and a
// global cap bounds total concurrency.
//

const config = require('@config/scrape')

let HostQueue = module.exports = {}

const jitter = (min, max) => Math.round(min + Math.random() * Math.max(0, max - min))

class Queue {
	constructor(options = {}) {
		this.globalConcurrency = options.globalConcurrency != null
			? options.globalConcurrency : config.maxConcurrentFetches
		this.perHostConcurrency = options.perHostConcurrency != null
			? options.perHostConcurrency : config.perHostConcurrency
		this.minSpacingMs = options.minSpacingMs != null
			? options.minSpacingMs : config.perHostMinSpacingMs
		this.maxSpacingMs = options.maxSpacingMs != null
			? options.maxSpacingMs : config.perHostMaxSpacingMs
		this.now = options.now || (() => Date.now())
		// Injectable so tests with a controlled clock don't spend real
		// wall-time waiting out politeness spacing.
		this.setTimeoutFn = options.setTimeoutFn || ((fn, ms) => setTimeout(fn, ms))
		this.clearTimeoutFn = options.clearTimeoutFn || ((handle) => clearTimeout(handle))

		this.pending = []            // FIFO
		this.hosts = new Map()       // hostKey -> { running, nextAllowedAt }
		this.running = 0
		this.seq = 0
		this.timer = null
		this.stopped = false
		this.drainWaiters = []
	}

	hostState(hostKey) {
		if (!this.hosts.has(hostKey)) {
			this.hosts.set(hostKey, { running: 0, nextAllowedAt: 0 })
		}
		return this.hosts.get(hostKey)
	}

	//
	// Queue a task for a host. Resolves/rejects with the task's own outcome.
	// spacingMs overrides the jitter floor (robots.txt Crawl-delay).
	//
	schedule(hostKey, taskFn, options = {}) {
		if (this.stopped) {
			return Promise.reject(new Error('HostQueue stopped'))
		}

		return new Promise((resolve, reject) => {
			this.pending.push({
				hostKey: hostKey || '(none)',
				taskFn,
				resolve,
				reject,
				spacingMs: options.spacingMs != null ? options.spacingMs : null,
				seq: this.seq++,
			})
			this.pump()
		})
	}

	//
	// Run whatever is eligible. FIFO scan rather than per-host round-robin, so
	// no host can starve another and order stays predictable.
	//
	pump() {
		if (this.stopped) { return }

		this.clearTimer()

		while (this.running < this.globalConcurrency) {
			const at = this.now()
			let index = -1
			let earliest = Infinity

			for (let i = 0; i < this.pending.length; i++) {
				const job = this.pending[i]
				const host = this.hostState(job.hostKey)

				if (host.running >= this.perHostConcurrency) { continue }

				if (host.nextAllowedAt > at) {
					earliest = Math.min(earliest, host.nextAllowedAt)
					continue
				}

				index = i
				break
			}

			if (index === -1) {
				// Nothing runnable now. Wake up when the earliest host is due.
				if (this.pending.length && earliest !== Infinity) {
					this.armTimer(earliest - at)
				}
				break
			}

			this.run(this.pending.splice(index, 1)[0])
		}

		this.checkDrained()
	}

	run(job) {
		const host = this.hostState(job.hostKey)
		host.running++
		this.running++

		const settle = () => {
			host.running--
			this.running--

			// Space the NEXT request to this host from the end of this one.
			const spacing = job.spacingMs != null
				? job.spacingMs
				: jitter(this.minSpacingMs, this.maxSpacingMs)
			host.nextAllowedAt = this.now() + spacing

			this.pump()
		}

		// A throwing taskFn must not wedge the pump, so the call itself is
		// wrapped rather than only its promise.
		let result = null
		try {
			result = Promise.resolve(job.taskFn())
		} catch (error) {
			result = Promise.reject(error)
		}

		result.then(
			value => { settle(); job.resolve(value) },
			error => { settle(); job.reject(error) }
		)
	}

	armTimer(delay) {
		this.clearTimer()
		this.timer = this.setTimeoutFn(() => {
			this.timer = null
			this.pump()
		}, Math.max(1, delay))

		//
		// Deliberately NOT unref'd.
		//
		// This timer is the only thing that will resume a host which is inside
		// its politeness gap. unref'ing it means that if nothing else holds the
		// event loop open, Node exits and every queued job is silently
		// abandoned with its promise never settling. Inside the server that is
		// masked by the listening socket, but it makes the queue unusable from
		// a standalone script and is a footgun either way.
		//
		// Shutdown is handled explicitly instead: stop() clears this timer and
		// rejects whatever is still queued, and it is wired to SIGINT/SIGTERM
		// in server/index.js.
		//
	}

	clearTimer() {
		if (this.timer) {
			this.clearTimeoutFn(this.timer)
			this.timer = null
		}
	}

	size() { return this.pending.length }
	inflight() { return this.running }

	stats() {
		return {
			pending: this.pending.length,
			running: this.running,
			hosts: this.hosts.size,
		}
	}

	checkDrained() {
		if (this.pending.length || this.running) { return }

		const waiters = this.drainWaiters
		this.drainWaiters = []
		waiters.forEach(resolve => resolve())
	}

	// Resolves once the queue is empty and nothing is in flight.
	drain() {
		if (!this.pending.length && !this.running) { return Promise.resolve() }
		return new Promise(resolve => this.drainWaiters.push(resolve))
	}

	//
	// Reject everything still queued. In-flight tasks are not cancelled here -
	// their own AbortControllers handle that.
	//
	stop() {
		this.stopped = true
		this.clearTimer()

		const pending = this.pending
		this.pending = []
		pending.forEach(job => job.reject(new Error('HostQueue stopped')))

		this.checkDrained()
	}
}

HostQueue.Queue = Queue

// Shared instance used by the scrape worker. Tests build their own.
HostQueue.shared = new Queue()