Optimize add_and_remove prep queries on registry push

CHANGES/488.bugfix

@@ -0,0 +1 @@
+Optimize registry push queries that prepare content for ``add_and_remove``: fix tag replacement filtering to use the tag name, avoid N+1 blob lookups for manifest lists, and use pulpcore's native async ``aadd_and_remove``.

pulp_container/app/registry_api.py

@@ -11,7 +11,6 @@
 import json
 import logging
 import re
-from itertools import chain
 from tempfile import NamedTemporaryFile
 from urllib.parse import parse_qs, urlencode, urljoin, urlparse, urlunparse
 
@@ -52,7 +51,7 @@
     Repository,
     UploadChunk,
 )
-from pulpcore.plugin.tasking import dispatch
+from pulpcore.plugin.tasking import aadd_and_remove, dispatch
 from pulpcore.plugin.util import get_domain, get_objects_for_user, get_url
 
 from pulp_container.app import models, serializers
@@ -80,7 +79,7 @@
     FileStorageRedirects,
     S3StorageRedirects,
 )
-from pulp_container.app.tasks import aadd_and_remove, download_image_data, recursive_remove_content
+from pulp_container.app.tasks import download_image_data, recursive_remove_content
 from pulp_container.app.token_verification import (
     RegistryAuthentication,
     RegistryPermission,
@@ -93,6 +92,7 @@
     extract_data_from_signature,
     filter_resource,
     get_accepted_media_types,
+    get_content_units_to_add,
     get_full_path,
     has_task_completed,
     validate_manifest,
@@ -1401,26 +1401,8 @@ def destroy(self, request, path, pk=None):
         return Response(status=202)
 
     def get_content_units_to_add(self, manifest, tag=None):
-        add_content_units = [str(manifest.pk)]
-        if tag:
-            add_content_units.append(str(tag.pk))
-        if manifest.media_type in (
-            models.MEDIA_TYPE.MANIFEST_LIST,
-            models.MEDIA_TYPE.INDEX_OCI,
-        ):
-            for listed_manifest in manifest.listed_manifests.all():
-                add_content_units.append(listed_manifest.pk)
-                add_content_units.append(listed_manifest.config_blob_id)
-                add_content_units.extend(listed_manifest.blobs.values_list("pk", flat=True))
-        elif manifest.media_type in (
-            models.MEDIA_TYPE.MANIFEST_V2,
-            models.MEDIA_TYPE.MANIFEST_OCI,
-        ):
-            add_content_units.append(manifest.config_blob_id)
-            add_content_units.extend(manifest.blobs.values_list("pk", flat=True))
-        else:
-            add_content_units.extend(manifest.blobs.values_list("pk", flat=True))
-        return add_content_units
+        """Collect content PKs to add for a manifest (and optional tag)."""
+        return get_content_units_to_add(manifest, tag)
 
     def fake_init_manifest(self, response, pk):
         """
@@ -1630,14 +1612,18 @@ def put(self, request, path, pk=None):
                 tag.touch()
 
             add_content_units = [str(tag.pk), str(manifest.pk)] + [
-                str(content.pk)
-                for content in chain(found_blobs, found_config_blobs, found_manifests)
+                str(pk)
+                for pk in found_blobs.values_list("pk", flat=True)
+                .union(found_config_blobs.values_list("pk", flat=True))
+                .union(found_manifests.values_list("pk", flat=True))
             ]
 
+            # Filter by tag name (string), not the Tag instance — Django would
+            # coerce the instance via __str__ and never match existing tags.
             tags_to_remove = models.Tag.objects.filter(
-                pk__in=repository.latest_version().content.all(), name=tag
+                pk__in=repository.latest_version().content.all(), name=tag.name
             ).exclude(tagged_manifest=manifest)
-            remove_content_units = [str(pk) for pk in tags_to_remove.values_list("pk")]
+            remove_content_units = [str(pk) for pk in tags_to_remove.values_list("pk", flat=True)]
 
             immediate_task = dispatch(
                 aadd_and_remove,

pulp_container/app/tasks/__init__.py

@@ -1,4 +1,4 @@
-from .download_image_data import aadd_and_remove, download_image_data  # noqa
+from .download_image_data import download_image_data  # noqa
 from .builder import build_image_from_containerfile, build_image  # noqa
 from .recursive_add import recursive_add_content  # noqa
 from .recursive_remove import recursive_remove_content  # noqa

pulp_container/app/tasks/download_image_data.py

@@ -1,10 +1,7 @@
 import json
 import logging
 
-from asgiref.sync import sync_to_async
-
 from pulpcore.plugin.stages import DeclarativeContent
-from pulpcore.plugin.tasking import add_and_remove
 
 from pulp_container.app.models import ContainerRemote, ContainerRepository, Tag
 from pulp_container.app.utils import determine_media_type_from_json
@@ -16,10 +13,6 @@
 log = logging.getLogger(__name__)
 
 
-async def aadd_and_remove(*args, **kwargs):
-    return await sync_to_async(add_and_remove)(*args, **kwargs)
-
-
 def download_image_data(repository_pk, remote_pk, raw_text_manifest_data, tag_name=None):
     repository = ContainerRepository.objects.get(pk=repository_pk)
     remote = ContainerRemote.objects.get(pk=remote_pk)

pulp_container/app/utils.py

@@ -383,3 +383,38 @@ def filter_resources(element_list, include_patterns, exclude_patterns):
     if exclude_patterns:
         element_list = filter(partial(exclude, patterns=exclude_patterns), element_list)
     return list(element_list)
+
+
+def get_content_units_to_add(manifest, tag=None):
+    """Collect content PKs to add for a manifest (and optional tag).
+
+    Uses bulk queries for listed manifests and related blobs to avoid N+1
+    lookups when preparing units for ``add_and_remove`` on the push/pull-through
+    paths.
+    """
+    # Local import avoids circular imports with models.
+    from pulp_container.app.models import BlobManifest
+
+    add_content_units = [str(manifest.pk)]
+    if tag:
+        add_content_units.append(str(tag.pk))
+    if manifest.media_type in (MEDIA_TYPE.MANIFEST_LIST, MEDIA_TYPE.INDEX_OCI):
+        # values_list avoids deferred-field loads that .only() can trigger on MasterModel.
+        listed = list(manifest.listed_manifests.values_list("pk", "config_blob_id"))
+        add_content_units.extend(str(pk) for pk, _ in listed)
+        add_content_units.extend(str(config_id) for _, config_id in listed if config_id)
+        listed_pks = [pk for pk, _ in listed]
+        if listed_pks:
+            add_content_units.extend(
+                str(pk)
+                for pk in BlobManifest.objects.filter(manifest__in=listed_pks)
+                .values_list("manifest_blob_id", flat=True)
+                .distinct()
+            )
+    elif manifest.media_type in (MEDIA_TYPE.MANIFEST_V2, MEDIA_TYPE.MANIFEST_OCI):
+        if manifest.config_blob_id:
+            add_content_units.append(str(manifest.config_blob_id))
+        add_content_units.extend(str(pk) for pk in manifest.blobs.values_list("pk", flat=True))
+    else:
+        add_content_units.extend(str(pk) for pk in manifest.blobs.values_list("pk", flat=True))
+    return add_content_units

pulp_container/tests/unit/test_add_and_remove_queries.py

@@ -0,0 +1,125 @@
+"""Unit tests for push-path content unit collection used with add_and_remove."""
+
+from django.db import connection
+from django.test import TestCase
+from django.test.utils import CaptureQueriesContext
+
+from pulp_container.app.models import Blob, BlobManifest, Manifest, ManifestListManifest, Tag
+from pulp_container.app.utils import get_content_units_to_add
+from pulp_container.constants import MEDIA_TYPE
+
+
+def _digest(n: int) -> str:
+    return f"sha256:{n:064x}"
+
+
+def _create_manifest(*, digest, media_type, config_blob=None, schema_version=2):
+    return Manifest.objects.create(
+        digest=digest,
+        media_type=media_type,
+        schema_version=schema_version,
+        config_blob=config_blob,
+        data="{}",
+    )
+
+
+class TestGetContentUnitsToAdd(TestCase):
+    """Exercise get_content_units_to_add query behavior."""
+
+    def test_simple_manifest_includes_config_and_blobs(self):
+        config = Blob.objects.create(digest=_digest(1))
+        layer = Blob.objects.create(digest=_digest(2))
+        manifest = _create_manifest(
+            digest=_digest(3),
+            media_type=MEDIA_TYPE.MANIFEST_V2,
+            config_blob=config,
+        )
+        BlobManifest.objects.create(manifest=manifest, manifest_blob=layer)
+        tag = Tag.objects.create(name="latest", tagged_manifest=manifest)
+
+        units = get_content_units_to_add(manifest, tag)
+
+        self.assertEqual(
+            set(units),
+            {str(manifest.pk), str(tag.pk), str(config.pk), str(layer.pk)},
+        )
+
+    def test_manifest_list_uses_bulk_blob_query(self):
+        config_a = Blob.objects.create(digest=_digest(10))
+        config_b = Blob.objects.create(digest=_digest(11))
+        layer_a = Blob.objects.create(digest=_digest(12))
+        layer_b = Blob.objects.create(digest=_digest(13))
+
+        child_a = _create_manifest(
+            digest=_digest(20),
+            media_type=MEDIA_TYPE.MANIFEST_V2,
+            config_blob=config_a,
+        )
+        child_b = _create_manifest(
+            digest=_digest(21),
+            media_type=MEDIA_TYPE.MANIFEST_V2,
+            config_blob=config_b,
+        )
+        BlobManifest.objects.create(manifest=child_a, manifest_blob=layer_a)
+        BlobManifest.objects.create(manifest=child_b, manifest_blob=layer_b)
+
+        index = _create_manifest(
+            digest=_digest(30),
+            media_type=MEDIA_TYPE.MANIFEST_LIST,
+        )
+        # image_manifest is the list; manifest_list is the listed child (through_fields order).
+        ManifestListManifest.objects.create(image_manifest=index, manifest_list=child_a)
+        ManifestListManifest.objects.create(image_manifest=index, manifest_list=child_b)
+
+        with CaptureQueriesContext(connection) as ctx:
+            units = get_content_units_to_add(index)
+
+        # One query for listed manifests, one bulk query for related blobs — not N+1.
+        self.assertEqual(len(ctx.captured_queries), 2, [q["sql"] for q in ctx.captured_queries])
+
+        self.assertEqual(
+            set(units),
+            {
+                str(index.pk),
+                str(child_a.pk),
+                str(child_b.pk),
+                str(config_a.pk),
+                str(config_b.pk),
+                str(layer_a.pk),
+                str(layer_b.pk),
+            },
+        )
+
+    def test_all_returned_pks_are_strings(self):
+        config = Blob.objects.create(digest=_digest(40))
+        manifest = _create_manifest(
+            digest=_digest(41),
+            media_type=MEDIA_TYPE.MANIFEST_OCI,
+            config_blob=config,
+        )
+        units = get_content_units_to_add(manifest)
+        self.assertTrue(all(isinstance(pk, str) for pk in units))
+
+
+class TestTagReplacementFilter(TestCase):
+    """Ensure tag replacement filters by Tag.name, not Tag.__str__."""
+
+    def test_filter_by_tag_name_matches_existing_tag(self):
+        manifest_old = _create_manifest(
+            digest=_digest(50),
+            media_type=MEDIA_TYPE.MANIFEST_V2,
+        )
+        manifest_new = _create_manifest(
+            digest=_digest(51),
+            media_type=MEDIA_TYPE.MANIFEST_V2,
+        )
+        old_tag = Tag.objects.create(name="latest", tagged_manifest=manifest_old)
+        new_tag = Tag.objects.create(name="latest", tagged_manifest=manifest_new)
+
+        # Correct filter used by the push path after optimization.
+        matched = Tag.objects.filter(name=new_tag.name).exclude(tagged_manifest=manifest_new)
+        self.assertQuerySetEqual(matched, [old_tag], ordered=False)
+
+        # Passing the Tag instance would coerce via MasterModel.__str__ and miss.
+        broken = Tag.objects.filter(name=new_tag).exclude(tagged_manifest=manifest_new)
+        self.assertFalse(broken.exists())