Custom memusage extension (#257)

* Custom extension to get memory usage

* Ignore bandit warning

* Simplify

* Add handler attributes

* Basic tests for memusage extension

* Simplify extension and tests

* More tests

* Make pylint happy

* Test get_virtual_size

* More tests

* Update readme

* Mention package requirement

* Rename module

Author: elacuesta

README.md

@@ -826,6 +826,32 @@ for a list of the accepted events and the arguments passed to their handlers.
   images, scripts, stylesheets, etc are not seen by Scrapy.
 
 
+## Memory usage extension
+
+The default Scrapy memory usage extension
+(`scrapy.extensions.memusage.MemoryUsage`) does not include the memory used by
+Playwright because the browser is launched as a separate process. The
+scrapy-playwright package provides a replacement extension which also considers
+the memory used by Playwright. This extension needs the
+[`psutil`](https://pypi.org/project/psutil/) package to work.
+
+Update the [EXTENSIONS](https://docs.scrapy.org/en/latest/topics/settings.html#std-setting-EXTENSIONS)
+setting to disable the built-in Scrapy extension and replace it with the one
+from the scrapy-playwright package:
+
+```python
+# settings.py
+EXTENSIONS = {
+    "scrapy.extensions.memusage.MemoryUsage": None,
+    "scrapy_playwright.memusage.ScrapyPlaywrightMemoryUsageExtension": 0,
+}
+```
+
+Refer to the
+[upstream docs](https://docs.scrapy.org/en/latest/topics/extensions.html#module-scrapy.extensions.memusage)
+for more information about supported settings.
+
+
 ## Examples
 
 **Click on a link, save the resulting page as PDF**

scrapy_playwright/handler.py

@@ -12,6 +12,7 @@
     Download,
     Error as PlaywrightError,
     Page,
+    Playwright as AsyncPlaywright,
     PlaywrightContextManager,
     Request as PlaywrightRequest,
     Response as PlaywrightResponse,
@@ -95,6 +96,9 @@ def from_settings(cls, settings: Settings) -> "Config":
 
 
 class ScrapyPlaywrightDownloadHandler(HTTPDownloadHandler):
+    playwright_context_manager: Optional[PlaywrightContextManager] = None
+    playwright: Optional[AsyncPlaywright] = None
+
     def __init__(self, crawler: Crawler) -> None:
         super().__init__(settings=crawler.settings, crawler=crawler)
         verify_installed_reactor("twisted.internet.asyncioreactor.AsyncioSelectorReactor")
@@ -294,8 +298,10 @@ async def _close(self) -> None:
         if hasattr(self, "browser"):
             logger.info("Closing browser")
             await self.browser.close()
-        await self.playwright_context_manager.__aexit__()
-        await self.playwright.stop()
+        if self.playwright_context_manager:
+            await self.playwright_context_manager.__aexit__()
+        if self.playwright:
+            await self.playwright.stop()
 
     def download_request(self, request: Request, spider: Spider) -> Deferred:
         if request.meta.get("playwright"):

scrapy_playwright/memusage.py

@@ -0,0 +1,56 @@
+from contextlib import suppress
+from importlib import import_module
+from typing import List
+
+from scrapy.exceptions import NotConfigured
+from scrapy.extensions.memusage import MemoryUsage
+
+from scrapy_playwright.handler import ScrapyPlaywrightDownloadHandler, logger
+
+
+_MIB_FACTOR = 1024**2
+
+
+class ScrapyPlaywrightMemoryUsageExtension(MemoryUsage):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        try:
+            self.psutil = import_module("psutil")
+        except ImportError as exc:
+            raise NotConfigured("The psutil module is not available") from exc
+
+    def _get_main_process_ids(self) -> List[int]:
+        try:
+            return [
+                handler.playwright_context_manager._connection._transport._proc.pid
+                for handler in self.crawler.engine.downloader.handlers._handlers.values()
+                if isinstance(handler, ScrapyPlaywrightDownloadHandler)
+                and handler.playwright_context_manager
+            ]
+        except Exception:
+            return []
+
+    def _get_descendant_processes(self, process) -> list:
+        children = process.children()
+        result = children.copy()
+        for child in children:
+            result.extend(self._get_descendant_processes(child))
+        return result
+
+    def _get_total_playwright_process_memory(self) -> int:
+        process_list = [self.psutil.Process(pid) for pid in self._get_main_process_ids()]
+        for proc in process_list.copy():
+            process_list.extend(self._get_descendant_processes(proc))
+        total_process_size = 0
+        for proc in process_list:
+            with suppress(Exception):  # might fail if the process exited in the meantime
+                total_process_size += proc.memory_info().rss
+        logger.debug(
+            "Total Playwright process memory: %i Bytes (%i MiB)",
+            total_process_size,
+            total_process_size / _MIB_FACTOR,
+        )
+        return total_process_size
+
+    def get_virtual_size(self) -> int:
+        return super().get_virtual_size() + self._get_total_playwright_process_memory()

tests/tests_asyncio/test_extensions.py

@@ -0,0 +1,89 @@
+from asyncio.subprocess import Process as AsyncioProcess
+from unittest import IsolatedAsyncioTestCase
+from unittest.mock import MagicMock, patch
+
+import pytest
+from playwright.async_api import PlaywrightContextManager
+from scrapy.exceptions import NotConfigured
+from scrapy.extensions.memusage import MemoryUsage
+
+from scrapy_playwright.memusage import ScrapyPlaywrightMemoryUsageExtension
+from scrapy_playwright.handler import ScrapyPlaywrightDownloadHandler
+
+
+SCHEMA_PID_MAP = {"http": 123, "https": 456}
+
+
+def mock_crawler_with_handlers() -> dict:
+    handlers = {}
+    for schema, pid in SCHEMA_PID_MAP.items():
+        process = MagicMock()
+        process.pid = pid
+        handlers[schema] = MagicMock(spec=ScrapyPlaywrightDownloadHandler)
+        handlers[schema].playwright_context_manager._connection._transport._proc = process
+    crawler = MagicMock()
+    crawler.engine.downloader.handlers._handlers = handlers
+    return crawler
+
+
+def raise_import_error(*args, **kwargs):
+    raise ImportError
+
+
+class MockMemoryInfo:
+    rss = 999
+
+
+@patch("scrapy.extensions.memusage.MailSender")
+class TestMemoryUsageExtension(IsolatedAsyncioTestCase):
+    async def test_process_availability(self, _MailSender):
+        """The main node process should be accessible from the context manager"""
+        ctx_manager = PlaywrightContextManager()
+        await ctx_manager.start()
+        assert isinstance(ctx_manager._connection._transport._proc, AsyncioProcess)
+        await ctx_manager.__aexit__()
+
+    @patch("scrapy_playwright.memusage.import_module", side_effect=raise_import_error)
+    async def test_psutil_not_available_extension_disabled(self, _import_module, _MailSender):
+        crawler = MagicMock()
+        with pytest.raises(NotConfigured):
+            ScrapyPlaywrightMemoryUsageExtension(crawler)
+
+    async def test_get_process_ids_ok(self, _MailSender):
+        crawler = mock_crawler_with_handlers()
+        extension = ScrapyPlaywrightMemoryUsageExtension(crawler)
+        assert extension._get_main_process_ids() == list(SCHEMA_PID_MAP.values())
+
+    async def test_get_process_ids_error(self, _MailSender):
+        crawler = mock_crawler_with_handlers()
+        crawler.engine.downloader.handlers._handlers = MagicMock()
+        crawler.engine.downloader.handlers._handlers.values.side_effect = raise_import_error
+        extension = ScrapyPlaywrightMemoryUsageExtension(crawler)
+        assert extension._get_main_process_ids() == []
+
+    async def test_get_descendant_processes(self, _MailSender):
+        p1 = MagicMock()
+        p2 = MagicMock()
+        p3 = MagicMock()
+        p4 = MagicMock()
+        p2.children.return_value = [p3, p4]
+        p1.children.return_value = [p2]
+        crawler = MagicMock()
+        extension = ScrapyPlaywrightMemoryUsageExtension(crawler)
+        assert extension._get_descendant_processes(p1) == [p2, p3, p4]
+
+    async def test_get_total_process_size(self, _MailSender):
+        crawler = MagicMock()
+        extension = ScrapyPlaywrightMemoryUsageExtension(crawler)
+        extension.psutil = MagicMock()
+        extension.psutil.Process.return_value.memory_info.return_value = MockMemoryInfo()
+        extension._get_main_process_ids = MagicMock(return_value=[1, 2, 3])
+        expected_size = MockMemoryInfo().rss * len(extension._get_main_process_ids())
+        assert extension._get_total_playwright_process_memory() == expected_size
+
+    async def test_get_virtual_size_sum(self, _MailSender):
+        crawler = MagicMock()
+        extension = ScrapyPlaywrightMemoryUsageExtension(crawler)
+        parent_cls_extension = MemoryUsage(crawler)
+        extension._get_total_playwright_process_memory = MagicMock(return_value=123)
+        assert extension.get_virtual_size() == parent_cls_extension.get_virtual_size() + 123

tox.ini

@@ -6,6 +6,7 @@ deps =
     pytest==7.4.0
     pytest_cov==4.1.0
     pytest_twisted==1.14
+    psutil==5.9.7
 commands =
     playwright install
     py.test -vv --reactor=asyncio \