mirror of https://github.com/scrapy/scrapy.git
66 lines
2.4 KiB
Python
66 lines
2.4 KiB
Python
"""
|
|
Request Limit Spider middleware
|
|
|
|
See documentation in docs/topics/spider-middleware.rst
|
|
"""
|
|
from itertools import imap
|
|
from scrapy.xlib.pydispatch import dispatcher
|
|
|
|
from scrapy import signals
|
|
from scrapy.project import crawler
|
|
from scrapy.exceptions import NotConfigured
|
|
from scrapy.conf import settings
|
|
from scrapy.http import Request
|
|
from scrapy import log
|
|
|
|
class RequestLimitMiddleware(object):
|
|
|
|
def __init__(self):
|
|
self.max_queue_size = settings.getint("REQUESTS_QUEUE_SIZE")
|
|
if not self.max_queue_size:
|
|
raise NotConfigured
|
|
|
|
self.max_pending = {}
|
|
self.dropped_count = {}
|
|
|
|
dispatcher.connect(self.spider_opened, signal=signals.spider_opened)
|
|
dispatcher.connect(self.spider_closed, signal=signals.spider_closed)
|
|
|
|
def spider_opened(self, spider):
|
|
self.max_pending[spider] = getattr(spider, 'requests_queue_size', self.max_queue_size)
|
|
self.dropped_count[spider] = 0
|
|
|
|
def spider_closed(self, spider):
|
|
dropped_count = self.dropped_count[spider]
|
|
if dropped_count:
|
|
max_pending = self.max_pending[spider]
|
|
log.msg('Dropped %d request(s) because the scheduler queue size limit (%d requests) was exceeded' % \
|
|
(dropped_count, max_pending), level=log.DEBUG, spider=spider)
|
|
del self.dropped_count[spider]
|
|
del self.max_pending[spider]
|
|
|
|
def process_spider_output(self, response, result, spider):
|
|
max_pending = self.max_pending.get(spider, 0)
|
|
if max_pending:
|
|
return imap(lambda v: self._limit_requests(v, spider, max_pending), result)
|
|
else:
|
|
return result
|
|
|
|
def _limit_requests(self, request_or_other, spider, max_pending):
|
|
if isinstance(request_or_other, Request):
|
|
free_slots = max_pending - self._pending_count(spider)
|
|
if free_slots > 0:
|
|
# Scheduler isn't saturated and it is fine to schedule more requests.
|
|
return request_or_other
|
|
else:
|
|
# Skip the request and give engine time to handle other tasks.
|
|
self.dropped_count[spider] += 1
|
|
return None
|
|
else:
|
|
# Return others (non-requests) as is.
|
|
return request_or_other
|
|
|
|
def _pending_count(self, spider):
|
|
pending = crawler.engine.scheduler.pending_requests.get(spider, [])
|
|
return len(pending)
|