add settings and review feedbacks on potential concurrent issues

Author: sylvinus

env.d/development/backend.defaults

@@ -62,7 +62,9 @@ STORAGE_MESSAGE_BLOBS_ENDPOINT_URL=http://objectstorage:9000
 STORAGE_MESSAGE_BLOBS_BUCKET_NAME=msg-blobs
 STORAGE_MESSAGE_BLOBS_ACCESS_KEY=st-messages
 STORAGE_MESSAGE_BLOBS_SECRET_KEY=password
-# Offload blobs older than N days
+# Tiered storage offload (disabled by default, enable to offload blobs to S3)
+TIERED_STORAGE_OFFLOAD_ENABLED=False
+# Offload blobs older than N days (0 = immediately)
 TIERED_STORAGE_OFFLOAD_AFTER_DAYS=3
 
 # OIDC

src/backend/core/management/commands/verify_tiered_storage.py

@@ -405,6 +405,10 @@ def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
         """
         Re-encrypt a single blob with the target key.
 
+        Uses select_for_update to prevent concurrent modifications.
+        For object storage blobs, keeps old content in memory so S3 can
+        be restored if the DB update fails.
+
         Args:
             blob: The blob to re-encrypt
             target_key_id: The encryption key ID to use for re-encryption
@@ -426,24 +430,25 @@ def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
             )
             return "success"
 
-        # Get the current encrypted/unencrypted content
         if blob.storage_location == BlobStorageLocationChoices.POSTGRES:
-            if blob.raw_content is None:
-                self.stdout.write(
-                    self.style.WARNING(
-                        f"  SKIP blob {blob.id}: no content in PostgreSQL"
-                    )
-                )
-                return "skipped"
+            with transaction.atomic():
+                blob = Blob.objects.select_for_update().get(id=blob.id)
+                if blob.encryption_key_id == target_key_id:
+                    return "skipped"
 
-            # Decrypt with old key (or passthrough if key_id=0)
-            decrypted = self.service.decrypt(bytes(blob.raw_content), old_key_id)
+                old_key_id = blob.encryption_key_id
 
-            # Re-encrypt with new key
-            encrypted, new_key_id = self.service.encrypt(decrypted)
+                if blob.raw_content is None:
+                    self.stdout.write(
+                        self.style.WARNING(
+                            f"  SKIP blob {blob.id}: no content in PostgreSQL"
+                        )
+                    )
+                    return "skipped"
+
+                decrypted = self.service.decrypt(bytes(blob.raw_content), old_key_id)
+                encrypted, new_key_id = self.service.encrypt(decrypted)
 
-            # Update blob
-            with transaction.atomic():
                 blob.raw_content = encrypted
                 blob.encryption_key_id = new_key_id
                 blob.save(update_fields=["raw_content", "encryption_key_id"])
@@ -463,31 +468,57 @@ def _re_encrypt_single_blob(self, blob: Blob, target_key_id: int) -> str:
                 )
                 return "skipped"
 
-            # Download and decrypt
             storage_key = self.service.compute_storage_key(bytes(blob.sha256))
+
+            # Keep old content in memory so we can restore S3 if the
+            # transaction fails after the S3 overwrite.
+            old_encrypted = None
+            s3_updated = False
+
             try:
-                with self.service.storage.open(storage_key, "rb") as f:
-                    encrypted_content = f.read()
-            except FileNotFoundError:
-                self.stderr.write(
-                    self.style.ERROR(
-                        f"  ERROR blob {blob.id}: not found in storage at {storage_key}"
-                    )
-                )
-                return "error"
+                with transaction.atomic():
+                    blob = Blob.objects.select_for_update().get(id=blob.id)
+                    if blob.encryption_key_id == target_key_id:
+                        return "skipped"
 
-            decrypted = self.service.decrypt(encrypted_content, old_key_id)
+                    old_key_id = blob.encryption_key_id
 
-            # Re-encrypt with new key
-            encrypted, new_key_id = self.service.encrypt(decrypted)
+                    try:
+                        with self.service.storage.open(storage_key, "rb") as f:
+                            old_encrypted = f.read()
+                    except FileNotFoundError:
+                        self.stderr.write(
+                            self.style.ERROR(
+                                f"  ERROR blob {blob.id}: not found in storage at {storage_key}"
+                            )
+                        )
+                        return "error"
 
-            # Upload new encrypted content (overwrites existing)
-            self.service.storage.save(storage_key, ContentFile(encrypted))
+                    decrypted = self.service.decrypt(old_encrypted, old_key_id)
+                    encrypted, new_key_id = self.service.encrypt(decrypted)
 
-            # Update blob metadata
-            with transaction.atomic():
-                blob.encryption_key_id = new_key_id
-                blob.save(update_fields=["encryption_key_id"])
+                    self.service.storage.save(storage_key, ContentFile(encrypted))
+                    s3_updated = True
+
+                    blob.encryption_key_id = new_key_id
+                    blob.save(update_fields=["encryption_key_id"])
+            except Exception:
+                # If S3 was overwritten but the transaction failed (DB error
+                # or commit failure), restore the old S3 content to prevent
+                # leaving the blob in a corrupted state.
+                if s3_updated and old_encrypted is not None:
+                    try:
+                        self.service.storage.save(
+                            storage_key, ContentFile(old_encrypted)
+                        )
+                    except Exception as restore_err:  # pylint: disable=broad-except
+                        self.stderr.write(
+                            self.style.ERROR(
+                                f"  CRITICAL: failed to restore S3 content for "
+                                f"blob {blob.id}: {restore_err}"
+                            )
+                        )
+                raise
 
             self.stdout.write(
                 f"  Re-encrypted blob {blob.id} (OBJECT_STORAGE): "

src/backend/core/models.py

@@ -1642,23 +1642,6 @@ def get_content(self) -> bytes:
             return pyzstd.decompress(compressed)
         raise ValueError(f"Unsupported compression type: {self.compression}")
 
-    def delete(self, *args, **kwargs):
-        """Delete blob, cleaning up object storage if orphaned."""
-        sha256_copy = bytes(self.sha256)
-        storage_location_copy = self.storage_location
-
-        # Delete from DB first
-        super().delete(*args, **kwargs)
-
-        # Then cleanup storage if was in object storage
-        if storage_location_copy == BlobStorageLocationChoices.OBJECT_STORAGE:
-            # pylint: disable-next=import-outside-toplevel
-            from core.services.tiered_storage import TieredStorageService
-
-            service = TieredStorageService()
-            if service.enabled:
-                service.delete_if_orphaned(sha256_copy)
-
 
 class Attachment(BaseModel):
     """Attachment model to link messages with blobs."""

src/backend/core/services/tiered_storage.py

@@ -30,13 +30,9 @@ class TieredStorageService:
 
     def __init__(self):
         """Initialize the service, checking if object storage is configured."""
-        # Check if message-blobs storage has endpoint_url configured
         self._storage = None
-        self.enabled = bool(
-            settings.STORAGES.get("message-blobs", {})
-            .get("OPTIONS", {})
-            .get("endpoint_url")
-        )
+        opts = settings.STORAGES.get("message-blobs", {}).get("OPTIONS", {})
+        self.enabled = bool(opts.get("endpoint_url") or opts.get("access_key"))
         # encryption_keys is a dict: {"1": "key1", "2": "key2"}
         self.encryption_keys = settings.MESSAGES_BLOB_ENCRYPTION_KEYS or {}
         self.active_key_id = settings.MESSAGES_BLOB_ENCRYPTION_ACTIVE_KEY_ID
@@ -245,6 +241,11 @@ def delete_if_orphaned(self, sha256_bytes: bytes) -> bool:
         """
         Delete storage object only if no other blobs reference it.
 
+        Counts ALL blobs with the same SHA256 (any storage_location), not just
+        OBJECT_STORAGE ones. This prevents a race condition where a concurrent
+        offload task is about to upload a blob with the same SHA256 that is
+        still in POSTGRES state.
+
         Args:
             sha256_bytes: SHA256 hash of the blob to potentially delete
 
@@ -256,11 +257,11 @@ def delete_if_orphaned(self, sha256_bytes: bytes) -> bool:
 
         from core.models import Blob
 
-        # Count remaining references
-        refs = Blob.objects.filter(
-            sha256=sha256_bytes,
-            storage_location=BlobStorageLocationChoices.OBJECT_STORAGE,
-        ).count()
+        # Count ALL remaining references (any storage location).
+        # A blob in POSTGRES with the same SHA256 could be offloaded later
+        # and would need this S3 object. Orphaned S3 objects from blobs that
+        # are never offloaded are cleaned up by the verify command.
+        refs = Blob.objects.filter(sha256=sha256_bytes).count()
 
         if refs > 0:
             logger.debug(

src/backend/core/services/tiered_storage_tasks.py

@@ -12,6 +12,7 @@
 from django.db import OperationalError, transaction
 from django.utils.timezone import now
 
+from botocore.exceptions import BotoCoreError
 from celery.utils.log import get_task_logger
 
 from core.enums import BlobStorageLocationChoices
@@ -22,6 +23,15 @@
 
 logger = get_task_logger(__name__)
 
+# Transient exceptions that should trigger a Celery retry.
+# BotoCoreError covers connection-level errors (timeouts, DNS, etc.).
+# ClientError is intentionally excluded - it covers HTTP API errors
+# (403, 404, etc.) which are usually permanent and shouldn't be retried.
+_TRANSIENT_EXCEPTIONS = (
+    OSError,
+    BotoCoreError,
+)
+
 
 @celery_app.task(bind=True)
 def offload_blobs_task(self) -> Dict[str, Any]:
@@ -36,9 +46,13 @@ def offload_blobs_task(self) -> Dict[str, Any]:
     Returns:
         Dict with status and number of blobs queued
     """
+    if not settings.TIERED_STORAGE_OFFLOAD_ENABLED:
+        logger.debug("Tiered storage offload disabled, skipping")
+        return {"status": "disabled", "queued": 0}
+
     service = TieredStorageService()
     if not service.enabled:
-        logger.debug("Tiered storage not enabled, skipping offload task")
+        logger.debug("Object storage not configured, skipping offload task")
         return {"status": "disabled", "queued": 0}
 
     offload_after_days = settings.TIERED_STORAGE_OFFLOAD_AFTER_DAYS
@@ -65,7 +79,7 @@ def offload_blobs_task(self) -> Dict[str, Any]:
     return {"status": "success", "queued": queued_count}
 
 
-@celery_app.task(bind=True)
+@celery_app.task(bind=True, max_retries=3)
 def offload_single_blob_task(self, blob_id: str) -> Dict[str, Any]:
     """
     Offload a single blob to object storage atomically.
@@ -76,12 +90,18 @@ def offload_single_blob_task(self, blob_id: str) -> Dict[str, Any]:
     3. Updates the blob record and clears raw_content
     4. All within a transaction for atomicity
 
+    Transient S3/network errors are automatically retried with exponential
+    backoff (up to 3 retries).
+
     Args:
         blob_id: UUID of the blob to offload
 
     Returns:
         Dict with status and blob_id
     """
+    if not settings.TIERED_STORAGE_OFFLOAD_ENABLED:
+        return {"status": "disabled", "blob_id": blob_id}
+
     service = TieredStorageService()
     if not service.enabled:
         return {"status": "disabled", "blob_id": blob_id}
@@ -128,6 +148,15 @@ def offload_single_blob_task(self, blob_id: str) -> Dict[str, Any]:
 
             return {"status": "success", "blob_id": blob_id, "key_id": key_id}
 
+    except _TRANSIENT_EXCEPTIONS as e:
+        logger.warning(
+            "Transient error offloading blob %s (attempt %d/%d): %s",
+            blob_id,
+            self.request.retries + 1,
+            self.max_retries + 1,
+            e,
+        )
+        raise self.retry(exc=e, countdown=60 * (2**self.request.retries)) from e
     except Exception as e:  # pylint: disable=broad-except
         logger.exception("Failed to offload blob %s: %s", blob_id, e)
         return {"status": "error", "blob_id": blob_id, "error": str(e)}

src/backend/core/signals.py

@@ -8,12 +8,14 @@
 from django.dispatch import receiver
 
 from core import models
+from core.enums import BlobStorageLocationChoices
 from core.services.identity.keycloak import (
     sync_mailbox_to_keycloak_user,
     sync_maildomain_to_keycloak_group,
 )
 from core.services.search import MESSAGE_INDEX, get_opensearch_client
 from core.services.search.tasks import index_message_task, reindex_thread_task
+from core.services.tiered_storage import TieredStorageService
 
 logger = logging.getLogger(__name__)
 
@@ -104,6 +106,19 @@ def index_thread_post_save(sender, instance, created, **kwargs):
         )
 
 
+@receiver(post_delete, sender=models.Blob)
+def cleanup_blob_storage(sender, instance, **kwargs):
+    """Clean up object storage when a blob is deleted.
+
+    Uses a post_delete signal instead of a model delete() override to ensure
+    cleanup also runs during CASCADE and QuerySet bulk deletes.
+    """
+    if instance.storage_location == BlobStorageLocationChoices.OBJECT_STORAGE:
+        service = TieredStorageService()
+        if service.enabled:
+            service.delete_if_orphaned(bytes(instance.sha256))
+
+
 @receiver(pre_delete, sender=models.Message)
 def delete_message_blobs(sender, instance, **kwargs):
     """Delete the blobs associated with a message."""