Merge pull request #1289 from 1st1/task_hook

Add Greenlet.add_spawn_callback() and Greenlet.remove_spawn_callback()

doc/api/gevent.greenlet.rst

@@ -176,6 +176,8 @@ yet and thus would evaluate to False.
 .. automethod:: Greenlet.rawlink
 .. automethod:: Greenlet.unlink
 .. automethod:: Greenlet.__str__
+.. automethod:: Greenlet.add_spawn_callback
+.. automethod:: Greenlet.remove_spawn_callback
 
 
 Raw greenlet Methods

src/gevent/_greenlet.pxd

@@ -172,3 +172,6 @@ cdef _killall(list greenlets, object exception)
 
 @cython.locals(done=list)
 cpdef joinall(greenlets, timeout=*, raise_error=*, count=*)
+
+cdef set _spawn_callbacks = None
+cdef _call_spawn_callbacks(gr)

src/gevent/greenlet.py

@@ -508,6 +508,7 @@ class Greenlet(greenlet):
     def start(self):
         """Schedule the greenlet to run in this loop iteration"""
         if self._start_event is None:
+            _call_spawn_callbacks(self)
             self._start_event = self.parent.loop.run_callback(self.switch)
 
     def start_later(self, seconds):
@@ -518,9 +519,46 @@ class Greenlet(greenlet):
         *seconds* later
         """
         if self._start_event is None:
+            _call_spawn_callbacks(self)
             self._start_event = self.parent.loop.timer(seconds)
             self._start_event.start(self.switch)
 
+    @staticmethod
+    def add_spawn_callback(callback):
+        """
+        add_spawn_callback(callback) -> None
+
+        Set up a *callback* to be invoked when :class:`Greenlet` objects
+        are started.
+
+        The invocation order of spawn callbacks is unspecified.  Adding the
+        same callback more than one time will not cause it to be called more
+        than once.
+
+        .. versionadded:: 1.3.8
+        """
+        global _spawn_callbacks
+        if _spawn_callbacks is None:  # pylint:disable=used-before-assignment
+            _spawn_callbacks = set()
+        _spawn_callbacks.add(callback)
+
+    @staticmethod
+    def remove_spawn_callback(callback):
+        """
+        remove_spawn_callback(callback) -> None
+
+        Remove *callback* function added with :meth:`Greenlet.add_spawn_callback`.
+        This function will not fail if *callback* has been already removed or
+        if *callback* was never added.
+
+        .. versionadded:: 1.3.8
+        """
+        global _spawn_callbacks
+        if _spawn_callbacks is not None:
+            _spawn_callbacks.discard(callback)
+            if not _spawn_callbacks:
+                _spawn_callbacks = None
+
     @classmethod
     def spawn(cls, *args, **kwargs):
         """
@@ -891,6 +929,15 @@ def _killall(greenlets, exception):
                 g.parent.handle_error(g, *sys_exc_info())
 
 
+def _call_spawn_callbacks(gr):
+    if _spawn_callbacks is not None:
+        for cb in _spawn_callbacks:
+            cb(gr)
+
+
+_spawn_callbacks = None
+
+
 def killall(greenlets, exception=GreenletExit, block=True, timeout=None):
     """
     Forceably terminate all the ``greenlets`` by causing them to raise ``exception``.

src/greentest/test__greenlet.py

@@ -699,6 +699,38 @@ class TestBasic(greentest.TestCase):
         while not raw.dead:
             gevent.sleep(0.01)
 
+    def test_add_spawn_callback(self):
+        called = {'#': 0}
+
+        def cb(gr):
+            called['#'] += 1
+            gr._called_test = True
+
+        gevent.Greenlet.add_spawn_callback(cb)
+        try:
+            g = gevent.spawn(lambda: None)
+            self.assertTrue(hasattr(g, '_called_test'))
+            g.join()
+            self.assertEqual(called['#'], 1)
+
+            g = gevent.spawn_later(1e-5, lambda: None)
+            self.assertTrue(hasattr(g, '_called_test'))
+            g.join()
+            self.assertEqual(called['#'], 2)
+
+            g = gevent.Greenlet(lambda: None)
+            g.start()
+            self.assertTrue(hasattr(g, '_called_test'))
+            g.join()
+            self.assertEqual(called['#'], 3)
+
+            gevent.Greenlet.remove_spawn_callback(cb)
+            g = gevent.spawn(lambda: None)
+            self.assertFalse(hasattr(g, '_called_test'))
+            g.join()
+            self.assertEqual(called['#'], 3)
+        finally:
+            gevent.Greenlet.remove_spawn_callback(cb)
 
     def test_getframe_value_error(self):
         def get():