src/python/grpcio/grpc/framework/foundation/logging_pool.py

   1 # Copyright 2015 gRPC authors.
   2 #
   3 # Licensed under the Apache License, Version 2.0 (the "License");
   4 # you may not use this file except in compliance with the License.
   5 # You may obtain a copy of the License at
   6 #
   7 #     http://www.apache.org/licenses/LICENSE-2.0
   8 #
   9 # Unless required by applicable law or agreed to in writing, software
  10 # distributed under the License is distributed on an "AS IS" BASIS,
  11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12 # See the License for the specific language governing permissions and
  13 # limitations under the License.
  14 """A thread pool that logs exceptions raised by tasks executed within it."""
  15
  16 import logging
  17
  18 from concurrent import futures
  19
  20 _LOGGER = logging.getLogger(__name__)
  21
  22
  23 def _wrap(behavior):
  24     """Wraps an arbitrary callable behavior in exception-logging."""
  25
  26     def _wrapping(*args, **kwargs):
  27         try:
  28             return behavior(*args, **kwargs)
  29         except Exception:
  30             _LOGGER.exception(
  31                 'Unexpected exception from %s executed in logging pool!',
  32                 behavior)
  33             raise
  34
  35     return _wrapping
  36
  37
  38 class _LoggingPool(object):
  39     """An exception-logging futures.ThreadPoolExecutor-compatible thread pool."""
  40
  41     def __init__(self, backing_pool):
  42         self._backing_pool = backing_pool
  43
  44     def __enter__(self):
  45         return self
  46
  47     def __exit__(self, exc_type, exc_val, exc_tb):
  48         self._backing_pool.shutdown(wait=True)
  49
  50     def submit(self, fn, *args, **kwargs):
  51         return self._backing_pool.submit(_wrap(fn), *args, **kwargs)
  52
  53     def map(self, func, *iterables, **kwargs):
  54         return self._backing_pool.map(_wrap(func),
  55                                       *iterables,
  56                                       timeout=kwargs.get('timeout', None))
  57
  58     def shutdown(self, wait=True):
  59         self._backing_pool.shutdown(wait=wait)
  60
  61
  62 def pool(max_workers):
  63     """Creates a thread pool that logs exceptions raised by the tasks within it.
  64
  65   Args:
  66     max_workers: The maximum number of worker threads to allow the pool.
  67
  68   Returns:
  69     A futures.ThreadPoolExecutor-compatible thread pool that logs exceptions
  70       raised by the tasks executed within it.
  71   """
  72     return _LoggingPool(futures.ThreadPoolExecutor(max_workers))