From 71a0f6ef85f2124d5fb25b55d6ec577a56265fe4 Mon Sep 17 00:00:00 2001 From: AnyISalIn Date: Mon, 21 Aug 2023 17:48:56 +0800 Subject: [PATCH] feat: replace threading.Lock() to FIFOLock Signed-off-by: AnyISalIn --- modules/call_queue.py | 5 ++--- modules/fifo_lock.py | 37 +++++++++++++++++++++++++++++++++++++ modules/progress.py | 7 ++++++- 3 files changed, 45 insertions(+), 4 deletions(-) create mode 100644 modules/fifo_lock.py diff --git a/modules/call_queue.py b/modules/call_queue.py index f2eb17d6166..ddf0d57383c 100644 --- a/modules/call_queue.py +++ b/modules/call_queue.py @@ -1,11 +1,10 @@ from functools import wraps import html -import threading import time -from modules import shared, progress, errors, devices +from modules import shared, progress, errors, devices, fifo_lock -queue_lock = threading.Lock() +queue_lock = fifo_lock.FIFOLock() def wrap_queued_call(func): diff --git a/modules/fifo_lock.py b/modules/fifo_lock.py new file mode 100644 index 00000000000..c35b3ae25a3 --- /dev/null +++ b/modules/fifo_lock.py @@ -0,0 +1,37 @@ +import threading +import collections + + +# reference: https://gist.github.com/vitaliyp/6d54dd76ca2c3cdfc1149d33007dc34a +class FIFOLock(object): + def __init__(self): + self._lock = threading.Lock() + self._inner_lock = threading.Lock() + self._pending_threads = collections.deque() + + def acquire(self, blocking=True): + with self._inner_lock: + lock_acquired = self._lock.acquire(False) + if lock_acquired: + return True + elif not blocking: + return False + + release_event = threading.Event() + self._pending_threads.append(release_event) + + release_event.wait() + return self._lock.acquire() + + def release(self): + with self._inner_lock: + if self._pending_threads: + release_event = self._pending_threads.popleft() + release_event.set() + + self._lock.release() + + __enter__ = acquire + + def __exit__(self, t, v, tb): + self.release() diff --git a/modules/progress.py b/modules/progress.py index e32b59dd5d2..a25a0113274 100644 --- a/modules/progress.py +++ b/modules/progress.py @@ -72,7 +72,12 @@ def progressapi(req: ProgressRequest): completed = req.id_task in finished_tasks if not active: - return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo="In queue..." if queued else "Waiting...") + textinfo = "Waiting..." + if queued: + sorted_queued = sorted(pending_tasks.keys(), key=lambda x: pending_tasks[x]) + queue_index = sorted_queued.index(req.id_task) + textinfo = "In queue: {}/{}".format(queue_index + 1, len(sorted_queued)) + return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo=textinfo) progress = 0