Spaces:
Sleeping
Sleeping
File size: 9,238 Bytes
754d92a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 |
# thanks to:
# https://github.com/oleglpts/PriorityThreadPoolExecutor/blob/master/PriorityThreadPoolExecutor/__init__.py
# https://github.com/oleglpts/PriorityThreadPoolExecutor/issues/4
import atexit
import itertools
import logging
import queue
import random
import sys
import threading
import weakref
from concurrent.futures import _base
from concurrent.futures.thread import BrokenThreadPool
from concurrent.futures.thread import ThreadPoolExecutor
from concurrent.futures.thread import _python_exit
from concurrent.futures.thread import _threads_queues
from concurrent.futures.thread import _WorkItem
from heapq import heappop
from heapq import heappush
logger = logging.getLogger(__name__)
########################################################################################################################
# Global variables #
########################################################################################################################
NULL_ENTRY = (sys.maxsize, _WorkItem(None, None, (), {}))
_shutdown = False
########################################################################################################################
# Before system exit procedure #
########################################################################################################################
def python_exit():
"""
Cleanup before system exit
"""
global _shutdown
_shutdown = True
items = list(_threads_queues.items())
for _t, q in items:
q.put(NULL_ENTRY)
for t, _q in items:
t.join()
# change default cleanup
atexit.unregister(_python_exit)
atexit.register(python_exit)
class PriorityQueue(queue.Queue):
"""Variant of Queue that retrieves open entries in priority order (lowest first).
Entries are typically tuples of the form: (priority number, data).
"""
REMOVED = "<removed-task>"
DEFAULT_PRIORITY = 100
def _init(self, maxsize):
self.queue = []
self.entry_finder = {}
self.counter = itertools.count()
def _qsize(self):
return len(self.queue)
def _put(self, item):
# heappush(self.queue, item)
try:
if item[1] in self.entry_finder:
self.remove(item[1])
count = next(self.counter)
entry = [item[0], count, item[1]]
self.entry_finder[item[1]] = entry
heappush(self.queue, entry)
except TypeError: # handle item==None
self._put((self.DEFAULT_PRIORITY, None))
def remove(self, task):
"""
This simply replaces the data with the REMOVED value,
which will get cleared out once _get reaches it.
"""
entry = self.entry_finder.pop(task)
entry[-1] = self.REMOVED
def _get(self):
while self.queue:
entry = heappop(self.queue)
if entry[2] is not self.REMOVED:
del self.entry_finder[entry[2]]
return entry
return None
def _worker(executor_reference, work_queue, initializer, initargs):
if initializer is not None:
try:
initializer(*initargs)
except BaseException:
_base.LOGGER.critical("Exception in initializer:", exc_info=True)
executor = executor_reference()
if executor is not None:
executor._initializer_failed()
return
try:
while True:
work_item = work_queue.get(block=True)
try:
if work_item[2] is not None:
work_item[2].run()
# Delete references to object. See issue16284
del work_item
# attempt to increment idle count
executor = executor_reference()
if executor is not None:
executor._idle_semaphore.release()
del executor
continue
executor = executor_reference()
# Exit if:
# - The interpreter is shutting down OR
# - The executor that owns the worker has been collected OR
# - The executor that owns the worker has been shutdown.
if _shutdown or executor is None or executor._shutdown:
# Flag the executor as shutting down as early as possible if it
# is not gc-ed yet.
if executor is not None:
executor._shutdown = True
# Notice other workers
work_queue.put(None)
return
del executor
finally:
work_queue.task_done()
except BaseException:
_base.LOGGER.critical("Exception in worker", exc_info=True)
class PriorityThreadPoolExecutor(ThreadPoolExecutor):
"""
Thread pool executor with priority queue (priorities must be different, lowest first)
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# change work queue type to queue.PriorityQueue
self._work_queue: PriorityQueue = PriorityQueue()
self._all_future = []
def submit(self, fn, *args, **kwargs):
"""
Sending the function to the execution queue
:param fn: function being executed
:type fn: callable
:param args: function's positional arguments
:param kwargs: function's keywords arguments
:return: future instance
:rtype: _base.Future
Added keyword:
- priority (integer later sys.maxsize)
"""
with self._shutdown_lock:
if self._broken:
raise BrokenThreadPool(self._broken)
if self._shutdown:
raise RuntimeError("cannot schedule new futures after shutdown")
if _shutdown:
raise RuntimeError(
"cannot schedule new futures after interpreter shutdown"
)
priority = kwargs.get("priority", random.randint(0, sys.maxsize - 1)) # noqa: S311
if "priority" in kwargs:
del kwargs["priority"]
f = _base.Future()
w = _WorkItem(f, fn, args, kwargs)
self._work_queue.put((priority, w))
self._adjust_thread_count()
self._all_future.append(f)
return f
def _adjust_thread_count(self):
# if idle threads are available, don't spin new threads
if self._idle_semaphore.acquire(timeout=0):
return
# When the executor gets lost, the weakref callback will wake up
# the worker threads.
def weakref_cb(_, q=self._work_queue):
q.put(None)
num_threads = len(self._threads)
if num_threads < self._max_workers:
thread_name = f"{self._thread_name_prefix or self}_{num_threads:d}"
t = threading.Thread(
name=thread_name,
target=_worker,
args=(
weakref.ref(self, weakref_cb),
self._work_queue,
self._initializer,
self._initargs,
),
)
t.start()
self._threads.add(t)
_threads_queues[t] = self._work_queue
def shutdown(self, wait=True, *, cancel_futures=False):
logger.debug("Shutting down executor %s", self._thread_name_prefix or self)
if wait:
logger.debug(
"Waiting for all tasks done %s", self._thread_name_prefix or self
)
self._work_queue.join()
logger.debug("All tasks done %s", self._thread_name_prefix or self)
with self._shutdown_lock:
self._shutdown = True
if cancel_futures:
# Drain all work items from the queue, and then cancel their
# associated futures.
while True:
try:
work_item = self._work_queue.get_nowait()
except queue.Empty:
break
if work_item is not None:
work_item.future.cancel()
# Send a wake-up to prevent threads calling
# _work_queue.get(block=True) from permanently blocking.
self._work_queue.put(None)
if wait:
logger.debug(
"Waiting for all thread done %s", self._thread_name_prefix or self
)
for t in self._threads:
self._work_queue.put(None)
t.join()
logger.debug("shutdown finish %s", self._thread_name_prefix or self)
def __del__(self):
for f in self._all_future:
if f.done() and not f.cancelled():
try:
f.result()
except Exception as e:
logger.warning("Exception in future %s: %s", f, e, exc_info=True)
|