Parallelize compute component comparisons (#1046)

Author: joseph-sentry

rollouts/__init__.py

@@ -14,3 +14,5 @@
 )
 
 NEW_TA_TASKS = Feature("new_ta_tasks")
+
+PARALLEL_COMPONENT_COMPARISON = Feature("parallel_component_comparison")

services/comparison_utils.py

@@ -0,0 +1,37 @@
+import sentry_sdk
+from shared.reports.readonly import ReadOnlyReport
+
+from database.models import CompareCommit
+from services.comparison import ComparisonContext, ComparisonProxy
+from services.comparison.types import Comparison, FullCommit
+from services.report import ReportService
+
+
+@sentry_sdk.trace
+def get_comparison_proxy(
+    comparison: CompareCommit,
+    report_service: ReportService,
+):
+    compare_commit = comparison.compare_commit
+    base_commit = comparison.base_commit
+
+    base_report = report_service.get_existing_report_for_commit(
+        base_commit, report_class=ReadOnlyReport
+    )
+    compare_report = report_service.get_existing_report_for_commit(
+        compare_commit, report_class=ReadOnlyReport
+    )
+    # No access to the PR so we have to assume the base commit did not need
+    # to be adjusted.
+    patch_coverage_base_commitid = base_commit.commitid
+    return ComparisonProxy(
+        Comparison(
+            head=FullCommit(commit=compare_commit, report=compare_report),
+            project_coverage_base=FullCommit(commit=base_commit, report=base_report),
+            patch_coverage_base_commitid=patch_coverage_base_commitid,
+            enriched_pull=None,
+        ),
+        context=ComparisonContext(
+            gh_app_installation_name=report_service.gh_app_installation_name
+        ),
+    )

tasks/compute_comparison.py

@@ -3,10 +3,10 @@
 
 import sentry_sdk
 from asgiref.sync import async_to_sync
+from celery import group
 from shared.celery_config import compute_comparison_task_name
 from shared.components import Component
 from shared.helpers.flag import Flag
-from shared.reports.readonly import ReadOnlyReport
 from shared.torngit.exceptions import TorngitRateLimitError
 from shared.yaml import UserYaml
 
@@ -16,12 +16,14 @@
 from database.models.reports import ReportLevelTotals, RepositoryFlag
 from helpers.comparison import minimal_totals
 from helpers.github_installation import get_installation_name_for_owner_for_task
+from rollouts import PARALLEL_COMPONENT_COMPARISON
 from services.archive import ArchiveService
-from services.comparison import ComparisonContext, ComparisonProxy, FilteredComparison
-from services.comparison.types import Comparison, FullCommit
+from services.comparison import ComparisonProxy, FilteredComparison
+from services.comparison_utils import get_comparison_proxy
 from services.report import ReportService
 from services.yaml import get_current_yaml, get_repo_yaml
 from tasks.base import BaseCodecovTask
+from tasks.compute_component_comparison import compute_component_comparison_task
 
 log = logging.getLogger(__name__)
 
@@ -54,10 +56,11 @@ def run_impl(
         installation_name_to_use = get_installation_name_for_owner_for_task(
             self.name, repo.owner
         )
-
-        comparison_proxy = self.get_comparison_proxy(
-            comparison, current_yaml, installation_name_to_use
+        report_service = ReportService(
+            current_yaml, gh_app_installation_name=installation_name_to_use
         )
+
+        comparison_proxy = get_comparison_proxy(comparison, report_service)
         if not comparison_proxy.has_head_report():
             comparison.error = CompareCommitError.missing_head_report.value
             comparison.state = CompareCommitState.error.value
@@ -241,10 +244,31 @@ def compute_component_comparisons(
                 component_count=len(components),
             ),
         )
-        for component in components:
-            self.compute_component_comparison(
-                db_session, comparison, comparison_proxy, component
-            )
+        if PARALLEL_COMPONENT_COMPARISON.check_value(
+            comparison.compare_commit.repoid, default=False
+        ):
+            self.parallel_compute_component_comparison(comparison.id, components)
+        else:
+            for component in components:
+                self.compute_component_comparison(
+                    db_session, comparison, comparison_proxy, component
+                )
+
+    @sentry_sdk.trace
+    def parallel_compute_component_comparison(
+        self,
+        comparison_id: int,
+        components: list[Component],
+    ):
+        task_group = group(
+            [
+                compute_component_comparison_task.s(
+                    comparison_id, component.component_id
+                )
+                for component in components
+            ]
+        )
+        task_group.apply_async()
 
     def compute_component_comparison(
         self,
@@ -288,38 +312,6 @@ def compute_component_comparison(
         db_session.add(component_comparison)
         db_session.flush()
 
-    @sentry_sdk.trace
-    def get_comparison_proxy(
-        self, comparison, current_yaml, installation_name_to_use: str | None = None
-    ):
-        compare_commit = comparison.compare_commit
-        base_commit = comparison.base_commit
-        report_service = ReportService(
-            current_yaml, gh_app_installation_name=installation_name_to_use
-        )
-        base_report = report_service.get_existing_report_for_commit(
-            base_commit, report_class=ReadOnlyReport
-        )
-        compare_report = report_service.get_existing_report_for_commit(
-            compare_commit, report_class=ReadOnlyReport
-        )
-        # No access to the PR so we have to assume the base commit did not need
-        # to be adjusted.
-        patch_coverage_base_commitid = base_commit.commitid
-        return ComparisonProxy(
-            Comparison(
-                head=FullCommit(commit=compare_commit, report=compare_report),
-                project_coverage_base=FullCommit(
-                    commit=base_commit, report=base_report
-                ),
-                patch_coverage_base_commitid=patch_coverage_base_commitid,
-                enriched_pull=None,
-            ),
-            context=ComparisonContext(
-                gh_app_installation_name=installation_name_to_use
-            ),
-        )
-
     @sentry_sdk.trace
     def store_results(self, comparison, impacted_files):
         repository = comparison.compare_commit.repository

tasks/compute_component_comparison.py

@@ -0,0 +1,96 @@
+from asgiref.sync import async_to_sync
+from shared.components import Component
+from shared.yaml import UserYaml
+from sqlalchemy.orm import Session
+
+from app import celery_app
+from database.models import CompareCommit, CompareComponent
+from helpers.github_installation import get_installation_name_for_owner_for_task
+from services.comparison import ComparisonProxy, FilteredComparison
+from services.comparison_utils import get_comparison_proxy
+from services.report import ReportService
+from services.yaml import get_current_yaml, get_repo_yaml
+from tasks.base import BaseCodecovTask
+
+
+def compute_component_comparison(
+    db_session: Session,
+    comparison: CompareCommit,
+    comparison_proxy: ComparisonProxy,
+    component: Component,
+):
+    component_comparison = (
+        db_session.query(CompareComponent)
+        .filter_by(
+            commit_comparison_id=comparison.id,
+            component_id=component.component_id,
+        )
+        .first()
+    )
+    if not component_comparison:
+        component_comparison = CompareComponent(
+            commit_comparison=comparison,
+            component_id=component.component_id,
+        )
+
+    # filter comparison by component
+    head_report = comparison_proxy.comparison.head.report
+    flags = component.get_matching_flags(head_report.flags.keys())
+    filtered: FilteredComparison = comparison_proxy.get_filtered_comparison(
+        flags=flags, path_patterns=component.paths
+    )
+
+    # component comparison totals
+    component_comparison.base_totals = (
+        filtered.project_coverage_base.report.totals.asdict()
+    )
+    component_comparison.head_totals = filtered.head.report.totals.asdict()
+    diff = comparison_proxy.get_diff()
+    if diff:
+        patch_totals = filtered.head.report.apply_diff(diff)
+        if patch_totals:
+            component_comparison.patch_totals = patch_totals.asdict()
+
+    db_session.add(component_comparison)
+    db_session.flush()
+
+
+class ComputeComponentComparisonTask(BaseCodecovTask):
+    def run_impl(
+        self,
+        db_session: Session,
+        comparison_id: int,
+        component_id: str,
+        *args,
+        **kwargs,
+    ):
+        comparison: CompareCommit = db_session.query(CompareCommit).get(comparison_id)
+        repo = comparison.compare_commit.repository
+        current_yaml = get_repo_yaml(repo)
+        installation_name_to_use = get_installation_name_for_owner_for_task(
+            self.name, repo.owner
+        )
+        report_service = ReportService(
+            current_yaml, gh_app_installation_name=installation_name_to_use
+        )
+        comparison_proxy = get_comparison_proxy(comparison, report_service)
+        head_commit = comparison_proxy.comparison.head.commit
+
+        yaml: UserYaml = async_to_sync(get_current_yaml)(
+            head_commit, comparison_proxy.repository_service
+        )
+
+        components = yaml.get_components()
+
+        component_dict = {c.component_id: c for c in components}
+        compute_component_comparison(
+            db_session, comparison, comparison_proxy, component_dict[component_id]
+        )
+
+
+RegisteredComputeComponentComparisonTask = celery_app.register_task(
+    ComputeComponentComparisonTask()
+)
+compute_component_comparison_task = celery_app.tasks[
+    RegisteredComputeComponentComparisonTask.name
+]