Handle edge cases in task diagnostics handling

When an Artifact fails to save (likely due to a memory report being
identical between two tasks by chance), we should be handling the case
instead of failing out faster.

closes #7020

Author: dralley

CHANGES/7020.bugfix

@@ -0,0 +1 @@
+Improve handling of some edge cases in task diagnostics code.

pulpcore/tasking/_util.py

@@ -10,8 +10,6 @@
 import tempfile
 from gettext import gettext as _
 
-from contextlib import suppress
-
 from django.conf import settings
 from django.db import connection, transaction, IntegrityError
 from django.db.models import Q
@@ -167,12 +165,16 @@ def __memory_diagnostic_decorator(task):
 
         stop_event.set()
         artifact = Artifact.init_and_validate(mem_diagnostics_file_path)
-        with suppress(IntegrityError):
+        try:
+            # it is possible for the diagnostic artifact (memory report) to be identical to
+            # a previous report, in which case we need to handle the case where saving a new
+            # artifact fails.
             artifact.save()
-            ProfileArtifact.objects.get_or_create(
-                artifact=artifact, name="memory_profile", task=task
-            )
-            _logger.info("Created memory diagnostic data.")
+        except IntegrityError:
+            artifact = Artifact.objects.get(sha256=artifact.sha256)
+
+        ProfileArtifact.objects.get_or_create(artifact=artifact, name="memory_profile", task=task)
+        _logger.info("Created memory diagnostic data.")
 
     return __memory_diagnostic_decorator
 
@@ -191,12 +193,18 @@ def __pyinstrument_diagnostic_decorator(task):
                 f.flush()
 
             artifact = Artifact.init_and_validate(str(profile_file_path))
-            with suppress(IntegrityError):
+            try:
+                # it is possible for the diagnostic artifact (memory report) to be identical to
+                # a previous report, in which case we need to handle the case where saving a new
+                # artifact fails.
                 artifact.save()
-                ProfileArtifact.objects.get_or_create(
-                    artifact=artifact, name="pyinstrument_profile", task=task
-                )
-                _logger.info("Created pyinstrument profile data.")
+            except IntegrityError:
+                artifact = Artifact.objects.get(sha256=artifact.sha256)
+
+            ProfileArtifact.objects.get_or_create(
+                artifact=artifact, name="pyinstrument_profile", task=task
+            )
+            _logger.info("Created pyinstrument profile data.")
         else:
             func(task)
 
@@ -217,12 +225,18 @@ def __memray_diagnostic_decorator(task):
                 func(task)
 
             artifact = Artifact.init_and_validate(str(profile_file_path))
-            with suppress(IntegrityError):
+            try:
+                # it is possible for the diagnostic artifact (memory report) to be identical to
+                # a previous report, in which case we need to handle the case where saving a new
+                # artifact fails.
                 artifact.save()
-                ProfileArtifact.objects.get_or_create(
-                    artifact=artifact, name="memray_profile", task=task
-                )
-                _logger.info("Created memray memory profile data.")
+            except IntegrityError:
+                artifact = Artifact.objects.get(sha256=artifact.sha256)
+
+            ProfileArtifact.objects.get_or_create(
+                artifact=artifact, name="memray_profile", task=task
+            )
+            _logger.info("Created memray memory profile data.")
         else:
             func(task)