feat: replace threading.Lock() to FIFOLock

Signed-off-by: AnyISalIn <anyisalin@gmail.com>
This commit is contained in:
AnyISalIn 2023-08-21 17:48:56 +08:00
parent 42b72fe246
commit 71a0f6ef85
3 changed files with 45 additions and 4 deletions

View File

@ -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):

37
modules/fifo_lock.py Normal file
View File

@ -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()

View File

@ -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