45 lines
1.5 KiB
Python
45 lines
1.5 KiB
Python
import gevent
|
|
from Debug import Debug
|
|
|
|
|
|
class GreenletManager:
|
|
# pool is either gevent.pool.Pool or GreenletManager.
|
|
# if pool is None, new gevent.pool.Pool() is created.
|
|
def __init__(self, pool=None):
|
|
self.greenlets = set()
|
|
if not pool:
|
|
pool = gevent.pool.Pool(None)
|
|
self.pool = pool
|
|
|
|
def _spawn_later(self, seconds, *args, **kwargs):
|
|
# If pool is another GreenletManager, delegate to it.
|
|
if hasattr(self.pool, 'spawnLater'):
|
|
return self.pool.spawnLater(seconds, *args, **kwargs)
|
|
|
|
# There's gevent.spawn_later(), but there isn't gevent.pool.Pool.spawn_later().
|
|
# Doing manually.
|
|
greenlet = self.pool.greenlet_class(*args, **kwargs)
|
|
self.pool.add(greenlet)
|
|
greenlet.start_later(seconds)
|
|
return greenlet
|
|
|
|
def _spawn(self, *args, **kwargs):
|
|
return self.pool.spawn(*args, **kwargs)
|
|
|
|
def spawnLater(self, *args, **kwargs):
|
|
greenlet = self._spawn_later(*args, **kwargs)
|
|
greenlet.link(lambda greenlet: self.greenlets.remove(greenlet))
|
|
self.greenlets.add(greenlet)
|
|
return greenlet
|
|
|
|
def spawn(self, *args, **kwargs):
|
|
greenlet = self._spawn(*args, **kwargs)
|
|
greenlet.link(lambda greenlet: self.greenlets.remove(greenlet))
|
|
self.greenlets.add(greenlet)
|
|
return greenlet
|
|
|
|
def stopGreenlets(self, reason="Stopping all greenlets"):
|
|
num = len(self.greenlets)
|
|
gevent.killall(list(self.greenlets), Debug.createNotifyType(reason), block=False)
|
|
return num
|