archival.ArchiveBox/archivebox/workers/semaphores.py

# import uuid
# from functools import wraps
# from django.db import connection, transaction
# from django.utils import timezone
# from huey.exceptions import TaskLockedException

# from archivebox.config import CONSTANTS

# class SqliteSemaphore:
#     def __init__(self, db_path, table_name, name, value=1, timeout=None):
#         self.db_path = db_path
#         self.table_name = table_name
#         self.name = name
#         self.value = value
#         self.timeout = timeout or 86400  # Set a max age for lock holders

#         # Ensure the table exists
#         with connection.cursor() as cursor:
#             cursor.execute(f"""
#                 CREATE TABLE IF NOT EXISTS {self.table_name} (
#                     id TEXT PRIMARY KEY,
#                     name TEXT,
#                     timestamp DATETIME
#                 )
#             """)

#     def acquire(self, name=None):
#         name = name or str(uuid.uuid4())
#         now = timezone.now()
#         expiration = now - timezone.timedelta(seconds=self.timeout)

#         with transaction.atomic():
#             # Remove expired locks
#             with connection.cursor() as cursor:
#                 cursor.execute(f"""
#                     DELETE FROM {self.table_name}
#                     WHERE name = %s AND timestamp < %s
#                 """, [self.name, expiration])

#             # Try to acquire the lock
#             with connection.cursor() as cursor:
#                 cursor.execute(f"""
#                     INSERT INTO {self.table_name} (id, name, timestamp)
#                     SELECT %s, %s, %s
#                     WHERE (
#                         SELECT COUNT(*) FROM {self.table_name}
#                         WHERE name = %s
#                     ) < %s
#                 """, [name, self.name, now, self.name, self.value])

#                 if cursor.rowcount > 0:
#                     return name

#         # If we couldn't acquire the lock, remove our attempted entry
#         with connection.cursor() as cursor:
#             cursor.execute(f"""
#                 DELETE FROM {self.table_name}
#                 WHERE id = %s AND name = %s
#             """, [name, self.name])

#         return None

#     def release(self, name):
#         with connection.cursor() as cursor:
#             cursor.execute(f"""
#                 DELETE FROM {self.table_name}
#                 WHERE id = %s AND name = %s
#             """, [name, self.name])
#         return cursor.rowcount > 0


# LOCKS_DB_PATH = CONSTANTS.DATABASE_FILE.parent / 'locks.sqlite3'


# def lock_task_semaphore(db_path, table_name, lock_name, value=1, timeout=None):
#     """
#     Lock which can be acquired multiple times (default = 1).

#     NOTE: no provisions are made for blocking, waiting, or notifying. This is
#     just a lock which can be acquired a configurable number of times.

#     Example:

#     # Allow up to 3 workers to run this task concurrently. If the task is
#     # locked, retry up to 2 times with a delay of 60s.
#     @huey.task(retries=2, retry_delay=60)
#     @lock_task_semaphore('path/to/db.sqlite3', 'semaphore_locks', 'my-lock', 3)
#     def my_task():
#         ...
#     """
#     sem = SqliteSemaphore(db_path, table_name, lock_name, value, timeout)
#     def decorator(fn):
#         @wraps(fn)
#         def inner(*args, **kwargs):
#             tid = sem.acquire()
#             if tid is None:
#                 raise TaskLockedException(f'unable to acquire lock {lock_name}')
#             try:
#                 return fn(*args, **kwargs)
#             finally:
#                 sem.release(tid)
#         return inner
#     return decorator