From c497392d5ba556ae0128579edb2a39449d3b5613 Mon Sep 17 00:00:00 2001
From: blissful <blissful@sunsetglow.net>
Date: Wed, 25 Oct 2023 12:23:37 -0400
Subject: [PATCH] add a lock() function for preventing concurrent writes

---
 rose/cache.py      | 20 ++++++++++++++++++++
 rose/cache.sql     |  7 +++++++
 rose/cache_test.py | 31 +++++++++++++++++++++++++++++--
 3 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/rose/cache.py b/rose/cache.py
index 1a0b1a6..42558ea 100644
--- a/rose/cache.py
+++ b/rose/cache.py
@@ -93,6 +93,26 @@ def migrate_database(c: Config) -> None:
         )
 
 
+@contextmanager
+def lock(c: Config, name: str, timeout: float = 1.0) -> Iterator[None]:
+    try:
+        valid_until = time.time() + timeout
+        while True:
+            with connect(c) as conn:
+                cursor = conn.execute("SELECT MAX(valid_until) FROM locks WHERE name = ?", (name,))
+                row = cursor.fetchone()
+                if not row or not row[0] or row[0] < time.time():
+                    conn.execute(
+                        "INSERT INTO locks (name, valid_until) VALUES (?, ?)", (name, valid_until)
+                    )
+                    break
+                time.sleep(max(0, row[0] - time.time()))
+        yield
+    finally:
+        with connect(c) as conn:
+            conn.execute("DELETE FROM locks WHERE name = ?", (name,))
+
+
 @dataclass
 class CachedArtist:
     name: str
diff --git a/rose/cache.sql b/rose/cache.sql
index d41ebf7..8bbccec 100644
--- a/rose/cache.sql
+++ b/rose/cache.sql
@@ -1,3 +1,10 @@
+CREATE TABLE locks (
+    name TEXT,
+    -- Unix epoch.
+    valid_until REAL NOT NULL,
+    PRIMARY KEY (name, valid_until)
+);
+
 CREATE TABLE release_type_enum (value TEXT PRIMARY KEY);
 INSERT INTO release_type_enum (value) VALUES
     ('album'),
diff --git a/rose/cache_test.py b/rose/cache_test.py
index 2d7cd15..43418be 100644
--- a/rose/cache_test.py
+++ b/rose/cache_test.py
@@ -1,5 +1,6 @@
 import hashlib
 import shutil
+import time
 from dataclasses import asdict
 from pathlib import Path
 
@@ -29,6 +30,7 @@
     list_genres,
     list_labels,
     list_releases,
+    lock,
     migrate_database,
     release_exists,
     track_exists,
@@ -40,7 +42,7 @@
 
 
 def test_schema(config: Config) -> None:
-    # Test that the schema successfully bootstraps.
+    """Test that the schema successfully bootstraps."""
     with CACHE_SCHEMA_PATH.open("rb") as fp:
         schema_hash = hashlib.sha256(fp.read()).hexdigest()
     migrate_database(config)
@@ -52,7 +54,7 @@ def test_schema(config: Config) -> None:
 
 
 def test_migration(config: Config) -> None:
-    # Test that "migrating" the database correctly migrates it.
+    """Test that "migrating" the database correctly migrates it."""
     config.cache_database_path.unlink()
     with connect(config) as conn:
         conn.execute(
@@ -80,6 +82,31 @@ def test_migration(config: Config) -> None:
         assert cursor.fetchone()[0] == 1
 
 
+def test_locks(config: Config) -> None:
+    """Test that taking locks works."""
+    lock_name = "lol"
+
+    # Test that the locking and timeout work.
+    start = time.time()
+    with lock(config, lock_name, timeout=0.2):
+        lock1_acq = time.time()
+        with lock(config, lock_name, timeout=0.2):
+            lock2_acq = time.time()
+    # Assert that we had to wait ~0.1sec to get the second lock.
+    assert lock1_acq - start < 0.05
+    assert lock2_acq - lock1_acq > 0.19
+
+    # Test that releasing a lock actually works.
+    start = time.time()
+    with lock(config, lock_name, timeout=0.2):
+        lock1_acq = time.time()
+    with lock(config, lock_name, timeout=0.2):
+        lock2_acq = time.time()
+    # Assert that we had to wait negligible time to get the second lock.
+    assert lock1_acq - start < 0.05
+    assert lock2_acq - lock1_acq < 0.05
+
+
 def test_update_cache_all(config: Config) -> None:
     """Test that the update all function works."""
     shutil.copytree(TEST_RELEASE_1, config.music_source_dir / TEST_RELEASE_1.name)