dagpool.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602
  1. # @file dagpool.py
  2. # @author Nat Goodspeed
  3. # @date 2016-08-08
  4. # @brief Provide DAGPool class
  5. from eventlet.event import Event
  6. from eventlet import greenthread
  7. import six
  8. import collections
  9. # value distinguished from any other Python value including None
  10. _MISSING = object()
  11. class Collision(Exception):
  12. """
  13. DAGPool raises Collision when you try to launch two greenthreads with the
  14. same key, or post() a result for a key corresponding to a greenthread, or
  15. post() twice for the same key. As with KeyError, str(collision) names the
  16. key in question.
  17. """
  18. pass
  19. class PropagateError(Exception):
  20. """
  21. When a DAGPool greenthread terminates with an exception instead of
  22. returning a result, attempting to retrieve its value raises
  23. PropagateError.
  24. Attributes:
  25. key
  26. the key of the greenthread which raised the exception
  27. exc
  28. the exception object raised by the greenthread
  29. """
  30. def __init__(self, key, exc):
  31. # initialize base class with a reasonable string message
  32. msg = "PropagateError({0}): {1}: {2}" \
  33. .format(key, exc.__class__.__name__, exc)
  34. super(PropagateError, self).__init__(msg)
  35. self.msg = msg
  36. # Unless we set args, this is unpickleable:
  37. # https://bugs.python.org/issue1692335
  38. self.args = (key, exc)
  39. self.key = key
  40. self.exc = exc
  41. def __str__(self):
  42. return self.msg
  43. class DAGPool(object):
  44. """
  45. A DAGPool is a pool that constrains greenthreads, not by max concurrency,
  46. but by data dependencies.
  47. This is a way to implement general DAG dependencies. A simple dependency
  48. tree (flowing in either direction) can straightforwardly be implemented
  49. using recursion and (e.g.)
  50. :meth:`GreenThread.imap() <eventlet.greenthread.GreenThread.imap>`.
  51. What gets complicated is when a given node depends on several other nodes
  52. as well as contributing to several other nodes.
  53. With DAGPool, you concurrently launch all applicable greenthreads; each
  54. will proceed as soon as it has all required inputs. The DAG is implicit in
  55. which items are required by each greenthread.
  56. Each greenthread is launched in a DAGPool with a key: any value that can
  57. serve as a Python dict key. The caller also specifies an iterable of other
  58. keys on which this greenthread depends. This iterable may be empty.
  59. The greenthread callable must accept (key, results), where:
  60. key
  61. is its own key
  62. results
  63. is an iterable of (key, value) pairs.
  64. A newly-launched DAGPool greenthread is entered immediately, and can
  65. perform any necessary setup work. At some point it will iterate over the
  66. (key, value) pairs from the passed 'results' iterable. Doing so blocks the
  67. greenthread until a value is available for each of the keys specified in
  68. its initial dependencies iterable. These (key, value) pairs are delivered
  69. in chronological order, *not* the order in which they are initially
  70. specified: each value will be delivered as soon as it becomes available.
  71. The value returned by a DAGPool greenthread becomes the value for its
  72. key, which unblocks any other greenthreads waiting on that key.
  73. If a DAGPool greenthread terminates with an exception instead of returning
  74. a value, attempting to retrieve the value raises :class:`PropagateError`,
  75. which binds the key of the original greenthread and the original
  76. exception. Unless the greenthread attempting to retrieve the value handles
  77. PropagateError, that exception will in turn be wrapped in a PropagateError
  78. of its own, and so forth. The code that ultimately handles PropagateError
  79. can follow the chain of PropagateError.exc attributes to discover the flow
  80. of that exception through the DAG of greenthreads.
  81. External greenthreads may also interact with a DAGPool. See :meth:`wait_each`,
  82. :meth:`waitall`, :meth:`post`.
  83. It is not recommended to constrain external DAGPool producer greenthreads
  84. in a :class:`GreenPool <eventlet.greenpool.GreenPool>`: it may be hard to
  85. provably avoid deadlock.
  86. .. automethod:: __init__
  87. .. automethod:: __getitem__
  88. """
  89. _Coro = collections.namedtuple("_Coro", ("greenthread", "pending"))
  90. def __init__(self, preload={}):
  91. """
  92. DAGPool can be prepopulated with an initial dict or iterable of (key,
  93. value) pairs. These (key, value) pairs are of course immediately
  94. available for any greenthread that depends on any of those keys.
  95. """
  96. try:
  97. # If a dict is passed, copy it. Don't risk a subsequent
  98. # modification to passed dict affecting our internal state.
  99. iteritems = six.iteritems(preload)
  100. except AttributeError:
  101. # Not a dict, just an iterable of (key, value) pairs
  102. iteritems = preload
  103. # Load the initial dict
  104. self.values = dict(iteritems)
  105. # track greenthreads
  106. self.coros = {}
  107. # The key to blocking greenthreads is the Event.
  108. self.event = Event()
  109. def waitall(self):
  110. """
  111. waitall() blocks the calling greenthread until there is a value for
  112. every DAGPool greenthread launched by :meth:`spawn`. It returns a dict
  113. containing all :class:`preload data <DAGPool>`, all data from
  114. :meth:`post` and all values returned by spawned greenthreads.
  115. See also :meth:`wait`.
  116. """
  117. # waitall() is an alias for compatibility with GreenPool
  118. return self.wait()
  119. def wait(self, keys=_MISSING):
  120. """
  121. *keys* is an optional iterable of keys. If you omit the argument, it
  122. waits for all the keys from :class:`preload data <DAGPool>`, from
  123. :meth:`post` calls and from :meth:`spawn` calls: in other words, all
  124. the keys of which this DAGPool is aware.
  125. wait() blocks the calling greenthread until all of the relevant keys
  126. have values. wait() returns a dict whose keys are the relevant keys,
  127. and whose values come from the *preload* data, from values returned by
  128. DAGPool greenthreads or from :meth:`post` calls.
  129. If a DAGPool greenthread terminates with an exception, wait() will
  130. raise :class:`PropagateError` wrapping that exception. If more than
  131. one greenthread terminates with an exception, it is indeterminate
  132. which one wait() will raise.
  133. If an external greenthread posts a :class:`PropagateError` instance,
  134. wait() will raise that PropagateError. If more than one greenthread
  135. posts PropagateError, it is indeterminate which one wait() will raise.
  136. See also :meth:`wait_each_success`, :meth:`wait_each_exception`.
  137. """
  138. # This is mostly redundant with wait_each() functionality.
  139. return dict(self.wait_each(keys))
  140. def wait_each(self, keys=_MISSING):
  141. """
  142. *keys* is an optional iterable of keys. If you omit the argument, it
  143. waits for all the keys from :class:`preload data <DAGPool>`, from
  144. :meth:`post` calls and from :meth:`spawn` calls: in other words, all
  145. the keys of which this DAGPool is aware.
  146. wait_each() is a generator producing (key, value) pairs as a value
  147. becomes available for each requested key. wait_each() blocks the
  148. calling greenthread until the next value becomes available. If the
  149. DAGPool was prepopulated with values for any of the relevant keys, of
  150. course those can be delivered immediately without waiting.
  151. Delivery order is intentionally decoupled from the initial sequence of
  152. keys: each value is delivered as soon as it becomes available. If
  153. multiple keys are available at the same time, wait_each() delivers
  154. each of the ready ones in arbitrary order before blocking again.
  155. The DAGPool does not distinguish between a value returned by one of
  156. its own greenthreads and one provided by a :meth:`post` call or *preload* data.
  157. The wait_each() generator terminates (raises StopIteration) when all
  158. specified keys have been delivered. Thus, typical usage might be:
  159. ::
  160. for key, value in dagpool.wait_each(keys):
  161. # process this ready key and value
  162. # continue processing now that we've gotten values for all keys
  163. By implication, if you pass wait_each() an empty iterable of keys, it
  164. returns immediately without yielding anything.
  165. If the value to be delivered is a :class:`PropagateError` exception object, the
  166. generator raises that PropagateError instead of yielding it.
  167. See also :meth:`wait_each_success`, :meth:`wait_each_exception`.
  168. """
  169. # Build a local set() and then call _wait_each().
  170. return self._wait_each(self._get_keyset_for_wait_each(keys))
  171. def wait_each_success(self, keys=_MISSING):
  172. """
  173. wait_each_success() filters results so that only success values are
  174. yielded. In other words, unlike :meth:`wait_each`, wait_each_success()
  175. will not raise :class:`PropagateError`. Not every provided (or
  176. defaulted) key will necessarily be represented, though naturally the
  177. generator will not finish until all have completed.
  178. In all other respects, wait_each_success() behaves like :meth:`wait_each`.
  179. """
  180. for key, value in self._wait_each_raw(self._get_keyset_for_wait_each(keys)):
  181. if not isinstance(value, PropagateError):
  182. yield key, value
  183. def wait_each_exception(self, keys=_MISSING):
  184. """
  185. wait_each_exception() filters results so that only exceptions are
  186. yielded. Not every provided (or defaulted) key will necessarily be
  187. represented, though naturally the generator will not finish until
  188. all have completed.
  189. Unlike other DAGPool methods, wait_each_exception() simply yields
  190. :class:`PropagateError` instances as values rather than raising them.
  191. In all other respects, wait_each_exception() behaves like :meth:`wait_each`.
  192. """
  193. for key, value in self._wait_each_raw(self._get_keyset_for_wait_each(keys)):
  194. if isinstance(value, PropagateError):
  195. yield key, value
  196. def _get_keyset_for_wait_each(self, keys):
  197. """
  198. wait_each(), wait_each_success() and wait_each_exception() promise
  199. that if you pass an iterable of keys, the method will wait for results
  200. from those keys -- but if you omit the keys argument, the method will
  201. wait for results from all known keys. This helper implements that
  202. distinction, returning a set() of the relevant keys.
  203. """
  204. if keys is not _MISSING:
  205. return set(keys)
  206. else:
  207. # keys arg omitted -- use all the keys we know about
  208. return set(six.iterkeys(self.coros)) | set(six.iterkeys(self.values))
  209. def _wait_each(self, pending):
  210. """
  211. When _wait_each() encounters a value of PropagateError, it raises it.
  212. In all other respects, _wait_each() behaves like _wait_each_raw().
  213. """
  214. for key, value in self._wait_each_raw(pending):
  215. yield key, self._value_or_raise(value)
  216. @staticmethod
  217. def _value_or_raise(value):
  218. # Most methods attempting to deliver PropagateError should raise that
  219. # instead of simply returning it.
  220. if isinstance(value, PropagateError):
  221. raise value
  222. return value
  223. def _wait_each_raw(self, pending):
  224. """
  225. pending is a set() of keys for which we intend to wait. THIS SET WILL
  226. BE DESTRUCTIVELY MODIFIED: as each key acquires a value, that key will
  227. be removed from the passed 'pending' set.
  228. _wait_each_raw() does not treat a PropagateError instance specially:
  229. it will be yielded to the caller like any other value.
  230. In all other respects, _wait_each_raw() behaves like wait_each().
  231. """
  232. while True:
  233. # Before even waiting, show caller any (key, value) pairs that
  234. # are already available. Copy 'pending' because we want to be able
  235. # to remove items from the original set while iterating.
  236. for key in pending.copy():
  237. value = self.values.get(key, _MISSING)
  238. if value is not _MISSING:
  239. # found one, it's no longer pending
  240. pending.remove(key)
  241. yield (key, value)
  242. if not pending:
  243. # Once we've yielded all the caller's keys, done.
  244. break
  245. # There are still more keys pending, so wait.
  246. self.event.wait()
  247. def spawn(self, key, depends, function, *args, **kwds):
  248. """
  249. Launch the passed *function(key, results, ...)* as a greenthread,
  250. passing it:
  251. - the specified *key*
  252. - an iterable of (key, value) pairs
  253. - whatever other positional args or keywords you specify.
  254. Iterating over the *results* iterable behaves like calling
  255. :meth:`wait_each(depends) <DAGPool.wait_each>`.
  256. Returning from *function()* behaves like
  257. :meth:`post(key, return_value) <DAGPool.post>`.
  258. If *function()* terminates with an exception, that exception is wrapped
  259. in :class:`PropagateError` with the greenthread's *key* and (effectively) posted
  260. as the value for that key. Attempting to retrieve that value will
  261. raise that PropagateError.
  262. Thus, if the greenthread with key 'a' terminates with an exception,
  263. and greenthread 'b' depends on 'a', when greenthread 'b' attempts to
  264. iterate through its *results* argument, it will encounter
  265. PropagateError. So by default, an uncaught exception will propagate
  266. through all the downstream dependencies.
  267. If you pass :meth:`spawn` a key already passed to spawn() or :meth:`post`, spawn()
  268. raises :class:`Collision`.
  269. """
  270. if key in self.coros or key in self.values:
  271. raise Collision(key)
  272. # The order is a bit tricky. First construct the set() of keys.
  273. pending = set(depends)
  274. # It's important that we pass to _wait_each() the same 'pending' set()
  275. # that we store in self.coros for this key. The generator-iterator
  276. # returned by _wait_each() becomes the function's 'results' iterable.
  277. newcoro = greenthread.spawn(self._wrapper, function, key,
  278. self._wait_each(pending),
  279. *args, **kwds)
  280. # Also capture the same (!) set in the new _Coro object for this key.
  281. # We must be able to observe ready keys being removed from the set.
  282. self.coros[key] = self._Coro(newcoro, pending)
  283. def _wrapper(self, function, key, results, *args, **kwds):
  284. """
  285. This wrapper runs the top-level function in a DAGPool greenthread,
  286. posting its return value (or PropagateError) to the DAGPool.
  287. """
  288. try:
  289. # call our passed function
  290. result = function(key, results, *args, **kwds)
  291. except Exception as err:
  292. # Wrap any exception it may raise in a PropagateError.
  293. result = PropagateError(key, err)
  294. finally:
  295. # function() has returned (or terminated with an exception). We no
  296. # longer need to track this greenthread in self.coros. Remove it
  297. # first so post() won't complain about a running greenthread.
  298. del self.coros[key]
  299. try:
  300. # as advertised, try to post() our return value
  301. self.post(key, result)
  302. except Collision:
  303. # if we've already post()ed a result, oh well
  304. pass
  305. # also, in case anyone cares...
  306. return result
  307. def spawn_many(self, depends, function, *args, **kwds):
  308. """
  309. spawn_many() accepts a single *function* whose parameters are the same
  310. as for :meth:`spawn`.
  311. The difference is that spawn_many() accepts a dependency dict
  312. *depends*. A new greenthread is spawned for each key in the dict. That
  313. dict key's value should be an iterable of other keys on which this
  314. greenthread depends.
  315. If the *depends* dict contains any key already passed to :meth:`spawn`
  316. or :meth:`post`, spawn_many() raises :class:`Collision`. It is
  317. indeterminate how many of the other keys in *depends* will have
  318. successfully spawned greenthreads.
  319. """
  320. # Iterate over 'depends' items, relying on self.spawn() not to
  321. # context-switch so no one can modify 'depends' along the way.
  322. for key, deps in six.iteritems(depends):
  323. self.spawn(key, deps, function, *args, **kwds)
  324. def kill(self, key):
  325. """
  326. Kill the greenthread that was spawned with the specified *key*.
  327. If no such greenthread was spawned, raise KeyError.
  328. """
  329. # let KeyError, if any, propagate
  330. self.coros[key].greenthread.kill()
  331. # once killed, remove it
  332. del self.coros[key]
  333. def post(self, key, value, replace=False):
  334. """
  335. post(key, value) stores the passed *value* for the passed *key*. It
  336. then causes each greenthread blocked on its results iterable, or on
  337. :meth:`wait_each(keys) <DAGPool.wait_each>`, to check for new values.
  338. A waiting greenthread might not literally resume on every single
  339. post() of a relevant key, but the first post() of a relevant key
  340. ensures that it will resume eventually, and when it does it will catch
  341. up with all relevant post() calls.
  342. Calling post(key, value) when there is a running greenthread with that
  343. same *key* raises :class:`Collision`. If you must post(key, value) instead of
  344. letting the greenthread run to completion, you must first call
  345. :meth:`kill(key) <DAGPool.kill>`.
  346. The DAGPool implicitly post()s the return value from each of its
  347. greenthreads. But a greenthread may explicitly post() a value for its
  348. own key, which will cause its return value to be discarded.
  349. Calling post(key, value, replace=False) (the default *replace*) when a
  350. value for that key has already been posted, by any means, raises
  351. :class:`Collision`.
  352. Calling post(key, value, replace=True) when a value for that key has
  353. already been posted, by any means, replaces the previously-stored
  354. value. However, that may make it complicated to reason about the
  355. behavior of greenthreads waiting on that key.
  356. After a post(key, value1) followed by post(key, value2, replace=True),
  357. it is unspecified which pending :meth:`wait_each([key...]) <DAGPool.wait_each>`
  358. calls (or greenthreads iterating over *results* involving that key)
  359. will observe *value1* versus *value2*. It is guaranteed that
  360. subsequent wait_each([key...]) calls (or greenthreads spawned after
  361. that point) will observe *value2*.
  362. A successful call to
  363. post(key, :class:`PropagateError(key, ExceptionSubclass) <PropagateError>`)
  364. ensures that any subsequent attempt to retrieve that key's value will
  365. raise that PropagateError instance.
  366. """
  367. # First, check if we're trying to post() to a key with a running
  368. # greenthread.
  369. # A DAGPool greenthread is explicitly permitted to post() to its
  370. # OWN key.
  371. coro = self.coros.get(key, _MISSING)
  372. if coro is not _MISSING and coro.greenthread is not greenthread.getcurrent():
  373. # oh oh, trying to post a value for running greenthread from
  374. # some other greenthread
  375. raise Collision(key)
  376. # Here, either we're posting a value for a key with no greenthread or
  377. # we're posting from that greenthread itself.
  378. # Has somebody already post()ed a value for this key?
  379. # Unless replace == True, this is a problem.
  380. if key in self.values and not replace:
  381. raise Collision(key)
  382. # Either we've never before posted a value for this key, or we're
  383. # posting with replace == True.
  384. # update our database
  385. self.values[key] = value
  386. # and wake up pending waiters
  387. self.event.send()
  388. # The comment in Event.reset() says: "it's better to create a new
  389. # event rather than reset an old one". Okay, fine. We do want to be
  390. # able to support new waiters, so create a new Event.
  391. self.event = Event()
  392. def __getitem__(self, key):
  393. """
  394. __getitem__(key) (aka dagpool[key]) blocks until *key* has a value,
  395. then delivers that value.
  396. """
  397. # This is a degenerate case of wait_each(). Construct a tuple
  398. # containing only this 'key'. wait_each() will yield exactly one (key,
  399. # value) pair. Return just its value.
  400. for _, value in self.wait_each((key,)):
  401. return value
  402. def get(self, key, default=None):
  403. """
  404. get() returns the value for *key*. If *key* does not yet have a value,
  405. get() returns *default*.
  406. """
  407. return self._value_or_raise(self.values.get(key, default))
  408. def keys(self):
  409. """
  410. Return a snapshot tuple of keys for which we currently have values.
  411. """
  412. # Explicitly return a copy rather than an iterator: don't assume our
  413. # caller will finish iterating before new values are posted.
  414. return tuple(six.iterkeys(self.values))
  415. def items(self):
  416. """
  417. Return a snapshot tuple of currently-available (key, value) pairs.
  418. """
  419. # Don't assume our caller will finish iterating before new values are
  420. # posted.
  421. return tuple((key, self._value_or_raise(value))
  422. for key, value in six.iteritems(self.values))
  423. def running(self):
  424. """
  425. Return number of running DAGPool greenthreads. This includes
  426. greenthreads blocked while iterating through their *results* iterable,
  427. that is, greenthreads waiting on values from other keys.
  428. """
  429. return len(self.coros)
  430. def running_keys(self):
  431. """
  432. Return keys for running DAGPool greenthreads. This includes
  433. greenthreads blocked while iterating through their *results* iterable,
  434. that is, greenthreads waiting on values from other keys.
  435. """
  436. # return snapshot; don't assume caller will finish iterating before we
  437. # next modify self.coros
  438. return tuple(six.iterkeys(self.coros))
  439. def waiting(self):
  440. """
  441. Return number of waiting DAGPool greenthreads, that is, greenthreads
  442. still waiting on values from other keys. This explicitly does *not*
  443. include external greenthreads waiting on :meth:`wait`,
  444. :meth:`waitall`, :meth:`wait_each`.
  445. """
  446. # n.b. if Event would provide a count of its waiters, we could say
  447. # something about external greenthreads as well.
  448. # The logic to determine this count is exactly the same as the general
  449. # waiting_for() call.
  450. return len(self.waiting_for())
  451. # Use _MISSING instead of None as the default 'key' param so we can permit
  452. # None as a supported key.
  453. def waiting_for(self, key=_MISSING):
  454. """
  455. waiting_for(key) returns a set() of the keys for which the DAGPool
  456. greenthread spawned with that *key* is still waiting. If you pass a
  457. *key* for which no greenthread was spawned, waiting_for() raises
  458. KeyError.
  459. waiting_for() without argument returns a dict. Its keys are the keys
  460. of DAGPool greenthreads still waiting on one or more values. In the
  461. returned dict, the value of each such key is the set of other keys for
  462. which that greenthread is still waiting.
  463. This method allows diagnosing a "hung" DAGPool. If certain
  464. greenthreads are making no progress, it's possible that they are
  465. waiting on keys for which there is no greenthread and no :meth:`post` data.
  466. """
  467. # We may have greenthreads whose 'pending' entry indicates they're
  468. # waiting on some keys even though values have now been posted for
  469. # some or all of those keys, because those greenthreads have not yet
  470. # regained control since values were posted. So make a point of
  471. # excluding values that are now available.
  472. available = set(six.iterkeys(self.values))
  473. if key is not _MISSING:
  474. # waiting_for(key) is semantically different than waiting_for().
  475. # It's just that they both seem to want the same method name.
  476. coro = self.coros.get(key, _MISSING)
  477. if coro is _MISSING:
  478. # Hmm, no running greenthread with this key. But was there
  479. # EVER a greenthread with this key? If not, let KeyError
  480. # propagate.
  481. self.values[key]
  482. # Oh good, there's a value for this key. Either the
  483. # greenthread finished, or somebody posted a value. Just say
  484. # the greenthread isn't waiting for anything.
  485. return set()
  486. else:
  487. # coro is the _Coro for the running greenthread with the
  488. # specified key.
  489. return coro.pending - available
  490. # This is a waiting_for() call, i.e. a general query rather than for a
  491. # specific key.
  492. # Start by iterating over (key, coro) pairs in self.coros. Generate
  493. # (key, pending) pairs in which 'pending' is the set of keys on which
  494. # the greenthread believes it's waiting, minus the set of keys that
  495. # are now available. Filter out any pair in which 'pending' is empty,
  496. # that is, that greenthread will be unblocked next time it resumes.
  497. # Make a dict from those pairs.
  498. return dict((key, pending)
  499. for key, pending in ((key, (coro.pending - available))
  500. for key, coro in six.iteritems(self.coros))
  501. if pending)