greenpool.py 8.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238
  1. import itertools
  2. import traceback
  3. from eventlet import event
  4. from eventlet import greenthread
  5. from eventlet import queue
  6. from eventlet import semaphore
  7. from eventlet.support import greenlets as greenlet
  8. __all__ = ['GreenPool', 'GreenPile']
  9. DEBUG = True
  10. class GreenPool(object):
  11. """The GreenPool class is a pool of green threads.
  12. """
  13. def __init__(self, size=1000):
  14. self.size = size
  15. self.coroutines_running = set()
  16. self.sem = semaphore.Semaphore(size)
  17. self.no_coros_running = event.Event()
  18. def resize(self, new_size):
  19. """ Change the max number of greenthreads doing work at any given time.
  20. If resize is called when there are more than *new_size* greenthreads
  21. already working on tasks, they will be allowed to complete but no new
  22. tasks will be allowed to get launched until enough greenthreads finish
  23. their tasks to drop the overall quantity below *new_size*. Until
  24. then, the return value of free() will be negative.
  25. """
  26. size_delta = new_size - self.size
  27. self.sem.counter += size_delta
  28. self.size = new_size
  29. def running(self):
  30. """ Returns the number of greenthreads that are currently executing
  31. functions in the GreenPool."""
  32. return len(self.coroutines_running)
  33. def free(self):
  34. """ Returns the number of greenthreads available for use.
  35. If zero or less, the next call to :meth:`spawn` or :meth:`spawn_n` will
  36. block the calling greenthread until a slot becomes available."""
  37. return self.sem.counter
  38. def spawn(self, function, *args, **kwargs):
  39. """Run the *function* with its arguments in its own green thread.
  40. Returns the :class:`GreenThread <eventlet.greenthread.GreenThread>`
  41. object that is running the function, which can be used to retrieve the
  42. results.
  43. If the pool is currently at capacity, ``spawn`` will block until one of
  44. the running greenthreads completes its task and frees up a slot.
  45. This function is reentrant; *function* can call ``spawn`` on the same
  46. pool without risk of deadlocking the whole thing.
  47. """
  48. # if reentering an empty pool, don't try to wait on a coroutine freeing
  49. # itself -- instead, just execute in the current coroutine
  50. current = greenthread.getcurrent()
  51. if self.sem.locked() and current in self.coroutines_running:
  52. # a bit hacky to use the GT without switching to it
  53. gt = greenthread.GreenThread(current)
  54. gt.main(function, args, kwargs)
  55. return gt
  56. else:
  57. self.sem.acquire()
  58. gt = greenthread.spawn(function, *args, **kwargs)
  59. if not self.coroutines_running:
  60. self.no_coros_running = event.Event()
  61. self.coroutines_running.add(gt)
  62. gt.link(self._spawn_done)
  63. return gt
  64. def _spawn_n_impl(self, func, args, kwargs, coro):
  65. try:
  66. try:
  67. func(*args, **kwargs)
  68. except (KeyboardInterrupt, SystemExit, greenlet.GreenletExit):
  69. raise
  70. except:
  71. if DEBUG:
  72. traceback.print_exc()
  73. finally:
  74. if coro is None:
  75. return
  76. else:
  77. coro = greenthread.getcurrent()
  78. self._spawn_done(coro)
  79. def spawn_n(self, function, *args, **kwargs):
  80. """Create a greenthread to run the *function*, the same as
  81. :meth:`spawn`. The difference is that :meth:`spawn_n` returns
  82. None; the results of *function* are not retrievable.
  83. """
  84. # if reentering an empty pool, don't try to wait on a coroutine freeing
  85. # itself -- instead, just execute in the current coroutine
  86. current = greenthread.getcurrent()
  87. if self.sem.locked() and current in self.coroutines_running:
  88. self._spawn_n_impl(function, args, kwargs, None)
  89. else:
  90. self.sem.acquire()
  91. g = greenthread.spawn_n(self._spawn_n_impl,
  92. function, args, kwargs, True)
  93. if not self.coroutines_running:
  94. self.no_coros_running = event.Event()
  95. self.coroutines_running.add(g)
  96. def waitall(self):
  97. """Waits until all greenthreads in the pool are finished working."""
  98. assert greenthread.getcurrent() not in self.coroutines_running, \
  99. "Calling waitall() from within one of the "\
  100. "GreenPool's greenthreads will never terminate."
  101. if self.running():
  102. self.no_coros_running.wait()
  103. def _spawn_done(self, coro):
  104. self.sem.release()
  105. if coro is not None:
  106. self.coroutines_running.remove(coro)
  107. # if done processing (no more work is waiting for processing),
  108. # we can finish off any waitall() calls that might be pending
  109. if self.sem.balance == self.size:
  110. self.no_coros_running.send(None)
  111. def waiting(self):
  112. """Return the number of greenthreads waiting to spawn.
  113. """
  114. if self.sem.balance < 0:
  115. return -self.sem.balance
  116. else:
  117. return 0
  118. def _do_map(self, func, it, gi):
  119. for args in it:
  120. gi.spawn(func, *args)
  121. gi.spawn(return_stop_iteration)
  122. def starmap(self, function, iterable):
  123. """This is the same as :func:`itertools.starmap`, except that *func* is
  124. executed in a separate green thread for each item, with the concurrency
  125. limited by the pool's size. In operation, starmap consumes a constant
  126. amount of memory, proportional to the size of the pool, and is thus
  127. suited for iterating over extremely long input lists.
  128. """
  129. if function is None:
  130. function = lambda *a: a
  131. gi = GreenMap(self.size)
  132. greenthread.spawn_n(self._do_map, function, iterable, gi)
  133. return gi
  134. def imap(self, function, *iterables):
  135. """This is the same as :func:`itertools.imap`, and has the same
  136. concurrency and memory behavior as :meth:`starmap`.
  137. It's quite convenient for, e.g., farming out jobs from a file::
  138. def worker(line):
  139. return do_something(line)
  140. pool = GreenPool()
  141. for result in pool.imap(worker, open("filename", 'r')):
  142. print result
  143. """
  144. return self.starmap(function, itertools.izip(*iterables))
  145. def return_stop_iteration():
  146. return StopIteration()
  147. class GreenPile(object):
  148. """GreenPile is an abstraction representing a bunch of I/O-related tasks.
  149. Construct a GreenPile with an existing GreenPool object. The GreenPile will
  150. then use that pool's concurrency as it processes its jobs. There can be
  151. many GreenPiles associated with a single GreenPool.
  152. A GreenPile can also be constructed standalone, not associated with any
  153. GreenPool. To do this, construct it with an integer size parameter instead
  154. of a GreenPool.
  155. It is not advisable to iterate over a GreenPile in a different greenthread
  156. than the one which is calling spawn. The iterator will exit early in that
  157. situation.
  158. """
  159. def __init__(self, size_or_pool=1000):
  160. if isinstance(size_or_pool, GreenPool):
  161. self.pool = size_or_pool
  162. else:
  163. self.pool = GreenPool(size_or_pool)
  164. self.waiters = queue.LightQueue()
  165. self.used = False
  166. self.counter = 0
  167. def spawn(self, func, *args, **kw):
  168. """Runs *func* in its own green thread, with the result available by
  169. iterating over the GreenPile object."""
  170. self.used = True
  171. self.counter += 1
  172. try:
  173. gt = self.pool.spawn(func, *args, **kw)
  174. self.waiters.put(gt)
  175. except:
  176. self.counter -= 1
  177. raise
  178. def __iter__(self):
  179. return self
  180. def next(self):
  181. """Wait for the next result, suspending the current greenthread until it
  182. is available. Raises StopIteration when there are no more results."""
  183. if self.counter == 0 and self.used:
  184. raise StopIteration()
  185. try:
  186. return self.waiters.get().wait()
  187. finally:
  188. self.counter -= 1
  189. # this is identical to GreenPile but it blocks on spawn if the results
  190. # aren't consumed, and it doesn't generate its own StopIteration exception,
  191. # instead relying on the spawning process to send one in when it's done
  192. class GreenMap(GreenPile):
  193. def __init__(self, size_or_pool):
  194. super(GreenMap, self).__init__(size_or_pool)
  195. self.waiters = queue.LightQueue(maxsize=self.pool.size)
  196. def next(self):
  197. try:
  198. val = self.waiters.get().wait()
  199. if isinstance(val, StopIteration):
  200. raise val
  201. else:
  202. return val
  203. finally:
  204. self.counter -= 1