greenpool.py 9.0 KB

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