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 from functools import wraps
import html import html
import threading
import time 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): 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 completed = req.id_task in finished_tasks
if not active: 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 progress = 0