python 協程 自定義互斥鎖

zhut96發表於2024-07-25
Python

最近在用python的一款非同步web框架sanic搭建web服務，遇到一個需要加特定鎖的場景：同一使用者併發處理訂單時需要排隊處理，但不同使用者不需要排隊。

如果僅僅使用async with asyncio.Lock()的話。會使所有請求都排隊處理。

import asyncio
import datetime

lock = asyncio.Lock()


async def place_order(user_id, order_id):
    async with lock:
        # 模擬下單處理
        print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Processing order {order_id} for user {user_id}")
        await asyncio.sleep(1)  # 假設處理需要 1 秒
        print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Order {order_id} for user {user_id} is done")


# 定義一個測試函式
async def test():
    # 建立四個任務，模擬兩個 UserID 的併發請求
    tasks = [
        asyncio.create_task(place_order(1, 101)),
        asyncio.create_task(place_order(1, 102)),
        asyncio.create_task(place_order(2, 201)),
        asyncio.create_task(place_order(2, 202)),
    ]
    # 等待所有任務完成
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    # 執行測試函式
    asyncio.run(test())

這顯然不是想要的結果，第二種方案是定義一個字典，key使用user_id，value為asyncio.Lock()，每次執行前從字典裡面獲取lock，相同的user_id將會使用同一個lock，那就實現了功能。

import asyncio
import datetime

locks = {}


async def place_order(user_id, order_id):
    if user_id not in locks:
        locks[user_id] = asyncio.Lock()
    async with locks[user_id]:
        # 模擬下單處理
        print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Processing order {order_id} for user {user_id}")
        await asyncio.sleep(1)  # 假設處理需要 1 秒
        print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Order {order_id} for user {user_id} is done")


# 定義一個測試函式
async def test():
    # 建立四個任務，模擬兩個 UserID 的併發請求
    tasks = [
        asyncio.create_task(place_order(1, 101)),
        asyncio.create_task(place_order(1, 102)),
        asyncio.create_task(place_order(2, 201)),
        asyncio.create_task(place_order(2, 202)),
    ]
    # 等待所有任務完成
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    # 執行測試函式
    asyncio.run(test())

但是這個方案會有缺點是，user_id執行完成之後沒有釋放資源，當請求的user_id變多之後，勢必會造成佔用過多的資源。繼續改進方案，將locks的value加一個計數器，當獲取lock時計數器加1，使用完之後計數器-1，當計數器變為小於等於0時，釋放locks對應的key。最後將這個功能封裝為一個類方便其他地方呼叫。

import asyncio

mutex_locks = {}


class MutexObj:
    def __init__(self):
        self.lock = asyncio.Lock()
        self.count = 0


class Mutex:
    def __init__(self, key: str):
        if key not in mutex_locks:
            mutex_locks[key] = MutexObj()
        self.__mutex_obj = mutex_locks[key]
        self.__key = key

    def lock(self):
        """
        獲取鎖
        :return:
        """
        self.__mutex_obj.count += 1
        return self.__mutex_obj.lock

    def release(self):
        """
        釋放鎖
        :return:
        """
        self.__mutex_obj.count -= 1
        if self.__mutex_obj.count <= 0:
            del mutex_locks[self.__key]
import asyncio
import datetime

from utils.mutex import Mutex, mutex_locks

locks = {}


async def place_order(user_id, order_id):
    mutex = Mutex(user_id)
    async with mutex.lock():
        try:
            # 模擬下單處理
            print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Processing order {order_id} for user {user_id}")
            await asyncio.sleep(1)  # 假設處理需要 1 秒
            print(f"{datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S:%f')} Order {order_id} for user {user_id} is done")
        except Exception as ex:
            print(ex)
        finally:
            mutex.release()
            print('=====================')
            print(mutex_locks)
            print('=====================')


# 定義一個測試函式
async def test():
    # 建立四個任務，模擬兩個 UserID 的併發請求
    tasks = [
        asyncio.create_task(place_order(1, 101)),
        asyncio.create_task(place_order(1, 102)),
        asyncio.create_task(place_order(2, 201)),
        asyncio.create_task(place_order(2, 202)),
    ]
    # 等待所有任務完成
    await asyncio.gather(*tasks)


if __name__ == '__main__':
    # 執行測試函式
    asyncio.run(test())

至此實現了我的需求，此方案只考慮了單應用場景，如果是分散式部署，需要更換方案如redis鎖，這裡暫且不考慮。如果有其他實現方式，歡迎留言交流。

相關文章