|
| 1 | +package ox.crawler |
| 2 | + |
| 3 | +import org.slf4j.LoggerFactory |
| 4 | +import ox.{discard, Fork, fork, supervised} |
| 5 | + |
| 6 | +import java.util.concurrent.{ArrayBlockingQueue, BlockingQueue} |
| 7 | +import scala.annotation.tailrec |
| 8 | + |
| 9 | +object Crawler: |
| 10 | + private val logger = LoggerFactory.getLogger(this.getClass) |
| 11 | + |
| 12 | + def crawl(crawlUrl: Url, http: Http, parseLinks: String => List[Url]): Map[Host, Int] = supervised { |
| 13 | + @tailrec |
| 14 | + def crawler(crawlerQueue: BlockingQueue[CrawlerMessage], data: CrawlerData): Map[Host, Int] = |
| 15 | + def handleMessage(msg: CrawlerMessage, data: CrawlerData): CrawlerData = msg match |
| 16 | + case Start(url) => crawlUrl(data, url) |
| 17 | + |
| 18 | + case CrawlResult(url, links) => |
| 19 | + val data2 = data.copy(inProgress = data.inProgress - url) |
| 20 | + |
| 21 | + links.foldLeft(data2) { case (d, link) => |
| 22 | + val d2 = d.copy(referenceCount = d.referenceCount.updated(link.host, d.referenceCount.getOrElse(link.host, 0) + 1)) |
| 23 | + crawlUrl(d2, link) |
| 24 | + } |
| 25 | + |
| 26 | + def crawlUrl(data: CrawlerData, url: Url): CrawlerData = |
| 27 | + if !data.visitedLinks.contains(url) then |
| 28 | + val (data2, workerQueue) = workerFor(data, url.host) |
| 29 | + workerQueue.put(url) |
| 30 | + data2.copy( |
| 31 | + visitedLinks = data.visitedLinks + url, |
| 32 | + inProgress = data.inProgress + url |
| 33 | + ) |
| 34 | + else data |
| 35 | + |
| 36 | + def workerFor(data: CrawlerData, host: Host): (CrawlerData, BlockingQueue[Url]) = |
| 37 | + data.workers.get(host) match |
| 38 | + case None => |
| 39 | + val workerQueue = new ArrayBlockingQueue[Url](32) |
| 40 | + worker(workerQueue, crawlerQueue).discard |
| 41 | + (data.copy(workers = data.workers + (host -> workerQueue)), workerQueue) |
| 42 | + case Some(queue) => (data, queue) |
| 43 | + |
| 44 | + val msg = crawlerQueue.take |
| 45 | + val data2 = handleMessage(msg, data) |
| 46 | + if data2.inProgress.isEmpty then data2.referenceCount else crawler(crawlerQueue, data2) |
| 47 | + end crawler |
| 48 | + |
| 49 | + def worker(workerQueue: BlockingQueue[Url], crawlerQueue: BlockingQueue[CrawlerMessage]): Fork[Unit] = |
| 50 | + def handleUrl(url: Url): Unit = |
| 51 | + val r = |
| 52 | + try parseLinks(http.get(url)) |
| 53 | + catch |
| 54 | + case e: Exception => |
| 55 | + logger.error(s"Cannot get contents of $url", e) |
| 56 | + List.empty[Url] |
| 57 | + fork(crawlerQueue.put(CrawlResult(url, r))).discard |
| 58 | + end handleUrl |
| 59 | + |
| 60 | + fork { |
| 61 | + while true do |
| 62 | + val url = workerQueue.take |
| 63 | + handleUrl(url) |
| 64 | + } |
| 65 | + end worker |
| 66 | + |
| 67 | + val crawlerQueue = new ArrayBlockingQueue[CrawlerMessage](32) |
| 68 | + crawlerQueue.put(Start(crawlUrl)) |
| 69 | + crawler(crawlerQueue, CrawlerData(Map(), Set(), Set(), Map())) |
| 70 | + } |
| 71 | + |
| 72 | + case class CrawlerData( |
| 73 | + referenceCount: Map[Host, Int], |
| 74 | + visitedLinks: Set[Url], |
| 75 | + inProgress: Set[Url], |
| 76 | + workers: Map[Host, BlockingQueue[Url]] |
| 77 | + ) |
| 78 | + |
| 79 | + sealed trait CrawlerMessage |
| 80 | + case class Start(url: Url) extends CrawlerMessage |
| 81 | + case class CrawlResult(url: Url, links: List[Url]) extends CrawlerMessage |
| 82 | +end Crawler |
0 commit comments