下面是一个完整的测试代码,用于测试Redis分布式锁的功能。我们将使用给定的Redis URL redis://10.80.0.211:6379/0 来测试锁的获取、释放以及并发控制。

import threading
import time
import random
from functools import wraps

# 使用提供的Redis分布式锁代码
import redis
import uuid
import time
from functools import wraps

class RedisJSONClient:

    REDIS_URL = ''

    _instance = None
    _redis_conn = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self, redis_url=None, **kwargs):
        if not self._redis_conn:
            self._connect(redis_url, **kwargs)

    def _connect(self, redis_url=None, **kwargs):
        try:
            if redis_url:
                self._redis_conn = redis.from_url(redis_url)
            else:
                self._redis_conn = redis.Redis(**kwargs)
            self._redis_conn.ping()
        except Exception as e:
            raise ConnectionError(f"Redis连接失败: {str(e)}")

    @property
    def client(self):
        if self._redis_conn is None:
            raise ConnectionError("Redis连接未初始化")
        return self._redis_conn

    def __del__(self):
        if self._redis_conn:
            self._redis_conn.close()


class RedisDistributedLock:

    LOCK_ACQUIRE_INTERVAL = 0.1

    def __init__(self, redis_client, lock_name, expire_time=600):
        self.redis_client = redis_client.client if hasattr(redis_client, 'client') else redis_client
        self.lock_name = lock_name
        self.expire_time = expire_time
        self.identifier = str(uuid.uuid4())

    def _try_acquire(self):
        return bool(
            self.redis_client.set(
                self.lock_name,
                self.identifier,
                ex=self.expire_time,
                nx=True
            )
        )

    def _acquire_with_timeout(self, timeout):
        end_time = time.time() + timeout
        while time.time() < end_time:
            if self._try_acquire():
                return True
            time.sleep(self.LOCK_ACQUIRE_INTERVAL)
        return False

    def acquire(self, blocking=True, timeout=None):
        if not blocking:
            return self._try_acquire()

        if timeout is not None:
            return self._acquire_with_timeout(timeout)

        while not self._try_acquire():
            time.sleep(self.LOCK_ACQUIRE_INTERVAL)
        return True

    def release(self):
        lua_script = """
        if redis.call("get", KEYS[1]) == ARGV[1] then
            return redis.call("del", KEYS[1])
        else
            return 0
        end
        """
        return self.redis_client.eval(lua_script, 1, self.lock_name, self.identifier)

    def __enter__(self):
        self.acquire(blocking=True)
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.release()


def ensure_execution(lock_name, expire_time=600):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            redis_client = RedisJSONClient(redis_url="redis://10.80.0.211:6379/0")
            with RedisDistributedLock(redis_client, lock_name, expire_time):
                return func(*args, **kwargs)
        return wrapper
    return decorator


# 测试代码开始
def test_distributed_lock():
    # 初始化Redis客户端
    redis_url = "redis://10.80.0.211:6379/0"
    redis_client = RedisJSONClient(redis_url=redis_url)
    
    # 测试锁的基本功能
    print("测试单个锁的获取和释放...")
    lock_name = "test_lock"
    lock = RedisDistributedLock(redis_client, lock_name, expire_time=10)
    
    # 获取锁
    assert lock.acquire(), "获取锁失败"
    print("锁获取成功")
    
    # 尝试再次获取锁(应该失败)
    assert not lock._try_acquire(), "锁被重复获取"
    print("锁无法被重复获取 - 测试通过")
    
    # 释放锁
    assert lock.release() == 1, "锁释放失败"
    print("锁释放成功")
    
    # 再次获取锁(应该成功)
    assert lock.acquire(), "锁再次获取失败"
    print("锁可以再次获取 - 测试通过")
    lock.release()
    
    # 测试上下文管理器
    print("\n测试上下文管理器...")
    with RedisDistributedLock(redis_client, lock_name):
        print("进入上下文管理器,锁已获取")
        # 尝试在另一个线程中获取锁
        def try_get_lock():
            another_lock = RedisDistributedLock(redis_client, lock_name)
            if not another_lock.acquire(blocking=False):
                print("另一个线程无法获取锁 - 测试通过")
            else:
                print("另一个线程获取到了锁 - 测试失败")
                another_lock.release()
        
        thread = threading.Thread(target=try_get_lock)
        thread.start()
        thread.join()
    print("退出上下文管理器,锁已释放")

# 测试并发情况下的锁
def test_concurrent_locks():
    redis_url = "redis://10.80.0.211:6379/0"
    redis_client = RedisJSONClient(redis_url=redis_url)
    lock_name = "concurrent_test_lock"
    
    shared_counter = 0
    threads = []
    num_threads = 5
    
    @ensure_execution(lock_name, expire_time=5)
    def increment_counter():
        nonlocal shared_counter
        # 模拟一些工作
        time.sleep(random.uniform(0.1, 0.3))
        shared_counter += 1
        print(f"线程 {threading.current_thread().name} 增加计数器到 {shared_counter}")
    
    print(f"\n测试并发情况下的锁,启动 {num_threads} 个线程...")
    for i in range(num_threads):
        t = threading.Thread(target=increment_counter, name=f"Thread-{i}")
        threads.append(t)
        t.start()
    
    for t in threads:
        t.join()
    
    print(f"最终计数器值: {shared_counter} (应为 {num_threads})")
    assert shared_counter == num_threads, "并发测试失败,计数器值不正确"

if __name__ == "__main__":
    # 测试基本功能
    test_distributed_lock()
    
    # 测试并发功能
    test_concurrent_locks()
    
    print("\n所有测试完成!")

代码说明:

  1. 基本功能测试 (test_distributed_lock):

    • 测试锁的获取、释放功能
    • 验证锁不能被重复获取
    • 测试上下文管理器的正确性
    • 验证另一个线程在锁被持有时无法获取锁
  2. 并发测试 (test_concurrent_locks):

    • 使用@ensure_execution装饰器确保函数在锁保护下执行
    • 创建多个线程并发执行计数器增加操作
    • 验证最终计数器值是否正确(应该等于线程数量)
  3. Redis连接:

    • 使用提供的Redis URL redis://10.80.0.211:6379/0
    • 通过RedisJSONClient建立连接

预期输出:

  1. 基本功能测试会打印锁的获取、释放状态
  2. 并发测试会显示多个线程有序地增加计数器
  3. 最终计数器值应该等于线程数量(证明锁有效防止了并发冲突)

如果Redis服务器正常运行且可访问,这个测试应该能验证分布式锁的所有关键功能。

Logo

北京人形旗下天工造物具身智能开源社区,聚焦具身天工与慧思开物两大平台

更多推荐