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

These methods allow to set up code to intercept Greenlet objects
creation.  This is useful to implement custom tracing of gevent tasks
and custom variants of `gevent.local()`.

The performance impact is minuscule: one extra fast "if" check in
`Greenlet.__init__`.

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

@@ -117,6 +117,7 @@ cdef class Greenlet(greenlet):
     cdef bint __start_pending(self)
     cdef bint __never_started_or_killed(self)
     cdef bint __start_completed(self)
+    cdef __call_spawn_callbacks(self)
     cdef __handle_death_before_start(self, tuple args)
 
     cdef __cancel_start(self)
@@ -149,6 +150,7 @@ cdef Waiter
 cdef wait
 cdef iwait
 cdef reraise
+cdef set _spawn_callbacks = None
 cpdef GEVENT_CONFIG
 
 

src/gevent/greenlet.py

@@ -264,6 +264,9 @@ def __init__(self, run=None, *args, **kwargs):
             self.spawn_tree_locals = None
             self._spawning_stack_frames = None
 
+        if _spawn_callbacks is not None:
+            self.__call_spawn_callbacks()
+
     @Lazy
     def spawning_stack(self):
         # Store this in the __dict__. We don't use it from the C
@@ -344,6 +347,11 @@ def dead(self):
             "Boolean indicating that the greenlet is dead and will not run again."
             return self.__start_cancelled_by_kill() or self.__started_but_aborted() or greenlet.dead.__get__(self)
 
+    def __call_spawn_callbacks(self):
+        if _spawn_callbacks is not None:
+            for cb in _spawn_callbacks:
+                cb(self)
+
     def __never_started_or_killed(self):
         return self._start_event is None
 
@@ -521,6 +529,41 @@ def start_later(self, seconds):
             self._start_event = self.parent.loop.timer(seconds)
             self._start_event.start(self.switch)
 
+    @staticmethod
+    def add_spawn_callback(cb):
+        """
+        add_spawn_callback(callback) -> None
+
+        Set up a *callback* to be invoked when a new :class:`Greenlet`
+        object is instantiated.
+
+        The invocation order of spawn callbacks is unspecified.  Adding one
+        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:
+            _spawn_callbacks = set()
+        _spawn_callbacks.add(cb)
+
+    @staticmethod
+    def remove_spawn_callback(cb):
+        """
+        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.
+
+        .. versionadded:: 1.3.8
+        """
+        global _spawn_callbacks
+        if _spawn_callbacks is not None:
+            _spawn_callbacks.discard(cb)
+            if not _spawn_callbacks:
+                _spawn_callbacks = None
+
     @classmethod
     def spawn(cls, *args, **kwargs):
         """
@@ -940,5 +983,7 @@ def _init():
 
 _init()
 
+_spawn_callbacks = None
+
 from gevent._util import import_c_accel
 import_c_accel(globals(), 'gevent._greenlet')

src/greentest/test__greenlet.py

@@ -699,6 +699,21 @@ def func():
         while not raw.dead:
             gevent.sleep(0.01)
 
+    def test_add_spawn_callback(self):
+        def cb(gr):
+            gr._called_test = True
+
+        gevent.Greenlet.add_spawn_callback(cb)
+        self.addCleanup(lambda: gevent.Greenlet.remove_spawn_callback(cb))
+
+        g = gevent.spawn(lambda: None)
+        self.assertTrue(hasattr(g, '_called_test'))
+        g.join()
+
+        gevent.Greenlet.remove_spawn_callback(cb)
+        g = gevent.spawn(lambda: None)
+        self.assertFalse(hasattr(g, '_called_test'))
+        g.join()
 
     def test_getframe_value_error(self):
         def get():