Skip to content

Commit 46c684b

Browse files
committed
feat: replace threading.Lock() to FIFOLock
Signed-off-by: AnyISalIn <[email protected]>
1 parent 42b72fe commit 46c684b

File tree

3 files changed

+46
-3
lines changed

3 files changed

+46
-3
lines changed

modules/call_queue.py

+2-2
Original file line numberDiff line numberDiff line change
@@ -3,9 +3,9 @@
33
import threading
44
import time
55

6-
from modules import shared, progress, errors, devices
6+
from modules import shared, progress, errors, devices, fifo_lock
77

8-
queue_lock = threading.Lock()
8+
queue_lock = fifo_lock.FIFOLock()
99

1010

1111
def wrap_queued_call(func):

modules/fifo_lock.py

+38
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,38 @@
1+
import threading
2+
import time
3+
import collections
4+
5+
6+
# reference: https://gist.github.com/vitaliyp/6d54dd76ca2c3cdfc1149d33007dc34a
7+
class FIFOLock(object):
8+
def __init__(self):
9+
self._lock = threading.Lock()
10+
self._inner_lock = threading.Lock()
11+
self._pending_threads = collections.deque()
12+
13+
def acquire(self, blocking=True):
14+
with self._inner_lock:
15+
lock_acquired = self._lock.acquire(False)
16+
if lock_acquired:
17+
return True
18+
elif not blocking:
19+
return False
20+
21+
release_event = threading.Event()
22+
self._pending_threads.append(release_event)
23+
24+
release_event.wait()
25+
return self._lock.acquire()
26+
27+
def release(self):
28+
with self._inner_lock:
29+
if self._pending_threads:
30+
release_event = self._pending_threads.popleft()
31+
release_event.set()
32+
33+
self._lock.release()
34+
35+
__enter__ = acquire
36+
37+
def __exit__(self, t, v, tb):
38+
self.release()

modules/progress.py

+6-1
Original file line numberDiff line numberDiff line change
@@ -72,7 +72,12 @@ def progressapi(req: ProgressRequest):
7272
completed = req.id_task in finished_tasks
7373

7474
if not active:
75-
return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo="In queue..." if queued else "Waiting...")
75+
textinfo = "Waiting..."
76+
if queued:
77+
sorted_queued = sorted(pending_tasks.keys(), key=lambda x: pending_tasks[x])
78+
queue_index = sorted_queued.index(req.id_task)
79+
textinfo = "In queue: {}/{}".format(queue_index + 1, len(sorted_queued))
80+
return ProgressResponse(active=active, queued=queued, completed=completed, id_live_preview=-1, textinfo=textinfo)
7681

7782
progress = 0
7883

0 commit comments

Comments
 (0)