smart_future.py

   1 #!/usr/bin/env python3
   2
   3 from __future__ import annotations
   4 import concurrent
   5 import concurrent.futures as fut
   6 import logging
   7 import traceback
   8 from typing import Callable, List, Set, TypeVar
   9
  10 from overrides import overrides
  11
  12 # This module is commonly used by others in here and should avoid
  13 # taking any unnecessary dependencies back on them.
  14 from deferred_operand import DeferredOperand
  15 import id_generator
  16
  17 logger = logging.getLogger(__name__)
  18
  19 T = TypeVar('T')
  20
  21
  22 def wait_any(
  23     futures: List[SmartFuture],
  24     *,
  25     callback: Callable = None,
  26     log_exceptions: bool = True,
  27 ):
  28     real_futures = []
  29     smart_future_by_real_future = {}
  30     completed_futures: Set[fut.Future] = set()
  31     for x in futures:
  32         assert type(x) == SmartFuture
  33         real_futures.append(x.wrapped_future)
  34         smart_future_by_real_future[x.wrapped_future] = x
  35
  36     while len(completed_futures) != len(real_futures):
  37         newly_completed_futures = concurrent.futures.as_completed(real_futures)
  38         for f in newly_completed_futures:
  39             if callback is not None:
  40                 callback()
  41             completed_futures.add(f)
  42             if log_exceptions and not f.cancelled():
  43                 exception = f.exception()
  44                 if exception is not None:
  45                     logger.warning(
  46                         f'Future {id(f)} raised an unhandled exception and exited.'
  47                     )
  48                     logger.exception(exception)
  49                     raise exception
  50             yield smart_future_by_real_future[f]
  51     if callback is not None:
  52         callback()
  53     return
  54
  55
  56 def wait_all(
  57     futures: List[SmartFuture],
  58     *,
  59     log_exceptions: bool = True,
  60 ) -> None:
  61     real_futures = []
  62     for x in futures:
  63         assert type(x) == SmartFuture
  64         real_futures.append(x.wrapped_future)
  65
  66     (done, not_done) = concurrent.futures.wait(
  67         real_futures, timeout=None, return_when=concurrent.futures.ALL_COMPLETED
  68     )
  69     if log_exceptions:
  70         for f in real_futures:
  71             if not f.cancelled():
  72                 exception = f.exception()
  73                 if exception is not None:
  74                     logger.warning(
  75                         f'Future {id(f)} raised an unhandled exception and exited.'
  76                     )
  77                     logger.exception(exception)
  78                     raise exception
  79     assert len(done) == len(real_futures)
  80     assert len(not_done) == 0
  81
  82
  83 class SmartFuture(DeferredOperand):
  84     """This is a SmartFuture, a class that wraps a normal Future and can
  85     then be used, mostly, like a normal (non-Future) identifier.
  86
  87     Using a FutureWrapper in expressions will block and wait until
  88     the result of the deferred operation is known.
  89     """
  90
  91     def __init__(self, wrapped_future: fut.Future) -> None:
  92         assert type(wrapped_future) == fut.Future
  93         self.wrapped_future = wrapped_future
  94         self.id = id_generator.get("smart_future_id")
  95
  96     def get_id(self) -> int:
  97         return self.id
  98
  99     def is_ready(self) -> bool:
 100         return self.wrapped_future.done()
 101
 102     # You shouldn't have to call this; instead, have a look at defining a
 103     # method on DeferredOperand base class.
 104     @overrides
 105     def _resolve(self, *, timeout=None) -> T:
 106         return self.wrapped_future.result(timeout)