Add cache lock in Shcemafile.get_or_create

Utilize a redis cache lock when getting or creating a new SchemaFile in
order to prevent race conditions.

Implement tests.

Refs. TS-2457

Author: Roffenlund

django/thunderstore/schema_server/models/file.py

@@ -8,8 +8,12 @@
 from django.db import models
 from django.utils import timezone
 
+from thunderstore.cache.utils import get_cache
 from thunderstore.core.mixins import S3FileMixin
 
+cache = get_cache("legacy")
+CACHE_LOCK_TIMEOUT = 30
+
 
 def get_schema_file_path(_, filename: str) -> str:
     return f"schema/sha256/{filename}"
@@ -39,27 +43,30 @@ def get_or_create(cls, content: bytes) -> "SchemaFile":
         hash.update(content)
         checksum = hash.hexdigest()
 
-        if existing := cls.objects.filter(checksum_sha256=checksum).first():
-            return existing
+        lock_key = f"lock.schemafile.{checksum}"
+        with cache.lock(lock_key, timeout=CACHE_LOCK_TIMEOUT, blocking_timeout=None):
+            if existing := cls.objects.filter(checksum_sha256=checksum).first():
+                return existing
+
+            gzipped = io.BytesIO()
+            with gzip.GzipFile(fileobj=gzipped, mode="wb") as f:
+                f.write(content)
+            timestamp = timezone.now()
 
-        gzipped = io.BytesIO()
-        with gzip.GzipFile(fileobj=gzipped, mode="wb") as f:
-            f.write(content)
-        timestamp = timezone.now()
+            file = ContentFile(
+                # TODO: This is immediately passed to BytesIO again, meaning
+                #       we're just wasting memory. Find a way to pass this to
+                #       the Django model without the inefficiency.
+                gzipped.getvalue(),
+                name=f"{checksum}.json.gz",
+            )
 
-        file = ContentFile(
-            # TODO: This is immediately passed to BytesIO again, meaning
-            #       we're just wasting memory. Find a way to pass this to
-            #       the Django model without the inefficiency.
-            gzipped.getvalue(),
-            name=f"{checksum}.json.gz",
-        )
-        return cls.objects.create(
-            data=file,
-            content_type="application/json",
-            content_encoding="gzip",
-            last_modified=timestamp,
-            checksum_sha256=checksum,
-            file_size=len(content),
-            gzip_size=file.size,
-        )
+            return cls.objects.create(
+                data=file,
+                content_type="application/json",
+                content_encoding="gzip",
+                last_modified=timestamp,
+                checksum_sha256=checksum,
+                file_size=len(content),
+                gzip_size=file.size,
+            )

django/thunderstore/schema_server/models/tests/test_file.py

@@ -1,6 +1,8 @@
+import concurrent.futures
 import time
 from datetime import timedelta
 from hashlib import sha256
+from unittest.mock import MagicMock, patch
 
 import pytest
 from django.utils import timezone
@@ -92,3 +94,39 @@ def test_schema_server_file_get_or_create_deduplication():
     file_b = SchemaFile.get_or_create(test_data_b)
     assert file_a != file_b
     assert SchemaFile.objects.count() == 2
+
+
+@pytest.mark.django_db
+def test_get_or_create_schema_file_cache_lock_acquired_and_released():
+    test_data = b"Hello world!"
+
+    mock_lock = MagicMock()
+    mock_lock.__enter__.return_value = True
+    mock_lock.__exit__.return_value = None
+
+    with patch("django_redis.cache.RedisCache.lock", return_value=mock_lock):
+        file = SchemaFile.get_or_create(test_data)
+
+    assert file is not None
+    mock_lock.__enter__.assert_called_once()
+    mock_lock.__exit__.assert_called_once()
+
+
+@pytest.mark.django_db(transaction=True)
+def test_get_or_create_schema_file_parallel():
+    test_data = b"Hello world!"
+
+    def call_get_or_create():
+        return SchemaFile.get_or_create(test_data)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+        futures = [executor.submit(call_get_or_create) for _ in range(3)]
+        results = [f.result() for f in futures]
+
+    pks = {file.pk for file in results}
+    assert len(pks) == 1
+
+    assert (
+        SchemaFile.objects.filter(checksum_sha256=results[0].checksum_sha256).count()
+        == 1
+    )