A simple worker-queue setup. Note that workers quit if queue is empty. That means you have to first put things into the queue before spawning the worker!
| 304 | |
| 305 | |
| 306 | class WorkerQueue: |
| 307 | ''' |
| 308 | A simple worker-queue setup. |
| 309 | |
| 310 | Note that workers quit if queue is empty. That means you have to first put |
| 311 | things into the queue before spawning the worker! |
| 312 | ''' |
| 313 | def __init__(self, max_workers): |
| 314 | self._queue = queue.Queue() |
| 315 | self._workers = [] |
| 316 | self._max_workers = max_workers |
| 317 | self._shutdown_handlers = [] |
| 318 | |
| 319 | # According to http://stackoverflow.com/a/27062830, those are |
| 320 | # threadsafe compared to increasing a simple integer variable. |
| 321 | self.num_done_tasks = itertools.count() |
| 322 | self.num_failed_tasks = itertools.count() |
| 323 | |
| 324 | def shutdown(self): |
| 325 | while self._shutdown_handlers: |
| 326 | try: |
| 327 | self._shutdown_handlers.pop()() |
| 328 | except Exception: |
| 329 | pass |
| 330 | |
| 331 | def _worker(self): |
| 332 | while True: |
| 333 | try: |
| 334 | func = self._queue.get(False) |
| 335 | except queue.Empty: |
| 336 | break |
| 337 | |
| 338 | try: |
| 339 | func(wq=self) |
| 340 | except Exception: |
| 341 | handle_cli_error() |
| 342 | next(self.num_failed_tasks) |
| 343 | finally: |
| 344 | self._queue.task_done() |
| 345 | next(self.num_done_tasks) |
| 346 | if not self._queue.unfinished_tasks: |
| 347 | self.shutdown() |
| 348 | |
| 349 | def spawn_worker(self): |
| 350 | if self._max_workers and len(self._workers) >= self._max_workers: |
| 351 | return |
| 352 | |
| 353 | t = click_threading.Thread(target=self._worker) |
| 354 | t.start() |
| 355 | self._workers.append(t) |
| 356 | |
| 357 | @contextlib.contextmanager |
| 358 | def join(self): |
| 359 | assert self._workers or not self._queue.unfinished_tasks |
| 360 | ui_worker = click_threading.UiWorker() |
| 361 | self._shutdown_handlers.append(ui_worker.shutdown) |
| 362 | _echo = click.echo |
| 363 |