We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
bulk_update() in content-stages can cause (very rare) deadlock
**Version**
3.14

**Describe the bug**
In high-concurrency environments, with overlapping content, calling bulk_update() can cause a deadlock. Specifically, this call:

https://github.com/pulp/pulpcore/blob/main/pulpcore/plugin/stages/content_stages.py#L158-L164

Ordering the list-to-be-updated does not, alas, protect us - because Postgres doesn't guarantee order when doing an update like this.

**To Reproduce**
We have only seen this "in the wild" once, syncing 8-10 repos with similar content at the same time with 10 workers available.

**Expected behavior**
Don't deadlock.

**Additional context**
This is the traceback from the initial description for

https://bugzilla.redhat.com/show_bug.cgi?id=2062526

We fixed the deadlock noted in https://bugzilla.redhat.com/show_bug.cgi?id=2062526#c2 under #2420 



--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `pulpcore/plugin/stages/content_stages.py`
Content:
```
1 from collections import defaultdict
2 
3 from asgiref.sync import sync_to_async
4 from django.core.exceptions import ObjectDoesNotExist
5 from django.db import IntegrityError, transaction
6 from django.db.models import Q
7 
8 from pulpcore.plugin.sync import sync_to_async_iterable
9 
10 from pulpcore.plugin.models import Content, ContentArtifact, ProgressReport
11 
12 from .api import Stage
13 
14 
15 class QueryExistingContents(Stage):
16     """
17     A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related
18     :class:`~pulpcore.plugin.models.ContentArtifact` objects too.
19 
20     This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`
21     and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each
22     :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one
23     :class:`~pulpcore.plugin.models.Artifact`.
24 
25     This stage inspects any "unsaved" Content unit objects and searches for existing saved Content
26     units inside Pulp with the same unit key. Any existing Content objects found, replace their
27     "unsaved" counterpart in the :class:`~pulpcore.plugin.stages.DeclarativeContent` object.
28 
29     Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to `self._out_q` after it has
30     been handled.
31 
32     This stage drains all available items from `self._in_q` and batches everything into one large
33     call to the db for efficiency.
34     """
35 
36     async def run(self):
37         """
38         The coroutine for this stage.
39 
40         Returns:
41             The coroutine for this stage.
42         """
43         async for batch in self.batches():
44             content_q_by_type = defaultdict(lambda: Q(pk__in=[]))
45             d_content_by_nat_key = defaultdict(list)
46             for d_content in batch:
47                 if d_content.content._state.adding:
48                     model_type = type(d_content.content)
49                     unit_q = d_content.content.q()
50                     content_q_by_type[model_type] = content_q_by_type[model_type] | unit_q
51                     d_content_by_nat_key[d_content.content.natural_key()].append(d_content)
52 
53             for model_type, content_q in content_q_by_type.items():
54                 try:
55                     await sync_to_async(model_type.objects.filter(content_q).touch)()
56                 except AttributeError:
57                     raise TypeError(
58                         "Plugins which declare custom ORM managers on their content classes "
59                         "should have those managers inherit from "
60                         "pulpcore.plugin.models.ContentManager."
61                     )
62                 async for result in sync_to_async_iterable(
63                     model_type.objects.filter(content_q).iterator()
64                 ):
65                     for d_content in d_content_by_nat_key[result.natural_key()]:
66                         d_content.content = result
67 
68             for d_content in batch:
69                 await self.put(d_content)
70 
71 
72 class ContentSaver(Stage):
73     """
74     A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related
75     :class:`~pulpcore.plugin.models.ContentArtifact` objects too.
76 
77     This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`
78     and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each
79     :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one
80     :class:`~pulpcore.plugin.models.Artifact`.
81 
82     Each "unsaved" Content objects is saved and a :class:`~pulpcore.plugin.models.ContentArtifact`
83     objects too.
84 
85     Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to after it has been handled.
86 
87     This stage drains all available items from `self._in_q` and batches everything into one large
88     call to the db for efficiency.
89     """
90 
91     async def run(self):
92         """
93         The coroutine for this stage.
94 
95         Returns:
96             The coroutine for this stage.
97         """
98         async for batch in self.batches():
99 
100             def process_batch():
101                 content_artifact_bulk = []
102                 to_update_ca_query = ContentArtifact.objects.none()
103                 to_update_ca_bulk = []
104                 to_update_ca_artifact = {}
105                 with transaction.atomic():
106                     self._pre_save(batch)
107                     # Process the batch in dc.content.natural_keys order.
108                     # This prevents deadlocks when we're processing the same/similar content
109                     # in concurrent workers.
110                     batch.sort(key=lambda x: "".join(map(str, x.content.natural_key())))
111                     for d_content in batch:
112                         # Are we saving to the database for the first time?
113                         content_already_saved = not d_content.content._state.adding
114                         if not content_already_saved:
115                             try:
116                                 with transaction.atomic():
117                                     d_content.content.save()
118                             except IntegrityError as e:
119                                 try:
120                                     d_content.content = d_content.content.__class__.objects.get(
121                                         d_content.content.q()
122                                     )
123                                 except ObjectDoesNotExist:
124                                     raise e
125                             else:
126                                 for d_artifact in d_content.d_artifacts:
127                                     if not d_artifact.artifact._state.adding:
128                                         artifact = d_artifact.artifact
129                                     else:
130                                         # set to None for on-demand synced artifacts
131                                         artifact = None
132                                     content_artifact = ContentArtifact(
133                                         content=d_content.content,
134                                         artifact=artifact,
135                                         relative_path=d_artifact.relative_path,
136                                     )
137                                     content_artifact_bulk.append(content_artifact)
138                                 continue
139                         # When the Content already exists, check if ContentArtifacts need to be
140                         # updated
141                         for d_artifact in d_content.d_artifacts:
142                             if not d_artifact.artifact._state.adding:
143                                 # the artifact is already present in the database; update references
144                                 # Creating one large query and one large dictionary
145                                 to_update_ca_query |= ContentArtifact.objects.filter(
146                                     content=d_content.content,
147                                     relative_path=d_artifact.relative_path,
148                                 )
149                                 key = (d_content.content.pk, d_artifact.relative_path)
150                                 to_update_ca_artifact[key] = d_artifact.artifact
151                     # Query db once and update each object in memory for bulk_update call
152                     for content_artifact in to_update_ca_query.iterator():
153                         key = (content_artifact.content_id, content_artifact.relative_path)
154                         # Maybe remove dict elements after to reduce memory?
155                         content_artifact.artifact = to_update_ca_artifact[key]
156                         to_update_ca_bulk.append(content_artifact)
157 
158                     # Sort the lists we're about to do bulk updates/creates on.
159                     # We know to_update_ca_bulk entries already are in the DB, so we can enforce
160                     # order just using pulp_id.
161                     to_update_ca_bulk.sort(key=lambda x: x.pulp_id)
162                     content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))
163 
164                     ContentArtifact.objects.bulk_update(to_update_ca_bulk, ["artifact"])
165                     ContentArtifact.objects.bulk_get_or_create(content_artifact_bulk)
166                     self._post_save(batch)
167 
168             await sync_to_async(process_batch)()
169             for declarative_content in batch:
170                 await self.put(declarative_content)
171 
172     def _pre_save(self, batch):
173         """
174         A hook plugin-writers can override to save related objects prior to content unit saving.
175 
176         This is run within the same transaction as the content unit saving.
177 
178         Args:
179             batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of
180                 :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.
181 
182         """
183         pass
184 
185     def _post_save(self, batch):
186         """
187         A hook plugin-writers can override to save related objects after content unit saving.
188 
189         This is run within the same transaction as the content unit saving.
190 
191         Args:
192             batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of
193                 :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.
194 
195         """
196         pass
197 
198 
199 class ResolveContentFutures(Stage):
200     """
201     This stage resolves the futures in :class:`~pulpcore.plugin.stages.DeclarativeContent`.
202 
203     Futures results are set to the found/created :class:`~pulpcore.plugin.models.Content`.
204 
205     This is useful when data downloaded from the plugin API needs to be parsed by FirstStage to
206     create additional :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be send down
207     the pipeline. Consider an example where content type `Foo` references additional instances of a
208     different content type `Bar`. Consider this code in FirstStage::
209 
210         # Create d_content and d_artifact for a `foo_a`
211         foo_a = DeclarativeContent(...)
212         # Send it in the pipeline
213         await self.put(foo_a)
214 
215         ...
216 
217         foo_a_content = await foo_a.resolution()  # awaits until the foo_a reaches this stage
218 
219     This creates a "looping" pattern, of sorts, where downloaded content at the end of the pipeline
220     can introduce new additional to-be-downloaded content at the beginning of the pipeline.
221     On the other hand, it can impose a substantial performance decrement of batching content in the
222     earlier stages.
223     If you want to drop a declarative content prematurely from the pipeline, use the function
224     `resolve()` to unblock the coroutines awaiting the attached future and do not hand the content
225     to the next stage.
226     As a rule of thumb, sending more items into the pipeline first and awaiting their resolution
227     later is better.
228     """
229 
230     async def run(self):
231         """
232         The coroutine for this stage.
233 
234         Returns:
235             The coroutine for this stage.
236         """
237         async for d_content in self.items():
238             d_content.resolve()
239             await self.put(d_content)
240 
241 
242 class ContentAssociation(Stage):
243     """
244     A Stages API stage that associates content units with `new_version`.
245 
246     This stage stores all content unit primary keys in memory before running. This is done to
247     compute the units already associated but not received from `self._in_q`. These units are passed
248     via `self._out_q` to the next stage as a :class:`django.db.models.query.QuerySet`.
249 
250     This stage creates a ProgressReport named 'Associating Content' that counts the number of units
251     associated. Since it's a stream the total count isn't known until it's finished.
252 
253     If `mirror` was enabled, then content units may also be un-assocated (removed) from
254     `new_version`. A ProgressReport named 'Un-Associating Content' is created that counts the number
255     of units un-associated.
256 
257     Args:
258         new_version (:class:`~pulpcore.plugin.models.RepositoryVersion`): The repo version this
259             stage associates content with.
260         mirror (bool): Whether or not to "mirror" the stream of DeclarativeContent - whether content
261             not in the stream should be removed from the repository.
262         args: unused positional arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.
263         kwargs: unused keyword arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.
264     """
265 
266     def __init__(self, new_version, mirror, *args, **kwargs):
267         super().__init__(*args, **kwargs)
268         self.new_version = new_version
269         self.allow_delete = mirror
270 
271     async def run(self):
272         """
273         The coroutine for this stage.
274 
275         Returns:
276             The coroutine for this stage.
277         """
278         async with ProgressReport(message="Associating Content", code="associating.content") as pb:
279             to_delete = {
280                 i
281                 async for i in sync_to_async_iterable(
282                     self.new_version.content.values_list("pk", flat=True)
283                 )
284             }
285 
286             async for batch in self.batches():
287                 to_add = set()
288                 for d_content in batch:
289                     try:
290                         to_delete.remove(d_content.content.pk)
291                     except KeyError:
292                         to_add.add(d_content.content.pk)
293                         await self.put(d_content)
294 
295                 if to_add:
296                     await sync_to_async(self.new_version.add_content)(
297                         Content.objects.filter(pk__in=to_add)
298                     )
299                     await pb.aincrease_by(len(to_add))
300 
301             if self.allow_delete:
302                 async with ProgressReport(
303                     message="Un-Associating Content", code="unassociating.content"
304                 ) as pb:
305                     if to_delete:
306                         await sync_to_async(self.new_version.remove_content)(
307                             Content.objects.filter(pk__in=to_delete)
308                         )
309                         await pb.aincrease_by(len(to_delete))
310
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/pulpcore/plugin/stages/content_stages.py b/pulpcore/plugin/stages/content_stages.py
--- a/pulpcore/plugin/stages/content_stages.py
+++ b/pulpcore/plugin/stages/content_stages.py
@@ -155,13 +155,39 @@
                         content_artifact.artifact = to_update_ca_artifact[key]
                         to_update_ca_bulk.append(content_artifact)
 
-                    # Sort the lists we're about to do bulk updates/creates on.
-                    # We know to_update_ca_bulk entries already are in the DB, so we can enforce
-                    # order just using pulp_id.
-                    to_update_ca_bulk.sort(key=lambda x: x.pulp_id)
-                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))
+                    # to_update_ca_bulk are the CAs that we know are already persisted.
+                    # We need to update their artifact_ids, and wish to do it in bulk to
+                    # avoid hundreds of round-trips to the database.
+                    #
+                    # To avoid deadlocks in high-concurrency environments with overlapping
+                    # content, we need to update the rows in some defined order. Unfortunately,
+                    # postgres doesn't support order-on-update - but it *does* support ordering
+                    # on select-for-update. So, we select-for-update, in pulp_id order, the
+                    # rows we're about to update as one db-call, and then do the update in a
+                    # second.
+                    ids = [k.pulp_id for k in to_update_ca_bulk]
+                    with transaction.atomic():
+                        # "len()" forces the QA to be evaluated. Using exist() or count() won't
+                        # work for us - Django is smart enough to either not-order, or even
+                        # not-emit, a select-for-update in these cases.
+                        #
+                        # To maximize performance, we make sure to only ask for pulp_ids, and
+                        # avoid instantiating a python-object for the affected CAs by using
+                        # values_list()
+                        len(
+                            ContentArtifact.objects.filter(pulp_id__in=ids)
+                            .only("pulp_id")
+                            .order_by("pulp_id")
+                            .select_for_update()
+                            .values_list()
+                        )
+                        ContentArtifact.objects.bulk_update(to_update_ca_bulk, ["artifact"])
 
-                    ContentArtifact.objects.bulk_update(to_update_ca_bulk, ["artifact"])
+                    # To avoid a similar deadlock issue when calling get_or_create, we sort the
+                    # "new" CAs to make sure inserts happen in a defined order. Since we can't
+                    # trust the pulp_id (by the time we go to create a CA, it may already exist,
+                    # and be replaced by the 'real' one), we sort by their "natural key".
+                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))
                     ContentArtifact.objects.bulk_get_or_create(content_artifact_bulk)
                     self._post_save(batch)

{"golden_diff": "diff --git a/pulpcore/plugin/stages/content_stages.py b/pulpcore/plugin/stages/content_stages.py\n--- a/pulpcore/plugin/stages/content_stages.py\n+++ b/pulpcore/plugin/stages/content_stages.py\n@@ -155,13 +155,39 @@\n                         content_artifact.artifact = to_update_ca_artifact[key]\n                         to_update_ca_bulk.append(content_artifact)\n \n-                    # Sort the lists we're about to do bulk updates/creates on.\n-                    # We know to_update_ca_bulk entries already are in the DB, so we can enforce\n-                    # order just using pulp_id.\n-                    to_update_ca_bulk.sort(key=lambda x: x.pulp_id)\n-                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))\n+                    # to_update_ca_bulk are the CAs that we know are already persisted.\n+                    # We need to update their artifact_ids, and wish to do it in bulk to\n+                    # avoid hundreds of round-trips to the database.\n+                    #\n+                    # To avoid deadlocks in high-concurrency environments with overlapping\n+                    # content, we need to update the rows in some defined order. Unfortunately,\n+                    # postgres doesn't support order-on-update - but it *does* support ordering\n+                    # on select-for-update. So, we select-for-update, in pulp_id order, the\n+                    # rows we're about to update as one db-call, and then do the update in a\n+                    # second.\n+                    ids = [k.pulp_id for k in to_update_ca_bulk]\n+                    with transaction.atomic():\n+                        # \"len()\" forces the QA to be evaluated. Using exist() or count() won't\n+                        # work for us - Django is smart enough to either not-order, or even\n+                        # not-emit, a select-for-update in these cases.\n+                        #\n+                        # To maximize performance, we make sure to only ask for pulp_ids, and\n+                        # avoid instantiating a python-object for the affected CAs by using\n+                        # values_list()\n+                        len(\n+                            ContentArtifact.objects.filter(pulp_id__in=ids)\n+                            .only(\"pulp_id\")\n+                            .order_by(\"pulp_id\")\n+                            .select_for_update()\n+                            .values_list()\n+                        )\n+                        ContentArtifact.objects.bulk_update(to_update_ca_bulk, [\"artifact\"])\n \n-                    ContentArtifact.objects.bulk_update(to_update_ca_bulk, [\"artifact\"])\n+                    # To avoid a similar deadlock issue when calling get_or_create, we sort the\n+                    # \"new\" CAs to make sure inserts happen in a defined order. Since we can't\n+                    # trust the pulp_id (by the time we go to create a CA, it may already exist,\n+                    # and be replaced by the 'real' one), we sort by their \"natural key\".\n+                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))\n                     ContentArtifact.objects.bulk_get_or_create(content_artifact_bulk)\n                     self._post_save(batch)\n", "issue": "bulk_update() in content-stages can cause (very rare) deadlock\n**Version**\r\n3.14\r\n\r\n**Describe the bug**\r\nIn high-concurrency environments, with overlapping content, calling bulk_update() can cause a deadlock. Specifically, this call:\r\n\r\nhttps://github.com/pulp/pulpcore/blob/main/pulpcore/plugin/stages/content_stages.py#L158-L164\r\n\r\nOrdering the list-to-be-updated does not, alas, protect us - because Postgres doesn't guarantee order when doing an update like this.\r\n\r\n**To Reproduce**\r\nWe have only seen this \"in the wild\" once, syncing 8-10 repos with similar content at the same time with 10 workers available.\r\n\r\n**Expected behavior**\r\nDon't deadlock.\r\n\r\n**Additional context**\r\nThis is the traceback from the initial description for\r\n\r\nhttps://bugzilla.redhat.com/show_bug.cgi?id=2062526\r\n\r\nWe fixed the deadlock noted in https://bugzilla.redhat.com/show_bug.cgi?id=2062526#c2 under #2420 \r\n\r\n\n", "before_files": [{"content": "from collections import defaultdict\n\nfrom asgiref.sync import sync_to_async\nfrom django.core.exceptions import ObjectDoesNotExist\nfrom django.db import IntegrityError, transaction\nfrom django.db.models import Q\n\nfrom pulpcore.plugin.sync import sync_to_async_iterable\n\nfrom pulpcore.plugin.models import Content, ContentArtifact, ProgressReport\n\nfrom .api import Stage\n\n\nclass QueryExistingContents(Stage):\n    \"\"\"\n    A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related\n    :class:`~pulpcore.plugin.models.ContentArtifact` objects too.\n\n    This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`\n    and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each\n    :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one\n    :class:`~pulpcore.plugin.models.Artifact`.\n\n    This stage inspects any \"unsaved\" Content unit objects and searches for existing saved Content\n    units inside Pulp with the same unit key. Any existing Content objects found, replace their\n    \"unsaved\" counterpart in the :class:`~pulpcore.plugin.stages.DeclarativeContent` object.\n\n    Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to `self._out_q` after it has\n    been handled.\n\n    This stage drains all available items from `self._in_q` and batches everything into one large\n    call to the db for efficiency.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for batch in self.batches():\n            content_q_by_type = defaultdict(lambda: Q(pk__in=[]))\n            d_content_by_nat_key = defaultdict(list)\n            for d_content in batch:\n                if d_content.content._state.adding:\n                    model_type = type(d_content.content)\n                    unit_q = d_content.content.q()\n                    content_q_by_type[model_type] = content_q_by_type[model_type] | unit_q\n                    d_content_by_nat_key[d_content.content.natural_key()].append(d_content)\n\n            for model_type, content_q in content_q_by_type.items():\n                try:\n                    await sync_to_async(model_type.objects.filter(content_q).touch)()\n                except AttributeError:\n                    raise TypeError(\n                        \"Plugins which declare custom ORM managers on their content classes \"\n                        \"should have those managers inherit from \"\n                        \"pulpcore.plugin.models.ContentManager.\"\n                    )\n                async for result in sync_to_async_iterable(\n                    model_type.objects.filter(content_q).iterator()\n                ):\n                    for d_content in d_content_by_nat_key[result.natural_key()]:\n                        d_content.content = result\n\n            for d_content in batch:\n                await self.put(d_content)\n\n\nclass ContentSaver(Stage):\n    \"\"\"\n    A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related\n    :class:`~pulpcore.plugin.models.ContentArtifact` objects too.\n\n    This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`\n    and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each\n    :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one\n    :class:`~pulpcore.plugin.models.Artifact`.\n\n    Each \"unsaved\" Content objects is saved and a :class:`~pulpcore.plugin.models.ContentArtifact`\n    objects too.\n\n    Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to after it has been handled.\n\n    This stage drains all available items from `self._in_q` and batches everything into one large\n    call to the db for efficiency.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for batch in self.batches():\n\n            def process_batch():\n                content_artifact_bulk = []\n                to_update_ca_query = ContentArtifact.objects.none()\n                to_update_ca_bulk = []\n                to_update_ca_artifact = {}\n                with transaction.atomic():\n                    self._pre_save(batch)\n                    # Process the batch in dc.content.natural_keys order.\n                    # This prevents deadlocks when we're processing the same/similar content\n                    # in concurrent workers.\n                    batch.sort(key=lambda x: \"\".join(map(str, x.content.natural_key())))\n                    for d_content in batch:\n                        # Are we saving to the database for the first time?\n                        content_already_saved = not d_content.content._state.adding\n                        if not content_already_saved:\n                            try:\n                                with transaction.atomic():\n                                    d_content.content.save()\n                            except IntegrityError as e:\n                                try:\n                                    d_content.content = d_content.content.__class__.objects.get(\n                                        d_content.content.q()\n                                    )\n                                except ObjectDoesNotExist:\n                                    raise e\n                            else:\n                                for d_artifact in d_content.d_artifacts:\n                                    if not d_artifact.artifact._state.adding:\n                                        artifact = d_artifact.artifact\n                                    else:\n                                        # set to None for on-demand synced artifacts\n                                        artifact = None\n                                    content_artifact = ContentArtifact(\n                                        content=d_content.content,\n                                        artifact=artifact,\n                                        relative_path=d_artifact.relative_path,\n                                    )\n                                    content_artifact_bulk.append(content_artifact)\n                                continue\n                        # When the Content already exists, check if ContentArtifacts need to be\n                        # updated\n                        for d_artifact in d_content.d_artifacts:\n                            if not d_artifact.artifact._state.adding:\n                                # the artifact is already present in the database; update references\n                                # Creating one large query and one large dictionary\n                                to_update_ca_query |= ContentArtifact.objects.filter(\n                                    content=d_content.content,\n                                    relative_path=d_artifact.relative_path,\n                                )\n                                key = (d_content.content.pk, d_artifact.relative_path)\n                                to_update_ca_artifact[key] = d_artifact.artifact\n                    # Query db once and update each object in memory for bulk_update call\n                    for content_artifact in to_update_ca_query.iterator():\n                        key = (content_artifact.content_id, content_artifact.relative_path)\n                        # Maybe remove dict elements after to reduce memory?\n                        content_artifact.artifact = to_update_ca_artifact[key]\n                        to_update_ca_bulk.append(content_artifact)\n\n                    # Sort the lists we're about to do bulk updates/creates on.\n                    # We know to_update_ca_bulk entries already are in the DB, so we can enforce\n                    # order just using pulp_id.\n                    to_update_ca_bulk.sort(key=lambda x: x.pulp_id)\n                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))\n\n                    ContentArtifact.objects.bulk_update(to_update_ca_bulk, [\"artifact\"])\n                    ContentArtifact.objects.bulk_get_or_create(content_artifact_bulk)\n                    self._post_save(batch)\n\n            await sync_to_async(process_batch)()\n            for declarative_content in batch:\n                await self.put(declarative_content)\n\n    def _pre_save(self, batch):\n        \"\"\"\n        A hook plugin-writers can override to save related objects prior to content unit saving.\n\n        This is run within the same transaction as the content unit saving.\n\n        Args:\n            batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of\n                :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.\n\n        \"\"\"\n        pass\n\n    def _post_save(self, batch):\n        \"\"\"\n        A hook plugin-writers can override to save related objects after content unit saving.\n\n        This is run within the same transaction as the content unit saving.\n\n        Args:\n            batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of\n                :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.\n\n        \"\"\"\n        pass\n\n\nclass ResolveContentFutures(Stage):\n    \"\"\"\n    This stage resolves the futures in :class:`~pulpcore.plugin.stages.DeclarativeContent`.\n\n    Futures results are set to the found/created :class:`~pulpcore.plugin.models.Content`.\n\n    This is useful when data downloaded from the plugin API needs to be parsed by FirstStage to\n    create additional :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be send down\n    the pipeline. Consider an example where content type `Foo` references additional instances of a\n    different content type `Bar`. Consider this code in FirstStage::\n\n        # Create d_content and d_artifact for a `foo_a`\n        foo_a = DeclarativeContent(...)\n        # Send it in the pipeline\n        await self.put(foo_a)\n\n        ...\n\n        foo_a_content = await foo_a.resolution()  # awaits until the foo_a reaches this stage\n\n    This creates a \"looping\" pattern, of sorts, where downloaded content at the end of the pipeline\n    can introduce new additional to-be-downloaded content at the beginning of the pipeline.\n    On the other hand, it can impose a substantial performance decrement of batching content in the\n    earlier stages.\n    If you want to drop a declarative content prematurely from the pipeline, use the function\n    `resolve()` to unblock the coroutines awaiting the attached future and do not hand the content\n    to the next stage.\n    As a rule of thumb, sending more items into the pipeline first and awaiting their resolution\n    later is better.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for d_content in self.items():\n            d_content.resolve()\n            await self.put(d_content)\n\n\nclass ContentAssociation(Stage):\n    \"\"\"\n    A Stages API stage that associates content units with `new_version`.\n\n    This stage stores all content unit primary keys in memory before running. This is done to\n    compute the units already associated but not received from `self._in_q`. These units are passed\n    via `self._out_q` to the next stage as a :class:`django.db.models.query.QuerySet`.\n\n    This stage creates a ProgressReport named 'Associating Content' that counts the number of units\n    associated. Since it's a stream the total count isn't known until it's finished.\n\n    If `mirror` was enabled, then content units may also be un-assocated (removed) from\n    `new_version`. A ProgressReport named 'Un-Associating Content' is created that counts the number\n    of units un-associated.\n\n    Args:\n        new_version (:class:`~pulpcore.plugin.models.RepositoryVersion`): The repo version this\n            stage associates content with.\n        mirror (bool): Whether or not to \"mirror\" the stream of DeclarativeContent - whether content\n            not in the stream should be removed from the repository.\n        args: unused positional arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.\n        kwargs: unused keyword arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.\n    \"\"\"\n\n    def __init__(self, new_version, mirror, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.new_version = new_version\n        self.allow_delete = mirror\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async with ProgressReport(message=\"Associating Content\", code=\"associating.content\") as pb:\n            to_delete = {\n                i\n                async for i in sync_to_async_iterable(\n                    self.new_version.content.values_list(\"pk\", flat=True)\n                )\n            }\n\n            async for batch in self.batches():\n                to_add = set()\n                for d_content in batch:\n                    try:\n                        to_delete.remove(d_content.content.pk)\n                    except KeyError:\n                        to_add.add(d_content.content.pk)\n                        await self.put(d_content)\n\n                if to_add:\n                    await sync_to_async(self.new_version.add_content)(\n                        Content.objects.filter(pk__in=to_add)\n                    )\n                    await pb.aincrease_by(len(to_add))\n\n            if self.allow_delete:\n                async with ProgressReport(\n                    message=\"Un-Associating Content\", code=\"unassociating.content\"\n                ) as pb:\n                    if to_delete:\n                        await sync_to_async(self.new_version.remove_content)(\n                            Content.objects.filter(pk__in=to_delete)\n                        )\n                        await pb.aincrease_by(len(to_delete))\n", "path": "pulpcore/plugin/stages/content_stages.py"}], "after_files": [{"content": "from collections import defaultdict\n\nfrom asgiref.sync import sync_to_async\nfrom django.core.exceptions import ObjectDoesNotExist\nfrom django.db import IntegrityError, transaction\nfrom django.db.models import Q\n\nfrom pulpcore.plugin.sync import sync_to_async_iterable\n\nfrom pulpcore.plugin.models import Content, ContentArtifact, ProgressReport\n\nfrom .api import Stage\n\n\nclass QueryExistingContents(Stage):\n    \"\"\"\n    A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related\n    :class:`~pulpcore.plugin.models.ContentArtifact` objects too.\n\n    This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`\n    and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each\n    :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one\n    :class:`~pulpcore.plugin.models.Artifact`.\n\n    This stage inspects any \"unsaved\" Content unit objects and searches for existing saved Content\n    units inside Pulp with the same unit key. Any existing Content objects found, replace their\n    \"unsaved\" counterpart in the :class:`~pulpcore.plugin.stages.DeclarativeContent` object.\n\n    Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to `self._out_q` after it has\n    been handled.\n\n    This stage drains all available items from `self._in_q` and batches everything into one large\n    call to the db for efficiency.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for batch in self.batches():\n            content_q_by_type = defaultdict(lambda: Q(pk__in=[]))\n            d_content_by_nat_key = defaultdict(list)\n            for d_content in batch:\n                if d_content.content._state.adding:\n                    model_type = type(d_content.content)\n                    unit_q = d_content.content.q()\n                    content_q_by_type[model_type] = content_q_by_type[model_type] | unit_q\n                    d_content_by_nat_key[d_content.content.natural_key()].append(d_content)\n\n            for model_type, content_q in content_q_by_type.items():\n                try:\n                    await sync_to_async(model_type.objects.filter(content_q).touch)()\n                except AttributeError:\n                    raise TypeError(\n                        \"Plugins which declare custom ORM managers on their content classes \"\n                        \"should have those managers inherit from \"\n                        \"pulpcore.plugin.models.ContentManager.\"\n                    )\n                async for result in sync_to_async_iterable(\n                    model_type.objects.filter(content_q).iterator()\n                ):\n                    for d_content in d_content_by_nat_key[result.natural_key()]:\n                        d_content.content = result\n\n            for d_content in batch:\n                await self.put(d_content)\n\n\nclass ContentSaver(Stage):\n    \"\"\"\n    A Stages API stage that saves :attr:`DeclarativeContent.content` objects and saves its related\n    :class:`~pulpcore.plugin.models.ContentArtifact` objects too.\n\n    This stage expects :class:`~pulpcore.plugin.stages.DeclarativeContent` units from `self._in_q`\n    and inspects their associated :class:`~pulpcore.plugin.stages.DeclarativeArtifact` objects. Each\n    :class:`~pulpcore.plugin.stages.DeclarativeArtifact` object stores one\n    :class:`~pulpcore.plugin.models.Artifact`.\n\n    Each \"unsaved\" Content objects is saved and a :class:`~pulpcore.plugin.models.ContentArtifact`\n    objects too.\n\n    Each :class:`~pulpcore.plugin.stages.DeclarativeContent` is sent to after it has been handled.\n\n    This stage drains all available items from `self._in_q` and batches everything into one large\n    call to the db for efficiency.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for batch in self.batches():\n\n            def process_batch():\n                content_artifact_bulk = []\n                to_update_ca_query = ContentArtifact.objects.none()\n                to_update_ca_bulk = []\n                to_update_ca_artifact = {}\n                with transaction.atomic():\n                    self._pre_save(batch)\n                    # Process the batch in dc.content.natural_keys order.\n                    # This prevents deadlocks when we're processing the same/similar content\n                    # in concurrent workers.\n                    batch.sort(key=lambda x: \"\".join(map(str, x.content.natural_key())))\n                    for d_content in batch:\n                        # Are we saving to the database for the first time?\n                        content_already_saved = not d_content.content._state.adding\n                        if not content_already_saved:\n                            try:\n                                with transaction.atomic():\n                                    d_content.content.save()\n                            except IntegrityError as e:\n                                try:\n                                    d_content.content = d_content.content.__class__.objects.get(\n                                        d_content.content.q()\n                                    )\n                                except ObjectDoesNotExist:\n                                    raise e\n                            else:\n                                for d_artifact in d_content.d_artifacts:\n                                    if not d_artifact.artifact._state.adding:\n                                        artifact = d_artifact.artifact\n                                    else:\n                                        # set to None for on-demand synced artifacts\n                                        artifact = None\n                                    content_artifact = ContentArtifact(\n                                        content=d_content.content,\n                                        artifact=artifact,\n                                        relative_path=d_artifact.relative_path,\n                                    )\n                                    content_artifact_bulk.append(content_artifact)\n                                continue\n                        # When the Content already exists, check if ContentArtifacts need to be\n                        # updated\n                        for d_artifact in d_content.d_artifacts:\n                            if not d_artifact.artifact._state.adding:\n                                # the artifact is already present in the database; update references\n                                # Creating one large query and one large dictionary\n                                to_update_ca_query |= ContentArtifact.objects.filter(\n                                    content=d_content.content,\n                                    relative_path=d_artifact.relative_path,\n                                )\n                                key = (d_content.content.pk, d_artifact.relative_path)\n                                to_update_ca_artifact[key] = d_artifact.artifact\n                    # Query db once and update each object in memory for bulk_update call\n                    for content_artifact in to_update_ca_query.iterator():\n                        key = (content_artifact.content_id, content_artifact.relative_path)\n                        # Maybe remove dict elements after to reduce memory?\n                        content_artifact.artifact = to_update_ca_artifact[key]\n                        to_update_ca_bulk.append(content_artifact)\n\n                    # to_update_ca_bulk are the CAs that we know are already persisted.\n                    # We need to update their artifact_ids, and wish to do it in bulk to\n                    # avoid hundreds of round-trips to the database.\n                    #\n                    # To avoid deadlocks in high-concurrency environments with overlapping\n                    # content, we need to update the rows in some defined order. Unfortunately,\n                    # postgres doesn't support order-on-update - but it *does* support ordering\n                    # on select-for-update. So, we select-for-update, in pulp_id order, the\n                    # rows we're about to update as one db-call, and then do the update in a\n                    # second.\n                    ids = [k.pulp_id for k in to_update_ca_bulk]\n                    with transaction.atomic():\n                        # \"len()\" forces the QA to be evaluated. Using exist() or count() won't\n                        # work for us - Django is smart enough to either not-order, or even\n                        # not-emit, a select-for-update in these cases.\n                        #\n                        # To maximize performance, we make sure to only ask for pulp_ids, and\n                        # avoid instantiating a python-object for the affected CAs by using\n                        # values_list()\n                        len(\n                            ContentArtifact.objects.filter(pulp_id__in=ids)\n                            .only(\"pulp_id\")\n                            .order_by(\"pulp_id\")\n                            .select_for_update()\n                            .values_list()\n                        )\n                        ContentArtifact.objects.bulk_update(to_update_ca_bulk, [\"artifact\"])\n\n                    # To avoid a similar deadlock issue when calling get_or_create, we sort the\n                    # \"new\" CAs to make sure inserts happen in a defined order. Since we can't\n                    # trust the pulp_id (by the time we go to create a CA, it may already exist,\n                    # and be replaced by the 'real' one), we sort by their \"natural key\".\n                    content_artifact_bulk.sort(key=lambda x: ContentArtifact.sort_key(x))\n                    ContentArtifact.objects.bulk_get_or_create(content_artifact_bulk)\n                    self._post_save(batch)\n\n            await sync_to_async(process_batch)()\n            for declarative_content in batch:\n                await self.put(declarative_content)\n\n    def _pre_save(self, batch):\n        \"\"\"\n        A hook plugin-writers can override to save related objects prior to content unit saving.\n\n        This is run within the same transaction as the content unit saving.\n\n        Args:\n            batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of\n                :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.\n\n        \"\"\"\n        pass\n\n    def _post_save(self, batch):\n        \"\"\"\n        A hook plugin-writers can override to save related objects after content unit saving.\n\n        This is run within the same transaction as the content unit saving.\n\n        Args:\n            batch (list of :class:`~pulpcore.plugin.stages.DeclarativeContent`): The batch of\n                :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be saved.\n\n        \"\"\"\n        pass\n\n\nclass ResolveContentFutures(Stage):\n    \"\"\"\n    This stage resolves the futures in :class:`~pulpcore.plugin.stages.DeclarativeContent`.\n\n    Futures results are set to the found/created :class:`~pulpcore.plugin.models.Content`.\n\n    This is useful when data downloaded from the plugin API needs to be parsed by FirstStage to\n    create additional :class:`~pulpcore.plugin.stages.DeclarativeContent` objects to be send down\n    the pipeline. Consider an example where content type `Foo` references additional instances of a\n    different content type `Bar`. Consider this code in FirstStage::\n\n        # Create d_content and d_artifact for a `foo_a`\n        foo_a = DeclarativeContent(...)\n        # Send it in the pipeline\n        await self.put(foo_a)\n\n        ...\n\n        foo_a_content = await foo_a.resolution()  # awaits until the foo_a reaches this stage\n\n    This creates a \"looping\" pattern, of sorts, where downloaded content at the end of the pipeline\n    can introduce new additional to-be-downloaded content at the beginning of the pipeline.\n    On the other hand, it can impose a substantial performance decrement of batching content in the\n    earlier stages.\n    If you want to drop a declarative content prematurely from the pipeline, use the function\n    `resolve()` to unblock the coroutines awaiting the attached future and do not hand the content\n    to the next stage.\n    As a rule of thumb, sending more items into the pipeline first and awaiting their resolution\n    later is better.\n    \"\"\"\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async for d_content in self.items():\n            d_content.resolve()\n            await self.put(d_content)\n\n\nclass ContentAssociation(Stage):\n    \"\"\"\n    A Stages API stage that associates content units with `new_version`.\n\n    This stage stores all content unit primary keys in memory before running. This is done to\n    compute the units already associated but not received from `self._in_q`. These units are passed\n    via `self._out_q` to the next stage as a :class:`django.db.models.query.QuerySet`.\n\n    This stage creates a ProgressReport named 'Associating Content' that counts the number of units\n    associated. Since it's a stream the total count isn't known until it's finished.\n\n    If `mirror` was enabled, then content units may also be un-assocated (removed) from\n    `new_version`. A ProgressReport named 'Un-Associating Content' is created that counts the number\n    of units un-associated.\n\n    Args:\n        new_version (:class:`~pulpcore.plugin.models.RepositoryVersion`): The repo version this\n            stage associates content with.\n        mirror (bool): Whether or not to \"mirror\" the stream of DeclarativeContent - whether content\n            not in the stream should be removed from the repository.\n        args: unused positional arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.\n        kwargs: unused keyword arguments passed along to :class:`~pulpcore.plugin.stages.Stage`.\n    \"\"\"\n\n    def __init__(self, new_version, mirror, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.new_version = new_version\n        self.allow_delete = mirror\n\n    async def run(self):\n        \"\"\"\n        The coroutine for this stage.\n\n        Returns:\n            The coroutine for this stage.\n        \"\"\"\n        async with ProgressReport(message=\"Associating Content\", code=\"associating.content\") as pb:\n            to_delete = {\n                i\n                async for i in sync_to_async_iterable(\n                    self.new_version.content.values_list(\"pk\", flat=True)\n                )\n            }\n\n            async for batch in self.batches():\n                to_add = set()\n                for d_content in batch:\n                    try:\n                        to_delete.remove(d_content.content.pk)\n                    except KeyError:\n                        to_add.add(d_content.content.pk)\n                        await self.put(d_content)\n\n                if to_add:\n                    await sync_to_async(self.new_version.add_content)(\n                        Content.objects.filter(pk__in=to_add)\n                    )\n                    await pb.aincrease_by(len(to_add))\n\n            if self.allow_delete:\n                async with ProgressReport(\n                    message=\"Un-Associating Content\", code=\"unassociating.content\"\n                ) as pb:\n                    if to_delete:\n                        await sync_to_async(self.new_version.remove_content)(\n                            Content.objects.filter(pk__in=to_delete)\n                        )\n                        await pb.aincrease_by(len(to_delete))\n", "path": "pulpcore/plugin/stages/content_stages.py"}]}