diff --git a/scrapy/core/downloader/__init__.py b/scrapy/core/downloader/__init__.py index 59c3ad074..4695d75f4 100644 --- a/scrapy/core/downloader/__init__.py +++ b/scrapy/core/downloader/__init__.py @@ -75,6 +75,8 @@ def _get_concurrency_delay(concurrency, spider, settings): class Downloader(object): + DOWNLOAD_SLOT = 'download_slot' + def __init__(self, crawler): self.settings = crawler.settings self.signals = crawler.signals @@ -111,8 +113,8 @@ class Downloader(object): return key, self.slots[key] def _get_slot_key(self, request, spider): - if 'download_slot' in request.meta: - return request.meta['download_slot'] + if self.DOWNLOAD_SLOT in request.meta: + return request.meta[self.DOWNLOAD_SLOT] key = urlparse_cached(request).hostname or '' if self.ip_concurrency: @@ -122,7 +124,7 @@ class Downloader(object): def _enqueue_request(self, request, spider): key, slot = self._get_slot(request, spider) - request.meta['download_slot'] = key + request.meta[self.DOWNLOAD_SLOT] = key def _deactivate(response): slot.active.remove(request) diff --git a/scrapy/core/queues.py b/scrapy/core/queues.py deleted file mode 100644 index 96d582fc7..000000000 --- a/scrapy/core/queues.py +++ /dev/null @@ -1,15 +0,0 @@ -import uuid -import os.path - - -def unique_files_queue(queue_class): - - class UniqueFilesQueue(queue_class): - def __init__(self, path): - path = path + "-" + uuid.uuid4().hex - while os.path.exists(path): - path = path + "-" + uuid.uuid4().hex - - super().__init__(path) - - return UniqueFilesQueue diff --git a/scrapy/core/scheduler.py b/scrapy/core/scheduler.py index eb790a67e..d40f3aa0c 100644 --- a/scrapy/core/scheduler.py +++ b/scrapy/core/scheduler.py @@ -13,7 +13,7 @@ logger = logging.getLogger(__name__) class Scheduler(object): def __init__(self, dupefilter, jobdir=None, dqclass=None, mqclass=None, - logunser=False, stats=None, pqclass=None): + logunser=False, stats=None, pqclass=None, crawler=None): self.df = dupefilter self.dqdir = self._dqdir(jobdir) self.pqclass = pqclass @@ -21,6 +21,7 @@ class Scheduler(object): self.mqclass = mqclass self.logunser = logunser self.stats = stats + self.crawler = crawler @classmethod def from_crawler(cls, crawler): @@ -32,14 +33,15 @@ class Scheduler(object): mqclass = load_object(settings['SCHEDULER_MEMORY_QUEUE']) logunser = settings.getbool('LOG_UNSERIALIZABLE_REQUESTS', settings.getbool('SCHEDULER_DEBUG')) return cls(dupefilter, jobdir=job_dir(settings), logunser=logunser, - stats=crawler.stats, pqclass=pqclass, dqclass=dqclass, mqclass=mqclass) + stats=crawler.stats, pqclass=pqclass, dqclass=dqclass, + mqclass=mqclass, crawler=crawler) def has_pending_requests(self): return len(self) > 0 def open(self, spider): self.spider = spider - self.mqs = self.pqclass(self._newmq) + self.mqs = create_instance(self.pqclass, None, self.crawler, self._newmq) self.dqs = self._dq() if self.dqdir else None return self.df.open() @@ -111,7 +113,7 @@ class Scheduler(object): return self.mqclass() def _newdq(self, priority): - return self.dqclass(join(self.dqdir, 'p%s' % priority)) + return self.dqclass(join(self.dqdir, 'p%s' % (priority, ))) def _dq(self): activef = join(self.dqdir, 'active.json') @@ -120,7 +122,12 @@ class Scheduler(object): prios = json.load(f) else: prios = () - q = self.pqclass(self._newdq, startprios=prios) + + q = create_instance(self.pqclass, + None, + self.crawler, + self._newdq, + startprios=prios) if q: logger.info("Resuming crawl (%(queuesize)d requests scheduled)", {'queuesize': len(q)}, extra={'spider': self.spider}) diff --git a/scrapy/pqueues.py b/scrapy/pqueues.py new file mode 100644 index 000000000..75073b7a4 --- /dev/null +++ b/scrapy/pqueues.py @@ -0,0 +1,246 @@ +from collections import deque +import hashlib +import logging +from six import text_type +from six.moves.urllib.parse import urlparse + +from queuelib import PriorityQueue + +from scrapy.core.downloader import Downloader +from scrapy.http import Request +from scrapy.signals import request_reached_downloader, response_downloaded + + +logger = logging.getLogger(__name__) + + +SCHEDULER_SLOT_META_KEY = Downloader.DOWNLOAD_SLOT + + +def _get_from_request(request, key, default=None): + if isinstance(request, dict): + return request.get(key, default) + + if isinstance(request, Request): + return getattr(request, key, default) + + raise ValueError('Bad type of request "%s"' % (request.__class__, )) + + +def _scheduler_slot_read(request, default=None): + meta = _get_from_request(request, 'meta', dict()) + slot = meta.get(SCHEDULER_SLOT_META_KEY, default) + return slot + + +def _scheduler_slot_write(request, slot): + meta = _get_from_request(request, 'meta', None) + if not isinstance(meta, dict): + raise ValueError('No meta attribute in %s' % (request, )) + meta[SCHEDULER_SLOT_META_KEY] = slot + + +def _scheduler_slot(request): + + slot = _scheduler_slot_read(request, None) + if slot is None: + url = _get_from_request(request, 'url') + slot = urlparse(url).hostname or '' + _scheduler_slot_write(request, slot) + + return slot + + +def _pathable(x): + pathable_slot = "".join([c if c.isalnum() or c in '-._' else '_' for c in x]) + + """ + as we replace some letters we can get collision for different slots + add we add unique part + """ + unique_slot = hashlib.md5(x.encode('utf8')).hexdigest() + + return '-'.join([pathable_slot, unique_slot]) + + +class PrioritySlot: + __slots__ = ('priority', 'slot') + + def __init__(self, priority=0, slot=None): + self.priority = priority + self.slot = slot + + def __hash__(self): + return hash((self.priority, self.slot)) + + def __eq__(self, other): + return (self.priority, self.slot) == (other.priority, other.slot) + + def __lt__(self, other): + return (self.priority, self.slot) < (other.priority, other.slot) + + def __str__(self): + return '_'.join([text_type(self.priority), _pathable(text_type(self.slot))]) + + +class PriorityAsTupleQueue(PriorityQueue): + """ + Python structures is not directly (de)serialized (to)from json. + We need this modified queue to transform custom structure (from)to + json serializable structures + """ + def __init__(self, qfactory, startprios=()): + + super(PriorityAsTupleQueue, self).__init__( + qfactory, + [PrioritySlot(priority=p[0], slot=p[1]) for p in startprios] + ) + + def close(self): + startprios = super(PriorityAsTupleQueue, self).close() + return [(s.priority, s.slot) for s in startprios] + + def is_empty(self): + return not self.queues or len(self) == 0 + + +class SlotBasedPriorityQueue(object): + + def __init__(self, qfactory, startprios={}): + self.pqueues = dict() # slot -> priority queue + self.qfactory = qfactory # factory for creating new internal queues + + if not startprios: + return + + if not isinstance(startprios, dict): + raise ValueError("Looks like your priorities file malforfemed. " + "Possible reason: You run scrapy with previous " + "version. Interrupted it. Updated scrapy. And " + "run again.") + + for slot, prios in startprios.items(): + self.pqueues[slot] = PriorityAsTupleQueue(self.qfactory, prios) + + def pop_slot(self, slot): + queue = self.pqueues[slot] + request = queue.pop() + is_empty = queue.is_empty() + if is_empty: + del self.pqueues[slot] + + return request, is_empty + + def push_slot(self, request, priority): + slot = _scheduler_slot(request) + is_new = False + if slot not in self.pqueues: + is_new = True + self.pqueues[slot] = PriorityAsTupleQueue(self.qfactory) + self.pqueues[slot].push(request, PrioritySlot(priority=priority, slot=slot)) + return slot, is_new + + def close(self): + startprios = dict() + for slot, queue in self.pqueues.items(): + prios = queue.close() + startprios[slot] = prios + self.pqueues.clear() + return startprios + + def __len__(self): + return sum(len(x) for x in self.pqueues.values()) if self.pqueues else 0 + + +class RoundRobinPriorityQueue(SlotBasedPriorityQueue): + + def __init__(self, qfactory, startprios={}): + super(RoundRobinPriorityQueue, self).__init__(qfactory, startprios) + self._slots = deque() + for slot in self.pqueues: + self._slots.append(slot) + + def push(self, request, priority): + slot, is_new = self.push_slot(request, priority) + if is_new: + self._slots.append(slot) + + def pop(self): + if not self._slots: + return + + slot = self._slots.popleft() + request, is_empty = self.pop_slot(slot) + + if not is_empty: + self._slots.append(slot) + + return request + + def close(self): + self._slots.clear() + return super(RoundRobinPriorityQueue, self).close() + + +class DownloaderAwarePriorityQueue(SlotBasedPriorityQueue): + + _DOWNLOADER_AWARE_PQ_ID = 'DOWNLOADER_AWARE_PQ_ID' + + @classmethod + def from_crawler(cls, crawler, qfactory, startprios={}): + return cls(crawler, qfactory, startprios) + + def __init__(self, crawler, qfactory, startprios={}): + super(DownloaderAwarePriorityQueue, self).__init__(qfactory, startprios) + self._slots = {slot: 0 for slot in self.pqueues} + crawler.signals.connect(self.on_response_download, + signal=response_downloaded) + crawler.signals.connect(self.on_request_reached_downloader, + signal=request_reached_downloader) + + def mark(self, request): + meta = _get_from_request(request, 'meta', None) + if not isinstance(meta, dict): + raise ValueError('No meta attribute in %s' % (request, )) + meta[self._DOWNLOADER_AWARE_PQ_ID] = id(self) + + def check_mark(self, request): + return request.meta.get(self._DOWNLOADER_AWARE_PQ_ID, None) == id(self) + + def pop(self): + slots = [(d, s) for s,d in self._slots.items() if s in self.pqueues] + + if not slots: + return + + slot = min(slots)[1] + request, _ = self.pop_slot(slot) + self.mark(request) + return request + + def push(self, request, priority): + slot, _ = self.push_slot(request, priority) + if slot not in self._slots: + self._slots[slot] = 0 + + def on_response_download(self, response, request, spider): + if not self.check_mark(request): + return + + slot = _scheduler_slot_read(request) + if slot not in self._slots or self._slots[slot] <= 0: + raise ValueError('Get response for wrong slot "%s"' % (slot, )) + self._slots[slot] = self._slots[slot] - 1 + if self._slots[slot] == 0 and slot not in self.pqueues: + del self._slots[slot] + + def on_request_reached_downloader(self, request, spider): + if not self.check_mark(request): + return + + slot = _scheduler_slot_read(request) + self._slots[slot] = self._slots.get(slot, 0) + 1 + + def close(self): + self._slots.clear() + return super(DownloaderAwarePriorityQueue, self).close() diff --git a/tests/test_scheduler.py b/tests/test_scheduler.py new file mode 100644 index 000000000..fd86e8d8c --- /dev/null +++ b/tests/test_scheduler.py @@ -0,0 +1,315 @@ +import contextlib +import shutil +import tempfile +import unittest + +from scrapy.crawler import Crawler +from scrapy.core.scheduler import Scheduler +from scrapy.http import Request +from scrapy.pqueues import _scheduler_slot_read, _scheduler_slot_write +from scrapy.signals import request_reached_downloader, response_downloaded +from scrapy.spiders import Spider + +class MockCrawler(Crawler): + def __init__(self, priority_queue_cls, jobdir): + + settings = dict(LOG_UNSERIALIZABLE_REQUESTS=False, + SCHEDULER_DISK_QUEUE='scrapy.squeues.PickleLifoDiskQueue', + SCHEDULER_MEMORY_QUEUE='scrapy.squeues.LifoMemoryQueue', + SCHEDULER_PRIORITY_QUEUE=priority_queue_cls, + JOBDIR=jobdir, + DUPEFILTER_CLASS='scrapy.dupefilters.BaseDupeFilter') + super(MockCrawler, self).__init__(Spider, settings) + + +class SchedulerHandler: + priority_queue_cls = None + jobdir = None + + def create_scheduler(self): + self.mock_crawler = MockCrawler(self.priority_queue_cls, self.jobdir) + self.scheduler = Scheduler.from_crawler(self.mock_crawler) + self.spider = Spider(name='spider') + self.scheduler.open(self.spider) + + def close_scheduler(self): + self.scheduler.close('finished') + self.mock_crawler.stop() + + def setUp(self): + self.create_scheduler() + + def tearDown(self): + self.close_scheduler() + + +_PRIORITIES = [("http://foo.com/a", -2), + ("http://foo.com/d", 1), + ("http://foo.com/b", -1), + ("http://foo.com/c", 0), + ("http://foo.com/e", 2)] + + +_URLS = {"http://foo.com/a", "http://foo.com/b", "http://foo.com/c"} + + +class BaseSchedulerInMemoryTester(SchedulerHandler): + def test_length(self): + self.assertFalse(self.scheduler.has_pending_requests()) + self.assertEqual(len(self.scheduler), 0) + + for url in _URLS: + self.scheduler.enqueue_request(Request(url)) + + self.assertTrue(self.scheduler.has_pending_requests()) + self.assertEqual(len(self.scheduler), len(_URLS)) + + def test_dequeue(self): + for url in _URLS: + self.scheduler.enqueue_request(Request(url)) + + urls = set() + while self.scheduler.has_pending_requests(): + urls.add(self.scheduler.next_request().url) + + self.assertEqual(urls, _URLS) + + def test_dequeue_priorities(self): + for url, priority in _PRIORITIES: + self.scheduler.enqueue_request(Request(url, priority=priority)) + + priorities = list() + while self.scheduler.has_pending_requests(): + priorities.append(self.scheduler.next_request().priority) + + self.assertEqual(priorities, sorted([x[1] for x in _PRIORITIES], key=lambda x: -x)) + + +class BaseSchedulerOnDiskTester(SchedulerHandler): + + def setUp(self): + self.jobdir = tempfile.mkdtemp() + self.create_scheduler() + + def tearDown(self): + self.close_scheduler() + + shutil.rmtree(self.jobdir) + self.jobdir = None + + def test_length(self): + self.assertFalse(self.scheduler.has_pending_requests()) + self.assertEqual(len(self.scheduler), 0) + + for url in _URLS: + self.scheduler.enqueue_request(Request(url)) + + self.close_scheduler() + self.create_scheduler() + + self.assertTrue(self.scheduler.has_pending_requests()) + self.assertEqual(len(self.scheduler), len(_URLS)) + + def test_dequeue(self): + for url in _URLS: + self.scheduler.enqueue_request(Request(url)) + + self.close_scheduler() + self.create_scheduler() + + urls = set() + while self.scheduler.has_pending_requests(): + urls.add(self.scheduler.next_request().url) + + self.assertEqual(urls, _URLS) + + def test_dequeue_priorities(self): + for url, priority in _PRIORITIES: + self.scheduler.enqueue_request(Request(url, priority=priority)) + + self.close_scheduler() + self.create_scheduler() + + priorities = list() + while self.scheduler.has_pending_requests(): + priorities.append(self.scheduler.next_request().priority) + + self.assertEqual(priorities, sorted([x[1] for x in _PRIORITIES], key=lambda x: -x)) + + +class TestSchedulerInMemory(BaseSchedulerInMemoryTester, unittest.TestCase): + priority_queue_cls = 'queuelib.PriorityQueue' + + +class TestSchedulerOnDisk(BaseSchedulerOnDiskTester, unittest.TestCase): + priority_queue_cls = 'queuelib.PriorityQueue' + + +_SLOTS = [("http://foo.com/a", 'a'), + ("http://foo.com/b", 'a'), + ("http://foo.com/c", 'b'), + ("http://foo.com/d", 'b'), + ("http://foo.com/e", 'c'), + ("http://foo.com/f", 'c')] + + +class TestSchedulerWithRoundRobinInMemory(BaseSchedulerInMemoryTester, unittest.TestCase): + priority_queue_cls = 'scrapy.pqueues.RoundRobinPriorityQueue' + + def test_round_robin(self): + for url, slot in _SLOTS: + request = Request(url) + _scheduler_slot_write(request, slot) + self.scheduler.enqueue_request(request) + + slots = list() + while self.scheduler.has_pending_requests(): + slots.append(_scheduler_slot_read(self.scheduler.next_request())) + + for i in range(0, len(_SLOTS), 2): + self.assertNotEqual(slots[i], slots[i+1]) + + def test_is_meta_set(self): + url = "http://foo.com/a" + request = Request(url) + if _scheduler_slot_read(request): + _scheduler_slot_write(request, None) + self.scheduler.enqueue_request(request) + self.assertIsNotNone(_scheduler_slot_read(request, None), None) + + +class TestSchedulerWithRoundRobinOnDisk(BaseSchedulerOnDiskTester, unittest.TestCase): + priority_queue_cls = 'scrapy.pqueues.RoundRobinPriorityQueue' + + def test_round_robin(self): + for url, slot in _SLOTS: + request = Request(url) + _scheduler_slot_write(request, slot) + self.scheduler.enqueue_request(request) + + self.close_scheduler() + self.create_scheduler() + + slots = list() + while self.scheduler.has_pending_requests(): + slots.append(_scheduler_slot_read(self.scheduler.next_request())) + + for i in range(0, len(_SLOTS), 2): + self.assertNotEqual(slots[i], slots[i+1]) + + def test_is_meta_set(self): + url = "http://foo.com/a" + request = Request(url) + if _scheduler_slot_read(request): + _scheduler_slot_write(request, None) + self.scheduler.enqueue_request(request) + + self.close_scheduler() + self.create_scheduler() + + self.assertIsNotNone(_scheduler_slot_read(request, None), None) + + +@contextlib.contextmanager +def mkdtemp(): + dir = tempfile.mkdtemp() + try: + yield dir + finally: + shutil.rmtree(dir) + + +def _migration(): + + with mkdtemp() as tmp_dir: + prev_scheduler_handler = SchedulerHandler() + prev_scheduler_handler.priority_queue_cls = 'queuelib.PriorityQueue' + prev_scheduler_handler.jobdir = tmp_dir + + prev_scheduler_handler.create_scheduler() + for url in _URLS: + prev_scheduler_handler.scheduler.enqueue_request(Request(url)) + prev_scheduler_handler.close_scheduler() + + next_scheduler_handler = SchedulerHandler() + next_scheduler_handler.priority_queue_cls = 'scrapy.pqueues.RoundRobinPriorityQueue' + next_scheduler_handler.jobdir = tmp_dir + + next_scheduler_handler.create_scheduler() + + +class TestMigration(unittest.TestCase): + def test_migration(self): + self.assertRaises(ValueError, _migration) + + +class TestSchedulerWithDownloaderAwareInMemory(BaseSchedulerInMemoryTester, unittest.TestCase): + priority_queue_cls = 'scrapy.pqueues.DownloaderAwarePriorityQueue' + + def test_logic(self): + for url, slot in _SLOTS: + request = Request(url) + _scheduler_slot_write(request, slot) + self.scheduler.enqueue_request(request) + + slots = list() + requests = list() + while self.scheduler.has_pending_requests(): + request = self.scheduler.next_request() + slots.append(_scheduler_slot_read(request)) + self.mock_crawler.signals.send_catch_log( + signal=request_reached_downloader, + request=request, + spider=self.spider + ) + requests.append(request) + self.assertEqual(len(slots), len(_SLOTS)) + + for request in requests: + self.mock_crawler.signals.send_catch_log(signal=response_downloaded, + request=request, + response=None, + spider=self.spider) + + unique_slots = len(set(s for _, s in _SLOTS)) + for i in range(0, len(_SLOTS), unique_slots): + part = slots[i:i + unique_slots] + self.assertEqual(len(part), len(set(part))) + + +class TestSchedulerWithDownloaderAwareOnDisk(BaseSchedulerOnDiskTester, unittest.TestCase): + priority_queue_cls = 'scrapy.pqueues.DownloaderAwarePriorityQueue' + def test_logic(self): + for url, slot in _SLOTS: + request = Request(url) + _scheduler_slot_write(request, slot) + self.scheduler.enqueue_request(request) + + self.close_scheduler() + self.create_scheduler() + + slots = list() + requests = list() + while self.scheduler.has_pending_requests(): + request = self.scheduler.next_request() + slots.append(_scheduler_slot_read(request)) + self.mock_crawler.signals.send_catch_log( + signal=request_reached_downloader, + request=request, + spider=self.spider + ) + requests.append(request) + + self.assertEqual(self.scheduler.mqs._slots, {}) + self.assertEqual(len(slots), len(_SLOTS)) + + for request in requests: + self.mock_crawler.signals.send_catch_log(signal=response_downloaded, + request=request, + response=None, + spider=self.spider) + + unique_slots = len(set(s for _, s in _SLOTS)) + for i in range(0, len(_SLOTS), unique_slots): + part = slots[i:i + unique_slots] + self.assertEqual(len(part), len(set(part)))