SolidScribe/server/workers/ScrapeQueue.js

//
// The scrape worker.
//
// Owns the whole lifecycle for one URL: cache check -> robots -> polite fetch ->
// parse -> thumbnail -> persist -> notify. Saving a note enqueues URLs and
// returns immediately; nothing in the request path waits on the network.
//
// Structural fixes over the old implementation:
//   - Promise.allSettled semantics instead of a completion counter, so a single
//     rejection can no longer strand a batch forever.
//   - Every fetch goes through HostQueue, so there is no unbounded fan-out.
//   - Jobs for the same URL coalesce into one fetch, fanned out to every
//     attachment row that wanted it.
//   - Pending work is recoverable after a restart from attachment.scrape_state.
//

const config = require('@config/scrape')
const db = require('@config/database')
const UrlSafety = require('@helpers/UrlSafety')
const HttpFetch = require('@helpers/HttpFetch')
const HostQueue = require('@helpers/HostQueue')
const Robots = require('@helpers/Robots')
const Realtime = require('@helpers/Realtime')
const SiteScrape = require('@helpers/SiteScrape')
const ScrapeCache = require('@models/ScrapeCache')

const cheerio = require('cheerio')

let ScrapeQueue = module.exports = {}

// attachment.scrape_state
const STATE = ScrapeQueue.STATE = { PENDING: 0, FETCHING: 1, DONE: 2, FAILED: 3 }

let settings = { ...config }
let queue = HostQueue.shared
let started = false

// urlKey -> { promise, jobs: [] } so concurrent requests for one URL share a fetch.
const inFlight = new Map()
// userId -> pending job count, for backpressure.
const perUser = new Map()
// userId -> debounce timer, so a 10-link note triggers one refresh not ten.
const emitTimers = new Map()
//
// Every job promise currently outstanding. Needed because callers like
// recoverPending() and scanTextForWebsites() deliberately do not await their
// batches - without this, drain() could return while fire-and-forget work was
// still sitting in front of its first await.
//
const outstanding = new Set()

const EMIT_DEBOUNCE_MS = 1500

const now = () => Math.round(Date.now() / 1000)

//
// Test seam. Lets a test run with zero politeness spacing and no robots lookups
// while still exercising the real code path.
//
ScrapeQueue.configure = (overrides = {}) => {
	settings = { ...config, ...overrides }

	queue = new HostQueue.Queue({
		globalConcurrency: settings.maxConcurrentFetches,
		perHostConcurrency: settings.perHostConcurrency,
		minSpacingMs: settings.perHostMinSpacingMs,
		maxSpacingMs: settings.perHostMaxSpacingMs,
	})

	return settings
}

ScrapeQueue.__state = () => ({
	inFlight: inFlight.size,
	perUser: new Map(perUser),
	queue: queue.stats(),
	started,
})

ScrapeQueue.__reset = () => {
	inFlight.clear()
	outstanding.clear()
	perUser.clear()
	emitTimers.forEach(timer => clearTimeout(timer))
	emitTimers.clear()
	settings = { ...config }
	queue = HostQueue.shared
	started = false
}

ScrapeQueue.start = () => {
	started = true
	return ScrapeQueue
}

ScrapeQueue.stop = () => {
	started = false
	queue.stop()
	emitTimers.forEach(timer => clearTimeout(timer))
	emitTimers.clear()
}

//
// Resolves once every queued and in-flight job has settled.
//
// Loops because jobs enqueue follow-up work: a page fetch spawns a thumbnail
// fetch, and a recovery batch is registered asynchronously.
//
ScrapeQueue.drain = async () => {
	for (let pass = 0; pass < 20; pass++) {
		await Promise.allSettled([...outstanding])
		await queue.drain()
		await Promise.allSettled([...inFlight.values()].map(entry => entry.promise))

		if (!outstanding.size && !inFlight.size && !queue.size() && !queue.inflight()) {
			return
		}
	}
}

ScrapeQueue.size = () => queue.size() + inFlight.size + outstanding.size

//
// Debounced notification. The client already refetches on these events, so one
// per user per window is enough - the old code emitted once per URL.
//
const notifyUser = (userId) => {
	if (emitTimers.has(userId)) { return }

	const timer = setTimeout(() => {
		emitTimers.delete(userId)
		Realtime.emitToUser(userId, 'update_counts')
		Realtime.emitToUser(userId, 'update_note_attachments')
	}, EMIT_DEBOUNCE_MS)

	if (timer.unref) { timer.unref() }
	emitTimers.set(userId, timer)
}

// Tests want the emit without waiting out the debounce.
ScrapeQueue.__flushNotifications = () => {
	for (const [userId, timer] of [...emitTimers.entries()]) {
		clearTimeout(timer)
		emitTimers.delete(userId)
		Realtime.emitToUser(userId, 'update_counts')
		Realtime.emitToUser(userId, 'update_note_attachments')
	}
}

const setAttachmentState = async (attachmentId, state) => {
	if (!attachmentId) { return }
	await db.promise().query(
		'UPDATE attachment SET scrape_state = ?, scrape_attempts = scrape_attempts + 1 WHERE id = ?',
		[state, attachmentId]
	)
}

//
// Write a successful result onto an attachment row.
//
// attachment.text stays PLAINTEXT - the FULLTEXT index depends on it.
//
const applyResultToAttachment = async (attachmentId, fields) => {
	if (!attachmentId) { return }

	await db.promise().query(
		`UPDATE attachment
		 SET text = ?, url = ?, last_indexed = ?, file_location = ?, scrape_state = ?,
			 scrape_attempts = scrape_attempts + 1
		 WHERE id = ?`,
		[fields.text, fields.url, now(), fields.fileLocation, STATE.DONE, attachmentId]
	)
}

//
// A failed scrape must not leave a sentinel string in the search index. The old
// code stored 'No site text' / 'Request Timeout' in attachment.text, which then
// got concatenated into the note's searchable content.
//
const applyFailureToAttachment = async (attachmentId, error) => {
	if (!attachmentId) { return }

	const kind = (error && error.kind) || 'unknown'

	// A blocked or permanently-gone URL has no preview to show, ever. Drop the
	// shell row rather than leaving an orphan showing the bare URL.
	const terminal = kind === 'blocked' || kind === 'robots_disallowed' ||
		(kind === 'http' && (error.status === 404 || error.status === 410))

	if (terminal) {
		await db.promise().query('DELETE FROM attachment WHERE id = ? AND scrape_state IN (?, ?)',
			[attachmentId, STATE.PENDING, STATE.FETCHING])
		return
	}

	await setAttachmentState(attachmentId, STATE.FAILED)
}

//
// Fetch and store the preview image. Never fatal: a page with no usable
// thumbnail is still a perfectly good link preview.
//
const fetchThumbnail = async (imageUrl, pageUrl) => {
	if (!imageUrl) { return null }

	// Lazy require: Attachment requires this module, so a top-level require
	// would be a cycle.
	const Attachment = require('@models/Attachment')

	try {
		return await queue.schedule(hostKeyFor(imageUrl), () =>
			Attachment.downloadThumbnail(imageUrl, pageUrl))
	} catch (error) {
		console.log('[scrape] thumbnail failed', imageUrl, error.kind || error.message)
		return null
	}
}

const hostKeyFor = (url) => {
	try { return UrlSafety.registrableHost(new URL(url).hostname) }
	catch (error) { return '(none)' }
}

//
// Fetch one URL and write it to the cache. Returns the cache row fields.
//
const performFetch = async (url) => {
	const cached = await ScrapeCache.get(url)

	// Fresh and usable - zero network I/O.
	if (ScrapeCache.isFresh(cached) && ScrapeCache.isUsable(cached)) {
		return { fromCache: true, fields: ScrapeCache.toAttachmentFields(cached) }
	}

	// Parked failure still inside its backoff window.
	if (ScrapeCache.isFresh(cached) && cached.fetch_state !== 'ok') {
		const error = new Error(cached.last_error || 'previously failed')
		error.kind = cached.fetch_state === 'blocked' ? 'blocked' : 'http'
		error.status = cached.status || null
		throw error
	}

	if (settings.checkRobots !== false) {
		const allowed = await Robots.isAllowed(url)
		if (!allowed) {
			const error = new Error(`robots.txt disallows ${url}`)
			error.kind = 'robots_disallowed'
			await ScrapeCache.recordFailure(url, error)
			throw error
		}
	}

	let spacingMs = null
	if (settings.checkRobots !== false) {
		const crawlDelay = await Robots.crawlDelayMs(url)
		if (Robots.exceedsCrawlDelayCap(crawlDelay)) {
			const error = new Error(`Crawl-delay of ${crawlDelay}ms exceeds cap`)
			error.kind = 'blocked'
			await ScrapeCache.recordFailure(url, error)
			throw error
		}
		spacingMs = crawlDelay
	}

	await ScrapeCache.markPending(url)

	const conditional = ScrapeCache.conditionalFor(cached)

	let response = null
	try {
		response = await queue.schedule(hostKeyFor(url), () => HttpFetch.fetchDocument(url, {
			timeoutMs: settings.pageTimeoutMs,
			budgetMs: settings.totalBudgetMs,
			maxBytes: settings.maxHtmlBytes,
			conditional,
		}), { spacingMs })
	} catch (error) {
		await ScrapeCache.recordFailure(url, error)
		throw error
	}

	// Unchanged since last time - reuse what we already stored.
	if (response.notModified && cached) {
		await ScrapeCache.touchNotModified(url)
		const refreshed = await ScrapeCache.get(url)
		return { fromCache: true, notModified: true, fields: ScrapeCache.toAttachmentFields(refreshed) }
	}

	const $ = cheerio.load(response.html || '')
	const parsed = SiteScrape.parse($, response.finalUrl || url)

	// Reuse an existing thumbnail file rather than refetching the image.
	let fileLocation = (cached && cached.image_url === parsed.imageUrl && cached.file_location)
		? cached.file_location
		: await fetchThumbnail(parsed.imageUrl, response.finalUrl || url)

	await ScrapeCache.putSuccess(url, {
		...parsed,
		status: response.status,
		etag: response.etag,
		lastModified: response.lastModified,
		contentType: response.contentType,
		fileLocation,
	})

	return {
		fromCache: false,
		fields: {
			text: parsed.searchText,
			fileLocation,
			// Keep the URL the user wrote. The canonical URL is recorded in the
			// cache row; rewriting attachment.url would change link behaviour.
			url,
		},
	}
}

//
// Coalesce: one network fetch per URL no matter how many notes asked for it.
//
const fetchCoalesced = (url) => {
	const key = UrlSafety.urlKey(url)
	if (!key) { return Promise.reject(Object.assign(new Error(`unusable url ${url}`), { kind: 'blocked' })) }

	const existing = inFlight.get(key)
	if (existing) { return existing.promise }

	const promise = performFetch(url).finally(() => inFlight.delete(key))
	inFlight.set(key, { promise, url })

	return promise
}

const releaseUser = (userId) => {
	const count = (perUser.get(userId) || 0) - 1
	if (count <= 0) { perUser.delete(userId) } else { perUser.set(userId, count) }
}

//
// Queue one URL. Resolves with the search text (or '' on failure) and never
// rejects, so a caller can await a whole batch without special handling.
//
ScrapeQueue.enqueueUrl = ({ userId, noteId, attachmentId, url }) => {
	const job = runJob({ userId, noteId, attachmentId, url })

	// Registered synchronously, so drain() sees the job even before it reaches
	// its first await.
	outstanding.add(job)
	job.finally(() => outstanding.delete(job))

	return job
}

const runJob = async ({ userId, noteId, attachmentId, url }) => {
	if (!settings.enabled) { return '' }

	const pending = perUser.get(userId) || 0
	if (pending >= settings.maxPendingPerUser) {
		// Left at scrape_state = 0, so recoverPending() picks it up later.
		// Deferred, not dropped.
		console.log(`[scrape] deferring ${url} - user ${userId} at pending cap`)
		return ''
	}
	perUser.set(userId, pending + 1)

	try {
		await setAttachmentState(attachmentId, STATE.FETCHING)

		const result = await fetchCoalesced(url)
		await applyResultToAttachment(attachmentId, result.fields)

		notifyUser(userId)

		return result.fields.text || ''
	} catch (error) {
		await applyFailureToAttachment(attachmentId, error)
		notifyUser(userId)

		console.log(`[scrape] ${url} failed: ${error.kind || 'error'} ${error.message}`)
		return ''
	} finally {
		releaseUser(userId)
	}
}

//
// Queue a batch. allSettled, so one bad URL cannot strand the rest - this is
// the structural replacement for the old completion counter.
//
ScrapeQueue.enqueueUrls = async (jobs) => {
	if (!jobs || !jobs.length) { return [] }

	const results = await Promise.allSettled(jobs.map(job => ScrapeQueue.enqueueUrl(job)))

	return results.map(result => result.status === 'fulfilled' ? result.value : '')
}

//
// Wait for one URL already queued, with a hard cap. Used by the bookmarklet so
// it keeps responding after the scrape without being able to hang forever.
//
ScrapeQueue.awaitUrl = async (promise, capMs) => {
	let timer = null

	const timeout = new Promise(resolve => {
		timer = setTimeout(() => resolve('__scrape_timeout__'), capMs || settings.totalBudgetMs)
		if (timer.unref) { timer.unref() }
	})

	try {
		return await Promise.race([promise, timeout])
	} finally {
		if (timer) { clearTimeout(timer) }
	}
}

//
// Crash recovery.
//
// The queue is in memory, so a pm2 reload loses it - but every enqueued URL
// already has an attachment row, which makes the work re-derivable with no
// extra table. Rows stuck mid-fetch are reset, rows still pending are
// re-enqueued, and anything older than a day is retired so nothing
// recirculates forever.
//
ScrapeQueue.recoverPending = async () => {
	if (!settings.enabled) { return { requeued: 0, retired: 0 } }

	// Anything left at FETCHING died with the previous process.
	await db.promise().query(
		'UPDATE attachment SET scrape_state = ? WHERE scrape_state = ?',
		[STATE.PENDING, STATE.FETCHING]
	)

	const cutoff = now() - (24 * 60 * 60)

	const [retiredResult] = await db.promise().query(
		`UPDATE attachment SET scrape_state = ?
		 WHERE scrape_state = ? AND enqueued_at IS NOT NULL AND enqueued_at < ?`,
		[STATE.FAILED, STATE.PENDING, cutoff]
	)

	const [rows] = await db.promise().query(
		`SELECT id, user_id, note_id, url FROM attachment
		 WHERE attachment_type = 1 AND visible = 1 AND scrape_state = ?
		 ORDER BY enqueued_at IS NULL, enqueued_at
		 LIMIT ?`,
		[STATE.PENDING, settings.maxQueueLength]
	)

	if (rows.length) {
		console.log(`[scrape] recovering ${rows.length} pending URL(s) after restart`)
	}

	// Fire and forget: boot must not wait on the network.
	// Fire and forget, but the individual jobs register themselves in
	// `outstanding`, so drain() can still wait for them.
	ScrapeQueue.enqueueUrls(rows.map(row => ({
		userId: row.user_id,
		noteId: row.note_id,
		attachmentId: row.id,
		url: row.url,
	}))).catch(error => console.log('[scrape] recovery failed', error.message))

	return { requeued: rows.length, retired: retiredResult.affectedRows || 0 }
}