Reduce collector memory use

components/collector/.vulture-whitelist.py

@@ -42,8 +42,8 @@
 lastBomImport  # unused variable (src/source_collectors/dependency_track/json_types.py:18)
 isLatest  # unused variable (src/source_collectors/dependency_track/json_types.py:21)
 latestVersion  # unused variable (src/source_collectors/dependency_track/json_types.py:29)
-directDependencies  # unused variable (src/source_collectors/dependency_track/json_types.py:39)
-repositoryMeta  # unused variable (src/source_collectors/dependency_track/json_types.py:40)
+directDependencies  # unused variable (src/source_collectors/dependency_track/json_types.py:38)
+repositoryMeta  # unused variable (src/source_collectors/dependency_track/json_types.py:39)
 latestVersion  # unused variable (src/source_collectors/dependency_track/security_warnings.py:15)
 vulnId  # unused variable (src/source_collectors/dependency_track/security_warnings.py:25)
 matrix  # unused variable (src/source_collectors/dependency_track/security_warnings.py:36)

components/collector/pyproject.toml

@@ -15,6 +15,7 @@ dependencies = [
     "aiohttp==3.14.5",
     "beautifulsoup4==4.15.0",
     "defusedxml==0.7.1",
+    "ijson==3.6.0",
     "packaging==26.3",
     "pymongo==4.19.0",
     "python-dateutil==3.0.0",

components/collector/src/source_collectors/dependency_track/base.py

@@ -6,18 +6,22 @@
 from itertools import batched
 from math import ceil
 from operator import itemgetter
-from typing import TYPE_CHECKING, Literal
+from typing import TYPE_CHECKING, Any, Literal
+
+import ijson  # type: ignore[import-untyped]  # ijson has no type information
 
 from base_collectors import TokenAuthenticationSourceCollector
 from collector_utilities.exceptions import CollectorError
 from collector_utilities.functions import add_query
 from collector_utilities.type import URL, Response
 from model import Entity, SourceResponses
 
+from .json_types import DependencyTrackComponent
+
 if TYPE_CHECKING:
-    from collections.abc import AsyncIterator, Iterable
+    from collections.abc import AsyncIterator, Awaitable, Callable, Iterable
 
-    from .json_types import DependencyTrackComponent, DependencyTrackProject
+    from .json_types import DependencyTrackProject
 
 
 class DependencyTrackComponentGraph:
@@ -88,44 +92,74 @@ async def _get_source_responses(self, *urls: URL) -> SourceResponses:
         """Extend to load multiple pages, if necessary."""
         responses = SourceResponses()
         for url in urls:
-            async for pages in self.__get_page_batches(url):
+            async for pages in self.__get_page_batches(url, self.__read_response):
                 responses.extend(SourceResponses(responses=pages))
         return responses
 
-    async def __get_page_batches(self, url: URL) -> AsyncIterator[list[Response]]:
-        """Yield all pages of the paginated URL, in batches, in page order.
+    async def _get_components(self, project_uuid: str) -> AsyncIterator[list[DependencyTrackComponent]]:
+        """Yield the components of the project, page by page, in page order.
+
+        Each component as returned by Dependency-Track contains the project, including the project's direct
+        dependencies, so the size of a page is quadratic in the number of direct dependencies of the project. To
+        not need memory for that, parse the components one by one while reading the response, and keep only the
+        attributes the collector needs.
+        """
+        components_api = URL(f"{await self._api_url()}/component/project/{project_uuid}")
+        async for pages in self.__get_page_batches(components_api, self.__read_components):
+            for page in pages:
+                yield page
+
+    async def __get_page_batches[T](
+        self, url: URL, read_page: Callable[[Response], Awaitable[T]]
+    ) -> AsyncIterator[list[T]]:
+        """Yield all pages of the paginated URL, read with the read page function, in batches, in page order.
 
         The first page has to be retrieved first, because its X-Total-Count header tells how many pages there are.
         The other pages are retrieved in batches of at most MAX_CONCURRENT_PAGES, because retrieving them one by one
         would need one round trip per PAGE_SIZE items, which times out for projects with many components. Batching
         also bounds the number of requests that are prepared at the same time, no matter how big the total count is.
         """
-        first_page = await self.__get_page(url, 1)  # Page numbers start at 1
-        page_count = ceil(int(first_page.headers.get("X-Total-Count", 0)) / self.PAGE_SIZE)
-        yield [first_page]
-        del first_page  # Free memory
+        first_response = await self.__get_page(url, 1)  # Page numbers start at 1
+        page_count = ceil(int(first_response.headers.get("X-Total-Count", 0)) / self.PAGE_SIZE)
+        yield [await read_page(first_response)]
+        del first_response  # Free memory
         for batch in batched(range(2, page_count + 1), self.MAX_CONCURRENT_PAGES, strict=False):
-            yield await asyncio.gather(*[self.__get_page(url, page_nr) for page_nr in batch])
+            yield await asyncio.gather(*[self.__get_and_read_page(url, page_nr, read_page) for page_nr in batch])
+
+    async def __get_and_read_page[T](self, url: URL, page_nr: int, read_page: Callable[[Response], Awaitable[T]]) -> T:
+        """Return one page of the paginated URL, read with the read page function."""
+        return await read_page(await self.__get_page(url, page_nr))
 
     async def __get_page(self, url: URL, page_nr: int) -> Response:
-        """Return one page of the paginated URL."""
+        """Return the response for one page of the paginated URL, without reading it."""
         paginated_url = add_query(url, f"pageSize={self.PAGE_SIZE}&pageNumber={page_nr}")
         response: Response = (await super()._get_source_responses(paginated_url))[0]
+        return response
+
+    @staticmethod
+    async def __read_response(response: Response) -> Response:
+        """Read the response and return it."""
         # Retrieving consecutive big responses without reading the response hangs the client, see
         # https://github.com/aio-libs/aiohttp/issues/2217
         await response.read()
         return response
 
-    async def _get_paginated_items(self, url: URL) -> AsyncIterator[list]:
-        """Yield the items of the paginated URL, page by page, in page order.
+    @classmethod
+    async def __read_components(cls, response: Response) -> list[DependencyTrackComponent]:
+        """Read the components from the response, one by one, keeping only the attributes the collector needs."""
+        return [cls.__slim_component(component) async for component in ijson.items_async(response.content, "item")]
 
-        Unlike _get_source_responses(), this doesn't keep all pages in memory. The response of a page can be garbage
-        collected as soon as the caller has processed its items, so callers can keep only the data they need, which
-        matters for projects with many components.
-        """
-        async for pages in self.__get_page_batches(url):
-            for page in pages:
-                yield await page.json(content_type=None)
+    @staticmethod
+    def __slim_component(component: dict[str, Any]) -> DependencyTrackComponent:
+        """Return a copy of the component with only the attributes the collector needs."""
+        slim_component = DependencyTrackComponent(name=component["name"], uuid=component["uuid"])
+        if "directDependencies" in component:
+            slim_component["directDependencies"] = component["directDependencies"]
+        if "repositoryMeta" in component:
+            slim_component["repositoryMeta"] = component["repositoryMeta"]
+        if "version" in component:
+            slim_component["version"] = component["version"]
+        return slim_component
 
     async def _get_projects_by_uuid(self) -> dict[str, DependencyTrackProject]:
         """Return a mapping of project UUIDs to projects."""

components/collector/src/source_collectors/dependency_track/dependencies.py

@@ -2,14 +2,14 @@
 
 from typing import TYPE_CHECKING
 
-from collector_utilities.type import URL
 from model import Entities, Entity, SourceResponses
 
 from .base import DependencyTrackComponentGraph, DependencyTrackLatestVersionStatusBase
-from .json_types import DependencyTrackComponent
 
 if TYPE_CHECKING:
-    from .json_types import DependencyTrackProject
+    from collector_utilities.type import URL
+
+    from .json_types import DependencyTrackComponent, DependencyTrackProject
 
 
 class DependencyTrackDependencies(DependencyTrackLatestVersionStatusBase):
@@ -22,43 +22,35 @@ async def _get_source_responses(self, *urls: URL) -> SourceResponses:
         that the responses for projects with many components don't need to be kept in memory. As the components are
         kept by the collector, the source responses are empty.
         """
-        api_url = await self._api_url()
-        self._components: list[tuple[DependencyTrackComponentGraph, list[DependencyTrackComponent]]] = []
+        self._components: list[
+            tuple[DependencyTrackProject, DependencyTrackComponentGraph, list[DependencyTrackComponent]]
+        ] = []
         for project_uuid, project in (await self._get_projects_by_uuid()).items():
             # Create one dependency graph per project so the root components can be looked up
             graph = DependencyTrackComponentGraph()
             components: list[DependencyTrackComponent] = []
-            async for page in self._get_paginated_items(URL(f"{api_url}/component/project/{project_uuid}")):
+            async for page in self._get_components(project_uuid):
                 graph.add_components(page)
-                components.extend(self._slim_component(component, project) for component in page)
-            self._components.append((graph, components))
+                for component in page:
+                    component.pop("directDependencies", None)  # Not needed once the component is in the graph
+                components.extend(page)
+            self._components.append((project, graph, components))
         return SourceResponses()
 
-    @staticmethod
-    def _slim_component(
-        component: DependencyTrackComponent, project: DependencyTrackProject
-    ) -> DependencyTrackComponent:
-        """Return a copy of the component with only the attributes needed to create the entity.
-
-        The copy refers to the project instead of having its own copy of the project, to save memory.
-        """
-        slim_component = DependencyTrackComponent(name=component["name"], project=project, uuid=component["uuid"])
-        if "repositoryMeta" in component:
-            slim_component["repositoryMeta"] = component["repositoryMeta"]
-        if "version" in component:
-            slim_component["version"] = component["version"]
-        return slim_component
-
     async def _parse_entities(self, responses: SourceResponses) -> Entities:
         """Parse the entities from the components."""
         entities = Entities()
-        for graph, components in self._components:
-            entities.extend([self._create_entity(component, graph) for component in components])
+        for project, graph, components in self._components:
+            entities.extend([self._create_entity(component, project, graph) for component in components])
         return entities
 
-    def _create_entity(self, component: DependencyTrackComponent, graph: DependencyTrackComponentGraph) -> Entity:
+    def _create_entity(
+        self,
+        component: DependencyTrackComponent,
+        project: DependencyTrackProject,
+        graph: DependencyTrackComponentGraph,
+    ) -> Entity:
         """Create an entity from the component."""
-        project = component["project"]
         current_version = component.get("version", "unknown")
         latest_version = component.get("repositoryMeta", {}).get("latestVersion", "unknown")
         return Entity(

components/collector/src/source_collectors/dependency_track/json_types.py

@@ -30,10 +30,9 @@ class DependencyTrackRepositoryMetaData(TypedDict):
 
 
 class DependencyTrackComponent(TypedDict):
-    """Component as returned by the Dependency-Track component API."""
+    """Component as returned by the Dependency-Track component API, with only the attributes the collector needs."""
 
     name: str
-    project: DependencyTrackProject
     uuid: str
     # Direct dependencies is a JSON string with a list of component identities, and null if the component has none
     directDependencies: NotRequired[str | None]

components/collector/src/source_collectors/dependency_track/security_warnings.py

@@ -52,12 +52,11 @@ async def _get_source_responses(self, *urls: URL) -> SourceResponses:
 
     async def _get_component_graph(self, project_uuid: str) -> DependencyTrackComponentGraph:
         """Return the dependency graph of the components of the project."""
-        components_api = URL(f"{await self._api_url()}/component/project/{project_uuid}")
         graph = DependencyTrackComponentGraph()
         # Add the components to the graph page by page, so that only the graph is kept in memory and not the
         # components. Don't add the responses to the source responses, so that only the findings are parsed into
         # entities.
-        async for components in self._get_paginated_items(components_api):
+        async for components in self._get_components(project_uuid):
             graph.add_components(components)
         return graph
 

components/collector/tests/source_collectors/dependency_track/base_test.py

@@ -2,24 +2,62 @@
 
 import json
 from typing import TYPE_CHECKING
+from unittest.mock import AsyncMock
 
-from source_collectors.dependency_track.json_types import (
-    DependencyTrackComponent,
-    DependencyTrackMetrics,
-    DependencyTrackProject,
-)
+from source_collectors.dependency_track.json_types import DependencyTrackMetrics, DependencyTrackProject
 
 from tests.source_collectors.source_collector_test_case import SourceCollectorTestCase
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from model.measurement import MetricMeasurement
 
 
+class StreamReaderStub:
+    """Stub for the aiohttp stream reader, to be able to stream the response content."""
+
+    def __init__(self, content: bytes) -> None:
+        self.content = content
+
+    async def read(self, size: int = -1) -> bytes:
+        """Return at most size bytes of the content, or all remaining content if size is negative."""
+        size = len(self.content) if size < 0 else size
+        chunk, self.content = self.content[:size], self.content[size:]
+        return chunk
+
+
 class DependencyTrackTestCase(SourceCollectorTestCase):
     """Base class for Dependency-Track collector unit tests."""
 
     SOURCE_TYPE = "dependency_track"
 
+    async def collect(self, **kwargs):
+        """Extend to give each response its own JSON, so the collector can either read or stream the response."""
+        if "get_request_side_effect" not in kwargs:
+            kwargs["get_request_side_effect"] = self.__get_responses(
+                kwargs.pop("get_request_json_return_value", None),
+                kwargs.pop("get_request_json_side_effect", None),
+                kwargs.pop("get_request_headers", None),
+            )
+        return await super().collect(**kwargs)
+
+    @staticmethod
+    def __get_responses(json_return_value, json_side_effect, headers) -> Callable[..., AsyncMock]:
+        """Return a get request side effect that creates a response for each request."""
+        json_items = iter(json_side_effect or [])
+
+        def get_response(*args, **kwargs) -> AsyncMock:  # noqa: ARG001
+            """Return a response with the next JSON item, if there are multiple, as JSON and as streamable content."""
+            json_item = json_return_value if json_side_effect is None else next(json_items)
+            response = AsyncMock()
+            response.json.return_value = json_item
+            response.content = StreamReaderStub(json.dumps(json_item).encode())
+            response.headers = headers or {}
+            return response
+
+        return get_response
+
     def setUp(self) -> None:
         """Extend to add the mandatory landing URL."""
         super().setUp()
@@ -44,15 +82,15 @@ def component(
         name: str = "component name",
         uuid: str = "component-uuid",
         direct_dependencies: list[str] | None = None,
-        project: DependencyTrackProject | None = None,
-    ) -> DependencyTrackComponent:
+        project: DependencyTrackProject | dict[str, object] | None = None,
+    ) -> dict[str, object]:
         """Create a Dependency-Track component fixture, with the UUIDs of its direct dependencies, if any."""
-        component = DependencyTrackComponent(
-            name=name,
-            project=self.projects()[0] if project is None else project,
-            uuid=uuid,
-            version="1.0",
-        )
+        component: dict[str, object] = {
+            "name": name,
+            "project": self.projects()[0] if project is None else project,  # Dependency-Track includes the project
+            "uuid": uuid,
+            "version": "1.0",
+        }
         if direct_dependencies is not None:
             component["directDependencies"] = json.dumps([{"uuid": child} for child in direct_dependencies])
         return component

components/collector/tests/source_collectors/dependency_track/test_dependencies.py

@@ -1,5 +1,6 @@
 """Unit tests for the Dependency-Track dependencies collector."""
 
+import json
 from typing import TYPE_CHECKING
 from unittest.mock import patch
 
@@ -10,15 +11,14 @@
 
 if TYPE_CHECKING:
     from model.measurement import MetricMeasurement
-    from source_collectors.dependency_track.json_types import DependencyTrackComponent
 
 
 class DependencyTrackDependenciesTest(DependencyTrackTestCase):
     """Unit tests for the Dependency-Track dependencies collector."""
 
     METRIC_TYPE = "dependencies"
 
-    def dependencies(self, latest_version: str) -> list[DependencyTrackComponent]:
+    def dependencies(self, latest_version: str) -> list[dict[str, object]]:
         """Create a list of dependencies as returned by Dependency-Track."""
         dependency = self.component()
         if latest_version:
@@ -44,7 +44,7 @@ def entities(self, latest_version: str, latest_version_status: str) -> list[dict
         """Create a list of expected entities."""
         return [self.entity(latest=latest_version, latest_version_status=latest_version_status)]
 
-    async def collect_components(self, components: list[DependencyTrackComponent]) -> MetricMeasurement:
+    async def collect_components(self, components: list[dict[str, object]]) -> MetricMeasurement:
         """Collect a measurement for the given components."""
         return await self.collect_measurement(get_request_json_side_effect=[self.projects(), components])
 
@@ -58,6 +58,18 @@ async def test_no_vulnerabilities(self):
         measurement = await self.collect_measurement(get_request_json_side_effect=[self.projects(), []])
         self.assert_measurement(measurement, value="0", entities=[])
 
+    async def test_project_with_many_direct_dependencies(self):
+        """Test that components are parsed correctly if the project included in each component is big.
+
+        Dependency-Track includes the project, with its direct dependencies, in each component. Make the project bigger
+        than the buffer used to read the response, so the components need to be parsed from multiple chunks.
+        """
+        direct_dependencies = json.dumps([{"uuid": f"direct-dependency-uuid-{index}"} for index in range(5000)])
+        project = dict(self.projects()[0], directDependencies=direct_dependencies)
+        components = [self.component(project=project), self.component(name="other", uuid="other-uuid", project=project)]
+        measurement = await self.collect_components(components)
+        self.assert_measurement(measurement, value="2", entities=[self.entity(), self.entity("other", "other-uuid")])
+
     async def test_updateable_dependencies(self):
         """Test one project with a component that can be updated."""
         measurement = await self.collect_measurement(