A framework for quick web archiving
選択できるのは25トピックまでです。 トピックは、先頭が英数字で、英数字とダッシュ('-')を使用した35文字以内のものにしてください。

429 行
16 KiB

  1. import qwarc.aiohttp
  2. from qwarc.const import *
  3. import qwarc.utils
  4. import qwarc.warc
  5. import aiohttp as _aiohttp
  6. if _aiohttp.__version__ != '2.3.10':
  7. raise ImportError('aiohttp must be version 2.3.10')
  8. import asyncio
  9. import collections
  10. import concurrent.futures
  11. import contextlib
  12. import io
  13. import itertools
  14. import logging
  15. import os
  16. import random
  17. import sqlite3
  18. import yarl
  19. class Item:
  20. itemType = None
  21. defaultResponseHandler = staticmethod(qwarc.utils.handle_response_default)
  22. def __init__(self, qwarcObj, itemValue, session, headers, warc):
  23. self.qwarcObj = qwarcObj
  24. self.itemValue = itemValue
  25. self.session = session
  26. self.headers = headers
  27. self.warc = warc
  28. if not hasattr(self, '_baseUrl'): # To allow subclasses to set the baseUrl before calling super().__init__
  29. self._baseUrl = None
  30. self.stats = {'tx': 0, 'rx': 0, 'requests': 0}
  31. self.logger = logging.LoggerAdapter(logging.getLogger(), {'itemType': self.itemType, 'itemValue': self.itemValue})
  32. self.childItems = []
  33. @property
  34. def baseUrl(self):
  35. return self._baseUrl
  36. @baseUrl.setter
  37. def baseUrl(self, baseUrl):
  38. if baseUrl is None:
  39. self._baseUrl = None
  40. elif isinstance(baseUrl, yarl.URL):
  41. self._baseUrl = baseUrl
  42. else:
  43. self._baseUrl = yarl.URL(baseUrl)
  44. def _merge_headers(self, headers, extraHeaders = []):
  45. d = {} # Preserves order from Python 3.7 (guaranteed) or CPython 3.6 (implementation detail)
  46. keys = {} # casefolded key -> d key
  47. for key, value in itertools.chain(self.headers, extraHeaders, headers):
  48. keyc = key.casefold()
  49. if value is None:
  50. if keyc in keys:
  51. del d[keys[keyc]]
  52. del keys[keyc]
  53. else:
  54. if keyc in keys and key != keys[keyc]:
  55. del d[keys[keyc]]
  56. d[key] = value
  57. keys[keyc] = key
  58. out = []
  59. for key, value in d.items():
  60. if isinstance(value, tuple):
  61. for value_ in value:
  62. out.append((key, value_))
  63. else:
  64. out.append((key, value))
  65. return out
  66. async def fetch(self, url, responseHandler = None, method = 'GET', data = None, headers = [], verify_ssl = True, timeout = 60, fromResponse = None):
  67. '''
  68. HTTP GET or POST a URL
  69. url: str or yarl.URL; if this is not a complete URL, it is evaluated relative to self.baseUrl
  70. responseHandler: None or a callable that determines how the response is handled; if None, self.defaultResponseHandler is used. See qwarc.utils.handle_response_default for details.
  71. method: str, must be 'GET' or 'POST'
  72. data: dict or list/tuple of lists/tuples of length two or bytes or file-like or None, the data to be sent in the request body
  73. headers: list of 2-tuples, additional or overriding headers for this request only
  74. To remove one of the default headers, pass a value of None.
  75. If a header appears multiple times, only the last one is used. To send a header multiple times, pass a tuple of values.
  76. verify_ssl: bool, whether the SSL/TLS certificate should be validated
  77. timeout: int or float, how long the fetch may take at most in total (sending request until finishing reading the response)
  78. fromResponse: ClientResponse or None; if provided, use fromResponse.url for the url completion (instead of self.baseUrl) and add it as a Referer header
  79. Returns response (a ClientResponse object or a qwarc.utils.DummyClientResponse object)
  80. '''
  81. #TODO: Rewrite using 'async with self.session.get'
  82. url = yarl.URL(url) # Explicitly convert for normalisation, percent-encoding, etc.
  83. if not url.scheme or not url.host:
  84. if fromResponse is not None:
  85. url = fromResponse.url.join(url)
  86. elif not self.baseUrl:
  87. raise ValueError('Incomplete URL and no baseUrl to join it with')
  88. else:
  89. url = self.baseUrl.join(url)
  90. originalUrl = url
  91. if responseHandler is None:
  92. responseHandler = self.defaultResponseHandler
  93. assert method in ('GET', 'POST'), 'method must be GET or POST'
  94. headers = self._merge_headers(headers, extraHeaders = [('Referer', str(fromResponse.url))] if fromResponse is not None else [])
  95. history = []
  96. attempt = 0
  97. redirectLevel = 0
  98. while True:
  99. attempt += 1
  100. response = None
  101. exc = None
  102. action = ACTION_RETRY
  103. writeToWarc = True
  104. try:
  105. try:
  106. with _aiohttp.Timeout(timeout):
  107. self.logger.info(f'Fetching {url}')
  108. response = await self.session.request(method, url, data = data, headers = headers, allow_redirects = False, verify_ssl = verify_ssl)
  109. try:
  110. while True:
  111. ret = await response.content.read(1048576)
  112. if not ret:
  113. break
  114. except:
  115. # No calling the handleResponse callback here because this is really bad. The not-so-bad exceptions (e.g. an error during reading the response) will be caught further down.
  116. response.close()
  117. raise
  118. else:
  119. response.rawRequestData.seek(0, io.SEEK_END)
  120. tx = response.rawRequestData.tell()
  121. response.rawResponseData.seek(0, io.SEEK_END)
  122. rx = response.rawResponseData.tell()
  123. self.logger.info(f'Fetched {url}: {response.status} (tx {tx}, rx {rx})')
  124. self.stats['tx'] += tx
  125. self.stats['rx'] += rx
  126. self.stats['requests'] += 1
  127. except (asyncio.TimeoutError, _aiohttp.ClientError) as e:
  128. self.logger.warning(f'Request for {url} failed: {e!r}')
  129. action, writeToWarc = await responseHandler(url = url, attempt = attempt, response = response, exc = e, redirectLevel = redirectLevel, item = self)
  130. exc = e # Pass the exception outward for the history
  131. else:
  132. action, writeToWarc = await responseHandler(url = url, attempt = attempt, response = response, exc = None, redirectLevel = redirectLevel, item = self)
  133. if response and exc is None and writeToWarc:
  134. self.warc.write_client_response(response)
  135. history.append((response, exc))
  136. retResponse = response if exc is None else qwarc.utils.DummyClientResponse()
  137. if action in (ACTION_SUCCESS, ACTION_IGNORE):
  138. retResponse.qhistory = tuple(history)
  139. return retResponse
  140. elif action == ACTION_FOLLOW_OR_SUCCESS:
  141. redirectUrl = response.headers.get('Location') or response.headers.get('URI')
  142. if not redirectUrl:
  143. retResponse.qhistory = tuple(history)
  144. return retResponse
  145. url = url.join(yarl.URL(redirectUrl))
  146. if response.status in (301, 302, 303) and method == 'POST':
  147. method = 'GET'
  148. data = None
  149. attempt = 0
  150. redirectLevel += 1
  151. elif action == ACTION_RETRIES_EXCEEDED:
  152. self.logger.error(f'Request for {url} failed {attempt} times')
  153. retResponse.qhistory = tuple(history)
  154. return retResponse
  155. elif action == ACTION_TOO_MANY_REDIRECTS:
  156. self.logger.error(f'Request for {url} (from {originalUrl}) exceeded redirect limit')
  157. retResponse.qhistory = tuple(history)
  158. return retResponse
  159. elif action == ACTION_RETRY:
  160. # Nothing to do, just go to the next cycle
  161. pass
  162. finally:
  163. if response:
  164. await response.release()
  165. async def process(self):
  166. raise NotImplementedError
  167. @classmethod
  168. def generate(cls):
  169. yield from () # Generate no items by default
  170. def add_subitem(self, itemClassOrType, itemValue):
  171. if issubclass(itemClassOrType, Item):
  172. item = (itemClassOrType.itemType, itemValue)
  173. else:
  174. item = (itemClassOrType, itemValue)
  175. if item not in self.childItems:
  176. self.childItems.append(item)
  177. async def flush_subitems(self):
  178. await self.qwarcObj.flush_subitems(self)
  179. def clear_subitems(self):
  180. self.childItems = []
  181. @classmethod
  182. def get_subclasses(cls):
  183. for subclass in cls.__subclasses__():
  184. yield subclass
  185. yield from subclass.get_subclasses()
  186. class QWARC:
  187. def __init__(self, itemClasses, warcBasePath, dbPath, command, specFile, specDependencies, logFilename, concurrency = 1, memoryLimit = 0, minFreeDisk = 0, warcSizeLimit = 0, warcDedupe = False):
  188. '''
  189. itemClasses: iterable of Item
  190. warcBasePath: str, base name of the WARC files
  191. dbPath: str, path to the sqlite3 database file
  192. command: list, the command line used to invoke qwarc
  193. specFile: str, path to the spec file
  194. specDependencies: qwarc.utils.SpecDependencies
  195. logFilename: str, name of the log file written by this process
  196. concurrency: int, number of concurrently processed items
  197. memoryLimit: int, gracefully stop when the process uses more than memoryLimit bytes of RSS; 0 disables the memory check
  198. minFreeDisk: int, pause when there's less than minFreeDisk space on the partition where WARCs are written; 0 disables the disk space check
  199. warcSizeLimit: int, size of each WARC file; 0 if the WARCs should not be split
  200. '''
  201. self._itemClasses = itemClasses
  202. self._itemTypeMap = {cls.itemType: cls for cls in itemClasses}
  203. self._warcBasePath = warcBasePath
  204. self._dbPath = dbPath
  205. self._command = command
  206. self._specFile = specFile
  207. self._specDependencies = specDependencies
  208. self._logFilename = logFilename
  209. self._concurrency = concurrency
  210. self._memoryLimit = memoryLimit
  211. self._minFreeDisk = minFreeDisk
  212. self._warcSizeLimit = warcSizeLimit
  213. self._warcDedupe = warcDedupe
  214. self._reset_working_vars()
  215. def _reset_working_vars(self):
  216. # Working variables
  217. self._db = None
  218. self._tasks = set()
  219. self._sleepTasks = set()
  220. self._sessions = [] # aiohttp.ClientSession instances
  221. self._freeSessions = collections.deque() # ClientSession instances that are currently free
  222. self._warc = None
  223. @contextlib.asynccontextmanager
  224. async def exclusive_db_lock(self):
  225. c = self._db.cursor()
  226. while True:
  227. try:
  228. c.execute('BEGIN EXCLUSIVE')
  229. break
  230. except sqlite3.OperationalError as e:
  231. if str(e) != 'database is locked':
  232. raise
  233. await asyncio.sleep(1)
  234. try:
  235. yield c
  236. c.execute('COMMIT')
  237. except:
  238. c.execute('ROLLBACK')
  239. raise
  240. def _make_item(self, itemType, itemValue, session, headers):
  241. try:
  242. itemClass = self._itemTypeMap[itemType]
  243. except KeyError:
  244. raise RuntimeError(f'No such item type: {itemType!r}')
  245. return itemClass(self, itemValue, session, headers, self._warc)
  246. async def _wait_for_free_task(self):
  247. if not self._tasks:
  248. return
  249. done, pending = await asyncio.wait(self._tasks, return_when = concurrent.futures.FIRST_COMPLETED)
  250. for future in done:
  251. newStatus = STATUS_DONE
  252. if future.taskType == 'sleep':
  253. self._sleepTasks.remove(future)
  254. elif future.taskType == 'process':
  255. item = future.item
  256. try:
  257. future.result()
  258. except asyncio.CancelledError as e:
  259. # Got cancelled, nothing we can do about it, but let's log a warning if it's a process task
  260. if future.taskType == 'process':
  261. logging.error(f'Task for {future.itemType}:{future.itemValue} cancelled: {future!r}')
  262. newStatus = STATUS_ERROR
  263. except Exception as e:
  264. if future.taskType == 'process':
  265. logging.error(f'{future.itemType}:{future.itemValue} failed: {e!r} ({item.stats["requests"]} requests, {item.stats["tx"]} tx, {item.stats["rx"]} rx)', exc_info = e)
  266. newStatus = STATUS_ERROR
  267. else:
  268. if future.taskType == 'process':
  269. logging.info(f'{future.itemType}:{future.itemValue} done: {item.stats["requests"]} requests, {item.stats["tx"]} tx, {item.stats["rx"]} rx')
  270. if future.taskType != 'process':
  271. continue
  272. async with self.exclusive_db_lock() as cursor:
  273. cursor.execute('UPDATE items SET status = ? WHERE id = ?', (newStatus, future.id))
  274. await self._insert_subitems(item)
  275. self._freeSessions.append(item.session)
  276. self._tasks = pending
  277. async def _insert_subitems(self, item):
  278. async with self.exclusive_db_lock() as cursor:
  279. if item.childItems:
  280. it = iter(item.childItems)
  281. while True:
  282. values = [(t, v, STATUS_TODO) for t, v in itertools.islice(it, 100000)]
  283. if not values:
  284. break
  285. cursor.executemany('INSERT OR IGNORE INTO items (type, value, status) VALUES (?, ?, ?)', values)
  286. async def run(self, loop):
  287. for i in range(self._concurrency):
  288. session = _aiohttp.ClientSession(
  289. connector = qwarc.aiohttp.TCPConnector(loop = loop),
  290. request_class = qwarc.aiohttp.ClientRequest,
  291. response_class = qwarc.aiohttp.ClientResponse,
  292. loop = loop
  293. )
  294. self._sessions.append(session)
  295. self._freeSessions.append(session)
  296. self._warc = qwarc.warc.WARC(self._warcBasePath, self._warcSizeLimit, self._warcDedupe, self._command, self._specFile, self._specDependencies, self._logFilename)
  297. self._db = sqlite3.connect(self._dbPath, timeout = 1)
  298. self._db.isolation_level = None # Transactions are handled manually below.
  299. self._db.execute('PRAGMA synchronous = OFF')
  300. async with self.exclusive_db_lock() as cursor:
  301. cursor.execute('SELECT name FROM sqlite_master WHERE type = "table" AND name = "items"')
  302. result = cursor.fetchone()
  303. if not result:
  304. self._create_db(cursor)
  305. self._insert_generated_items(cursor)
  306. try:
  307. while True:
  308. while len(self._tasks) >= self._concurrency:
  309. await self._wait_for_free_task()
  310. if os.path.exists('STOP'):
  311. logging.info('Gracefully shutting down due to STOP file')
  312. break
  313. if self._memoryLimit and qwarc.utils.uses_too_much_memory(self._memoryLimit):
  314. logging.info(f'Gracefully shutting down due to memory usage (current = {qwarc.utils.get_rss()} > limit = {self._memoryLimit})')
  315. break
  316. if self._minFreeDisk and qwarc.utils.too_little_disk_space(self._minFreeDisk):
  317. logging.info('Disk space is low, sleeping')
  318. sleepTask = asyncio.ensure_future(asyncio.sleep(random.uniform(self._concurrency / 2, self._concurrency * 1.5)))
  319. sleepTask.taskType = 'sleep'
  320. self._tasks.add(sleepTask)
  321. self._sleepTasks.add(sleepTask)
  322. continue
  323. async with self.exclusive_db_lock() as cursor:
  324. cursor.execute('SELECT id, type, value, status FROM items WHERE status = ? LIMIT 1', (STATUS_TODO,))
  325. result = cursor.fetchone()
  326. if not result:
  327. if cursor.execute('SELECT id, status FROM items WHERE status != ? LIMIT 1', (STATUS_DONE,)).fetchone():
  328. # There is currently no item to do, but there are still some in progress, so more TODOs may appear in the future.
  329. # It would be nice if we could just await wait_for_free_task() here, but that doesn't work because those TODOs might be in another process.
  330. # So instead, we insert a dummy task which just sleeps a bit. Average sleep time is equal to concurrency, i.e. one check per second.
  331. #TODO: The average sleep time is too large if there are only few sleep tasks; scale with len(sleepTasks)/self._concurrency?
  332. sleepTask = asyncio.ensure_future(asyncio.sleep(random.uniform(self._concurrency / 2, self._concurrency * 1.5)))
  333. sleepTask.taskType = 'sleep'
  334. self._tasks.add(sleepTask)
  335. self._sleepTasks.add(sleepTask)
  336. continue
  337. else:
  338. # Really nothing to do anymore
  339. break
  340. id_, itemType, itemValue, status = result
  341. cursor.execute('UPDATE items SET status = ? WHERE id = ?', (STATUS_INPROGRESS, id_))
  342. session = self._freeSessions.popleft()
  343. item = self._make_item(itemType, itemValue, session, DEFAULT_HEADERS)
  344. task = asyncio.ensure_future(item.process())
  345. #TODO: Is there a better way to add custom information to a task/coroutine object?
  346. task.taskType = 'process'
  347. task.id = id_
  348. task.itemType = itemType
  349. task.itemValue = itemValue
  350. task.item = item
  351. self._tasks.add(task)
  352. for sleepTask in self._sleepTasks:
  353. sleepTask.cancel()
  354. while len(self._tasks):
  355. await self._wait_for_free_task()
  356. logging.info('Done')
  357. except (Exception, KeyboardInterrupt) as e:
  358. # Kill all tasks
  359. for task in self._tasks:
  360. task.cancel()
  361. await asyncio.wait(self._tasks, return_when = concurrent.futures.ALL_COMPLETED)
  362. raise
  363. finally:
  364. for session in self._sessions:
  365. session.close()
  366. self._warc.close()
  367. self._db.close()
  368. self._reset_working_vars()
  369. async def flush_subitems(self, item):
  370. await self._insert_subitems(item)
  371. item.clear_subitems()
  372. def _create_db(self, cursor):
  373. cursor.execute('CREATE TABLE items (id INTEGER PRIMARY KEY, type TEXT, value TEXT, status INTEGER)')
  374. cursor.execute('CREATE INDEX items_status_idx ON items (status)')
  375. cursor.execute('CREATE UNIQUE INDEX items_type_value_idx ON items (type, value)')
  376. def _insert_generated_items(self, cursor):
  377. it = itertools.chain((cls.itemType, value, STATUS_TODO) for cls in self._itemClasses for value in cls.generate())
  378. while True:
  379. values = tuple(itertools.islice(it, 100000))
  380. if not values:
  381. break
  382. cursor.executemany('INSERT OR IGNORE INTO items (type, value, status) VALUES (?, ?, ?)', values)