Writing a simple distributed lock from scratch

[rednafi]

import redis
import time
import uuid
from contextlib import contextmanager

# Initialize Redis client
redis_client = redis.Redis(host='localhost', port=6379, db=0)

def acquire_lock(lock_name, lock_timeout=5, acquire_timeout=10):
    """Attempt to acquire a Redis lock."""
    identifier = str(uuid.uuid4())  # A unique identifier for this lock owner
    lock_name = f"lock:{lock_name}"  # Prefix the lock name
    end = time.time() + acquire_timeout

    while time.time() < end:
        # Try to acquire the lock using NX and EX to set if not exists, with a timeout
        if redis_client.set(lock_name, identifier, ex=lock_timeout, nx=True):
            return identifier
        time.sleep(0.01)  # Short sleep to avoid spamming the Redis server

    raise TimeoutError(f"Could not acquire lock for {lock_name}")

def release_lock(lock_name, identifier):
    """Release a Redis lock if the lock is still held by the current identifier."""
    lock_name = f"lock:{lock_name}"
    pipe = redis_client.pipeline(True)

    try:
        pipe.watch(lock_name)
        if pipe.get(lock_name) == identifier.encode():
            pipe.multi()
            pipe.delete(lock_name)
            pipe.execute()
            return True
        pipe.unwatch()
    except redis.exceptions.WatchError:
        pass
    return False

@contextmanager
def redis_lock(lock_name, lock_timeout=5, acquire_timeout=10):
    """Context manager to handle acquiring and releasing Redis locks."""
    identifier = acquire_lock(lock_name, lock_timeout, acquire_timeout)
    try:
        yield
    finally:
        release_lock(lock_name, identifier)

# Critical section: payment processing function
def process_payment(order_id, thread_name):
    try:
        with redis_lock(f"order:{order_id}:lock", lock_timeout=5, acquire_timeout=10):
            print(f"[{thread_name}] Processing payment for order {order_id}")
            time.sleep(3)  # Simulate payment processing time
            print(f"[{thread_name}] Payment processed for order {order_id}")
    except TimeoutError as e:
        print(f"[{thread_name}] {str(e)}")

# Simulate multiple threads trying to process the same order
if __name__ == "__main__":
    import threading

    order_id = 123

    # Create multiple threads trying to process the same order concurrently
    threads = []
    for i in range(5):
        thread_name = f"Thread-{i+1}"
        thread = threading.Thread(target=process_payment, args=(order_id, thread_name))
        threads.append(thread)

    # Start all threads
    for thread in threads:
        thread.start()

    # Wait for all threads to finish
    for thread in threads:
        thread.join()