Enhance grouping algo

Signed-off-by: Tushar Goel <tushar.goel.dav@gmail.com>

Author: TG1999

vulnerabilities/pipelines/v2_improvers/mark_unfurl_version_range.py

@@ -9,6 +9,7 @@
 
 from django.db import transaction
 from django.db.models import Exists
+from django.db.models import Min
 from django.db.models import OuterRef
 from django.db.models import Q
 
@@ -51,7 +52,7 @@ def mark_all_impacts_unfurled(self):
             impacted_packages=impacted_packages,
         )
 
-        for advisory_id in advisories_qs.iterator(chunk_size=100):
+        for advisory_id in advisories_qs.iterator(chunk_size=1000):
             batch.append(advisory_id)
 
             if len(batch) >= batch_size:
@@ -104,9 +105,12 @@ def latest_advisories_with_all_impacts_unfurled_attempted(
             _all_impacts_unfurled_successfully=False,
             is_latest=True,
         )
-        .annotate(has_unattempted_impacts=Exists(impacts_not_attempted))
+        .annotate(
+            has_unattempted_impacts=Exists(impacts_not_attempted),
+            first_base_purl=Min("impacted_packages__base_purl"),
+        )
         .filter(has_unattempted_impacts=False)
-        .order_by("_all_impacts_unfurled", "datasource_id")
+        .order_by("_all_impacts_unfurled", "first_base_purl")
         .values_list("id", flat=True)
     )
 

vulnerabilities/pipes/group_advisories.py

@@ -7,6 +7,8 @@
 # See https://aboutcode.org for more information about nexB OSS projects.
 #
 
+import hashlib
+import json
 from collections import defaultdict
 from typing import List
 
@@ -16,6 +18,7 @@
 from vulnerabilities.models import AdvisorySetMember
 from vulnerabilities.models import AdvisoryV2
 from vulnerabilities.models import Group
+from vulnerabilities.utils import normalize_list
 
 
 @transaction.atomic
@@ -138,25 +141,137 @@ def group_advisory_for_package(package, logger=None):
         return
 
 
+def compute_advisory_content_hash(adv, version_less_purl_str: str):
+    """
+    Compute a content hash for an advisory.
+
+    ``version_less_purl_str`` is pre-computed by the caller once and reused
+    across all advisories — avoids re-constructing PackageURL N times.
+    The impacted_packages relation must already be prefetched with
+    ``affecting_packages`` and ``fixed_by_packages`` before calling this.
+    """
+    affected = []
+    fixed = []
+
+    for impact in adv.impacted_packages.all():
+        if impact.base_purl != version_less_purl_str:
+            continue
+        for pkg in impact.affecting_packages.all():
+            if pkg.package_url:
+                affected.append(pkg.package_url)
+        for pkg in impact.fixed_by_packages.all():
+            if pkg.package_url:
+                fixed.append(pkg.package_url)
+
+    normalized_data = {
+        "affected_packages": normalize_list(affected),
+        "fixed_packages": normalize_list(fixed),
+    }
+    normalized_json = json.dumps(normalized_data, separators=(",", ":"), sort_keys=True)
+    return hashlib.sha256(normalized_json.encode("utf-8")).hexdigest()
+
+
+def get_merged_identifier_groups(advisories, alias_map: dict):
+    """
+    Merge advisories based on shared advisory_id or alias.
+
+    ``alias_map`` is a dict[adv.id -> list[AdvisoryAlias]] pre-built by the
+    caller from a single bulk query — eliminates per-advisory alias lookups.
+
+    Uses a union-find (DSU) structure instead of the original O(n²) list-scan
+    merge, reducing merge cost to O(n·α(n)).
+    """
+    from vulnerabilities.models import Group
+
+    advisories = list(advisories)
+    if not advisories:
+        return []
+
+    parent = list(range(len(advisories)))
+    rank = [0] * len(advisories)
+
+    def find(x):
+        while parent[x] != x:
+            parent[x] = parent[parent[x]]
+            x = parent[x]
+        return x
+
+    def union(x, y):
+        rx, ry = find(x), find(y)
+        if rx == ry:
+            return
+        if rank[rx] < rank[ry]:
+            rx, ry = ry, rx
+        parent[ry] = rx
+        if rank[rx] == rank[ry]:
+            rank[rx] += 1
+
+    identifier_to_idx: dict[str, int] = {}
+
+    for i, adv in enumerate(advisories):
+        identifiers = [adv.advisory_id] + [alias.alias for alias in alias_map.get(adv.id, [])]
+        for ident in identifiers:
+            if ident in identifier_to_idx:
+                union(i, identifier_to_idx[ident])
+            else:
+                identifier_to_idx[ident] = i
+
+    root_to_group: dict[int, list] = defaultdict(list)
+    for i, adv in enumerate(advisories):
+        root_to_group[find(i)].append(adv)
+
+    final_groups: list[Group] = []
+
+    for group_members in root_to_group.values():
+        aliases = set()
+        for adv in group_members:
+            aliases.update(alias_map.get(adv.id, []))
+
+        primary = max(
+            group_members,
+            key=lambda a: a.precedence if a.precedence is not None else -1,
+        )
+        secondaries = [a for a in group_members if a is not primary]
+        final_groups.append(Group(aliases=aliases, primary=primary, secondaries=secondaries))
+
+    return final_groups
+
+
 def merge_advisories(advisories, package):
     """
-    Merge advisories based on their content hash and identifiers.
+    Merge advisories based on content hash and identifiers.
+
+    Builds the alias map once up-front from the already-prefetched queryset
+    so every downstream call shares a single in-memory dict.
     """
-    from vulnerabilities.utils import compute_advisory_content_hash
-    from vulnerabilities.utils import get_merged_identifier_groups
+    from packageurl import PackageURL
 
     advisories = list(advisories)
+    if not advisories:
+        return []
+
+    version_less_purl_str = str(
+        PackageURL(
+            type=package.type,
+            namespace=package.namespace,
+            name=package.name,
+            qualifiers=package.qualifiers,
+            subpath=package.subpath,
+        )
+    )
 
-    content_hash_map = defaultdict(list)
+    alias_map: dict[int, list] = defaultdict(list)
+    for adv in advisories:
+        alias_map[adv.id] = list(adv.aliases.all())
 
+    content_hash_map: dict[str, list] = defaultdict(list)
     for adv in advisories:
-        content_hash = compute_advisory_content_hash(adv, package)
+        content_hash = compute_advisory_content_hash(adv, version_less_purl_str)
         content_hash_map[content_hash].append(adv)
 
-    final_groups: List[Group] = []
-
+    final_groups: list[Group] = []
     for group in content_hash_map.values():
-        groups = get_merged_identifier_groups(group)
+        groups = get_merged_identifier_groups(group, alias_map)
         final_groups.extend(groups)
 
     return final_groups

vulnerabilities/tests/test_advisory_merge.py

@@ -8,6 +8,7 @@
 #
 
 import hashlib
+from collections import defaultdict
 
 import pytest
 
@@ -19,10 +20,10 @@
 from vulnerabilities.models import ImpactedPackage
 from vulnerabilities.models import PackageV2
 from vulnerabilities.pipes.group_advisories import delete_and_save_advisory_set
+from vulnerabilities.pipes.group_advisories import get_merged_identifier_groups
 from vulnerabilities.pipes.group_advisories import merge_advisories
 from vulnerabilities.utils import compute_advisory_content_hash
 from vulnerabilities.utils import get_advisories_from_groups
-from vulnerabilities.utils import get_merged_identifier_groups
 
 
 @pytest.mark.django_db
@@ -91,7 +92,11 @@ def test_identifier_merging(self):
         adv1.aliases.add(alias)
         adv2.aliases.add(alias)
 
-        groups = get_merged_identifier_groups([adv1, adv2])
+        alias_map: dict[int, list] = defaultdict(list)
+        for adv in [adv1, adv2]:
+            alias_map[adv.id] = list(adv.aliases.all())
+
+        groups = get_merged_identifier_groups([adv1, adv2], alias_map=alias_map)
 
         assert len(groups) == 1
         identifiers, primary, secondary = groups[0]
@@ -112,7 +117,11 @@ def test_transitive_merge(self):
         a2.aliases.add(alias_2)
         a3.aliases.add(alias_2)
 
-        groups = get_merged_identifier_groups([a1, a2, a3])
+        alias_map: dict[int, list] = defaultdict(list)
+        for adv in [a1, a2, a3]:
+            alias_map[adv.id] = list(adv.aliases.all())
+
+        groups = get_merged_identifier_groups([a1, a2, a3], alias_map=alias_map)
 
         assert len(groups) == 1
 
@@ -125,7 +134,11 @@ def test_primary_selection_by_precedence(self):
         a1.aliases.add(alias_1)
         a2.aliases.add(alias_1)
 
-        groups = get_merged_identifier_groups([a1, a2])
+        alias_map: dict[int, list] = defaultdict(list)
+        for adv in [a1, a2]:
+            alias_map[adv.id] = list(adv.aliases.all())
+
+        groups = get_merged_identifier_groups([a1, a2], alias_map=alias_map)
         _, primary, _ = groups[0]
 
         assert primary == a2
@@ -134,7 +147,11 @@ def test_get_advisories_from_groups(self):
         adv = self.create_advisory("GHSA-ABC-123", ["1.0"])
         adv.aliases.create(alias="CVE-999")
 
-        groups = get_merged_identifier_groups([adv])
+        alias_map: dict[int, list] = defaultdict(list)
+        for adv in [adv]:
+            alias_map[adv.id] = list(adv.aliases.all())
+
+        groups = get_merged_identifier_groups([adv], alias_map=alias_map)
         result = get_advisories_from_groups(groups)
 
         assert result[0].identifier == "GHSA-ABC-123"