From be7c8ca5581dacc55a5fc414d0eb742a0b4b5af3 Mon Sep 17 00:00:00 2001
From: fit2bot <68588906+fit2bot@users.noreply.github.com>
Date: Wed, 8 Dec 2021 17:55:08 +0800
Subject: [PATCH] fix: redis lock bug (#7340)

Co-authored-by: feng626 <1304903146@qq.com>
Co-authored-by: feng626 <57284900+feng626@users.noreply.github.com>
---
 apps/common/utils/lock.py | 13 ++++++++++++-
 1 file changed, 12 insertions(+), 1 deletion(-)

diff --git a/apps/common/utils/lock.py b/apps/common/utils/lock.py
index f52c28b68..f2983e11b 100644
--- a/apps/common/utils/lock.py
+++ b/apps/common/utils/lock.py
@@ -1,7 +1,10 @@
 from functools import wraps
 import threading
 
-from redis_lock import Lock as RedisLock, NotAcquired
+from redis_lock import (
+    Lock as RedisLock, NotAcquired, UNLOCK_SCRIPT,
+    EXTEND_SCRIPT, RESET_SCRIPT, RESET_ALL_SCRIPT
+)
 from redis import Redis
 from django.db import transaction
 
@@ -50,6 +53,7 @@ class DistributedLock(RedisLock):
             auto_renewal = False
 
         super().__init__(redis_client=redis, name='{'+name+'}', expire=expire, auto_renewal=auto_renewal)
+        self.register_scripts(redis)
         self._release_on_transaction_commit = release_on_transaction_commit
         self._release_raise_exc = release_raise_exc
         self._reentrant = reentrant
@@ -73,6 +77,13 @@ class DistributedLock(RedisLock):
                 return func(*args, **kwds)
         return inner
 
+    @classmethod
+    def register_scripts(cls, redis_client):
+        cls.unlock_script = redis_client.register_script(UNLOCK_SCRIPT)
+        cls.extend_script = redis_client.register_script(EXTEND_SCRIPT)
+        cls.reset_script = redis_client.register_script(RESET_SCRIPT)
+        cls.reset_all_script = redis_client.register_script(RESET_ALL_SCRIPT)
+
     def locked_by_me(self):
         if self.locked():
             if self.get_owner_id() == self.id: