opt/tools/thread_pool.py

from concurrent.futures import ThreadPoolExecutor
import threading


class MyThreadPoolExecutor(ThreadPoolExecutor):
    def __init__(self, max_workers=None, thread_name_prefix='', *args, **kwargs):
        super().__init__(max_workers, thread_name_prefix, *args, **kwargs)
        self._idle_workers = max_workers
        self._lock = threading.Lock()

    def submit(self, fn, *args, **kwargs):
        with self._lock:
            self._idle_workers -= 1
        future = super().submit(self._wrap_task, fn, *args, **kwargs)
        return future

    def _wrap_task(self, fn, *args, **kwargs):
        try:
            result = fn(*args, **kwargs)
        finally:
            with self._lock:
                self._idle_workers += 1
        return result

    def get_idle_worker_count(self):
        with self._lock:
            return self._idle_workers