Update community caches in parallel

Refactor task for updating chunked community packages to update each
community in parallel. Introduce a new task which is spawned by the main
update task to fire on each community.

Remove function responsible for updating all community caches and
introduce a new function doing the same thing but for a single
community, taking the community as a parameter.

Update cache tests because the old cache update function was removed.

Add update_single_community_cache to celery tests.

Add a create_batch method to CommunityFactory to help with testing.

Add test for the newly updated celery task to assert that the child
tasks are indeed spawned per community.

Refs. TS-2378

Author: Roffenlund

django/thunderstore/community/factories.py

@@ -22,6 +22,12 @@ class Meta:
     name = factory.Sequence(lambda n: f"TestCommunity{n}")
     identifier = factory.Sequence(lambda n: f"test-community-{n}")
 
+    @classmethod
+    def create_batch(cls, size=1, **kwargs):
+        return [
+            cls(**{**kwargs, "identifier": f"test-community-{i}"}) for i in range(size)
+        ]
+
 
 class PackageCategoryFactory(DjangoModelFactory):
     class Meta:

django/thunderstore/core/tests/test_celery.py

@@ -32,6 +32,7 @@ def test_task():
     "thunderstore.core.tasks.celery_post",
     "thunderstore.cache.tasks.invalidate_cache",
     "thunderstore.repository.tasks.update_api_caches",
+    "thunderstore.repository.tasks.update_single_community_cache",
     "thunderstore.usermedia.tasks.celery_cleanup_expired_uploads",
     "thunderstore.schema_import.tasks.sync_ecosystem_schema",
     "thunderstore.repository.tasks.files.extract_package_version_file_tree",

django/thunderstore/repository/api/v1/tasks.py

@@ -32,11 +32,9 @@ def update_api_v1_indexes() -> None:
     APIV1PackageCache.drop_stale_cache()
 
 
-def update_api_v1_chunked_package_caches() -> None:
-    for community in Community.objects.iterator():
-        try:
-            APIV1ChunkedPackageCache.update_for_community(community)
-        except Exception as e:  # pragma: no cover
-            capture_exception(e)
-
+def update_api_v1_chunked_package_cache_for_community(community: Community) -> None:
+    try:
+        APIV1ChunkedPackageCache.update_for_community(community)
+    except Exception as e:  # pragma: no cover
+        capture_exception(e)
     APIV1ChunkedPackageCache.drop_stale_cache()

django/thunderstore/repository/api/v1/tests/test_caches.py

@@ -15,7 +15,7 @@
 from thunderstore.community.models import Community, CommunitySite, PackageListing
 from thunderstore.repository.api.v1.tasks import (
     update_api_v1_caches,
-    update_api_v1_chunked_package_caches,
+    update_api_v1_chunked_package_cache_for_community,
 )
 from thunderstore.repository.models import APIV1ChunkedPackageCache, APIV1PackageCache
 from thunderstore.repository.models.cache import get_package_listing_chunk
@@ -155,7 +155,7 @@ def test_api_v1_chunked_package_cache__builds_index_and_chunks(
     PackageListingFactory(community_=community)
     assert APIV1ChunkedPackageCache.get_latest_for_community(community) is None
 
-    update_api_v1_chunked_package_caches()
+    update_api_v1_chunked_package_cache_for_community(community)
     cache = APIV1ChunkedPackageCache.get_latest_for_community(community)
     assert cache is not None
     assert cache.index.data_url.startswith(settings.AWS_S3_ENDPOINT_URL)
@@ -171,10 +171,12 @@ def test_api_v1_chunked_package_cache__drops_stale_caches() -> None:
     """
     Caches are currently only soft deleted.
     """
-    PackageListingFactory()
+    listing = PackageListingFactory()
+    community = listing.community
+
     assert not APIV1ChunkedPackageCache.objects.exists()
 
-    update_api_v1_chunked_package_caches()
+    update_api_v1_chunked_package_cache_for_community(community)
     first_cache = APIV1ChunkedPackageCache.objects.get()
     assert not first_cache.is_deleted
 
@@ -183,7 +185,7 @@ def test_api_v1_chunked_package_cache__drops_stale_caches() -> None:
     assert not first_cache.is_deleted
 
     # Two caches exists, but neither is beyond the cutoff period.
-    update_api_v1_chunked_package_caches()
+    update_api_v1_chunked_package_cache_for_community(community)
     APIV1ChunkedPackageCache.drop_stale_cache()
     second_cache = APIV1ChunkedPackageCache.get_latest_for_community(
         first_cache.community,

django/thunderstore/repository/tasks/caches.py

@@ -1,14 +1,19 @@
-from celery import shared_task  # type: ignore
+import logging
 
+from celery import group, shared_task  # type: ignore
+
+from thunderstore.community.models import Community
 from thunderstore.core.settings import CeleryQueues
 from thunderstore.repository.api.experimental.views.package_index import (
     update_api_experimental_package_index,
 )
 from thunderstore.repository.api.v1.tasks import (
     update_api_v1_caches,
-    update_api_v1_chunked_package_caches,
+    update_api_v1_chunked_package_cache_for_community,
 )
 
+logger = logging.getLogger(__name__)
+
 
 @shared_task(
     name="thunderstore.repository.tasks.update_api_caches",
@@ -28,11 +33,37 @@ def update_experimental_package_index():
     update_api_experimental_package_index()
 
 
+@shared_task(
+    name="thunderstore.repository.tasks.update_single_community_cache",
+    queue=CeleryQueues.BackgroundLongRunning,
+    expires=60 * 5,
+    soft_time_limit=60 * 15,
+)
+def update_single_community_cache(community_pk: int):
+    """Update the package cache for a single community."""
+
+    community = Community.objects.get(pk=community_pk)
+    update_api_v1_chunked_package_cache_for_community(community)
+
+
 @shared_task(
     name="thunderstore.repository.tasks.update_chunked_package_caches",
     queue=CeleryQueues.BackgroundLongRunning,
     soft_time_limit=60 * 60 * 23,
     time_limit=60 * 60 * 24,
 )
 def update_chunked_community_package_caches():
-    update_api_v1_chunked_package_caches()
+    """Update the package caches for all communities in parallel."""
+
+    try:
+        communities = Community.objects.values_list("pk", flat=True).iterator()
+
+        task_group = group(
+            update_single_community_cache.s(community_pk)
+            for community_pk in communities
+        )
+
+        result = task_group.apply_async()
+        return result
+    except Exception as e:  # pragma: no cover
+        logger.error(f"Failed starting parallel update of community caches. Error: {e}")

django/thunderstore/repository/tasks/tests/test_caches.py

@@ -0,0 +1,24 @@
+import pytest
+
+from thunderstore.community.factories import CommunityFactory
+from thunderstore.repository.tasks import update_chunked_community_package_caches
+
+
+@pytest.mark.django_db
+def test_update_chunked_community_package_caches_in_parallel():
+    """
+    Test that the celery tasks creates sub-tasks for each community and they
+    run successfully.
+    """
+
+    size = 5
+    CommunityFactory.create_batch(size=size)
+    result = update_chunked_community_package_caches.apply_async()
+    result = result.get()
+    sub_task_results = result.children
+
+    assert result.successful()
+    for sub_task in sub_task_results:
+        assert sub_task.successful()
+
+    assert len(sub_task_results) == size