github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/utils/thread_pool_executor.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  # pytype: skip-file
    19  
    20  import queue
    21  import threading
    22  import weakref
    23  from concurrent.futures import _base
    24  
    25  
    26  class _WorkItem(object):
    27    def __init__(self, future, fn, args, kwargs):
    28      self._future = future
    29      self._fn = fn
    30      self._fn_args = args
    31      self._fn_kwargs = kwargs
    32  
    33    def run(self):
    34      if self._future.set_running_or_notify_cancel():
    35        # If the future wasn't cancelled, then attempt to execute it.
    36        try:
    37          self._future.set_result(self._fn(*self._fn_args, **self._fn_kwargs))
    38        except BaseException as exc:
    39          self._future.set_exception(exc)
    40  
    41  
    42  class _Worker(threading.Thread):
    43    def __init__(self, idle_worker_queue, work_item):
    44      super().__init__()
    45      self._idle_worker_queue = idle_worker_queue
    46      self._work_item = work_item
    47      self._wake_semaphore = threading.Semaphore(0)
    48      self._lock = threading.Lock()
    49      self._shutdown = False
    50  
    51    def run(self):
    52      while True:
    53        self._work_item.run()
    54        self._work_item = None
    55  
    56        self._idle_worker_queue.put(self)
    57        self._wake_semaphore.acquire()
    58        if self._work_item is None:
    59          return
    60  
    61    def assign_work(self, work_item):
    62      """Assigns the work item and wakes up the thread.
    63  
    64      This method must only be called while the worker is idle.
    65      """
    66      self._work_item = work_item
    67      self._wake_semaphore.release()
    68  
    69    def shutdown(self):
    70      """Wakes up this thread with a 'None' work item signalling to shutdown."""
    71      self._wake_semaphore.release()
    72  
    73  
    74  class UnboundedThreadPoolExecutor(_base.Executor):
    75    def __init__(self):
    76      self._idle_worker_queue = queue.Queue()
    77      self._max_idle_threads = 16
    78      self._workers = weakref.WeakSet()
    79      self._shutdown = False
    80      self._lock = threading.Lock()  # Guards access to _workers and _shutdown
    81  
    82    def submit(self, fn, *args, **kwargs):
    83      """Attempts to submit the work item.
    84  
    85      A runtime error is raised if the pool has been shutdown.
    86      """
    87      future = _base.Future()
    88      work_item = _WorkItem(future, fn, args, kwargs)
    89      with self._lock:
    90        if self._shutdown:
    91          raise RuntimeError(
    92              'Cannot schedule new tasks after thread pool has been shutdown.')
    93        try:
    94          self._idle_worker_queue.get(block=False).assign_work(work_item)
    95  
    96          # If we have more idle threads then the max allowed, shutdown a thread.
    97          if self._idle_worker_queue.qsize() > self._max_idle_threads:
    98            try:
    99              self._idle_worker_queue.get(block=False).shutdown()
   100            except queue.Empty:
   101              pass
   102        except queue.Empty:
   103          worker = _Worker(self._idle_worker_queue, work_item)
   104          worker.daemon = True
   105          worker.start()
   106          self._workers.add(worker)
   107      return future
   108  
   109    def shutdown(self, wait=True):
   110      with self._lock:
   111        if self._shutdown:
   112          return
   113  
   114        self._shutdown = True
   115  
   116        for worker in self._workers:
   117          worker.shutdown()
   118  
   119        if wait:
   120          for worker in self._workers:
   121            worker.join()
   122  
   123  
   124  class _SharedUnboundedThreadPoolExecutor(UnboundedThreadPoolExecutor):
   125    def shutdown(self, wait=True):
   126      # Prevent shutting down the shared thread pool
   127      pass
   128  
   129  
   130  _SHARED_UNBOUNDED_THREAD_POOL_EXECUTOR = _SharedUnboundedThreadPoolExecutor()
   131  
   132  
   133  def shared_unbounded_instance():
   134    return _SHARED_UNBOUNDED_THREAD_POOL_EXECUTOR