59 lines
2.3 KiB
Python
59 lines
2.3 KiB
Python
import gevent
|
|
import gevent.pool as gpool
|
|
|
|
from .provider import CONTEXT_ATTR
|
|
|
|
GEVENT_VERSION = gevent.version_info[0:3]
|
|
|
|
|
|
class TracingMixin(object):
|
|
def __init__(self, *args, **kwargs):
|
|
# get the current Context if available
|
|
current_g = gevent.getcurrent()
|
|
ctx = getattr(current_g, CONTEXT_ATTR, None)
|
|
|
|
# create the Greenlet as usual
|
|
super(TracingMixin, self).__init__(*args, **kwargs)
|
|
|
|
# the context is always available made exception of the main greenlet
|
|
if ctx:
|
|
# create a new context that inherits the current active span
|
|
new_ctx = ctx.clone()
|
|
setattr(self, CONTEXT_ATTR, new_ctx)
|
|
|
|
|
|
class TracedGreenlet(TracingMixin, gevent.Greenlet):
|
|
"""
|
|
``Greenlet`` class that is used to replace the original ``gevent``
|
|
class. This class is supposed to do ``Context`` replacing operation, so
|
|
that any greenlet inherits the context from the parent Greenlet.
|
|
When a new greenlet is spawned from the main greenlet, a new instance
|
|
of ``Context`` is created. The main greenlet is not affected by this behavior.
|
|
|
|
There is no need to inherit this class to create or optimize greenlets
|
|
instances, because this class replaces ``gevent.greenlet.Greenlet``
|
|
through the ``patch()`` method. After the patch, extending the gevent
|
|
``Greenlet`` class means extending automatically ``TracedGreenlet``.
|
|
"""
|
|
def __init__(self, *args, **kwargs):
|
|
super(TracedGreenlet, self).__init__(*args, **kwargs)
|
|
|
|
|
|
class TracedIMapUnordered(TracingMixin, gpool.IMapUnordered):
|
|
def __init__(self, *args, **kwargs):
|
|
super(TracedIMapUnordered, self).__init__(*args, **kwargs)
|
|
|
|
|
|
if GEVENT_VERSION >= (1, 3) or GEVENT_VERSION < (1, 1):
|
|
# For gevent <1.1 and >=1.3, IMap is its own class, so we derive
|
|
# from TracingMixin
|
|
class TracedIMap(TracingMixin, gpool.IMap):
|
|
def __init__(self, *args, **kwargs):
|
|
super(TracedIMap, self).__init__(*args, **kwargs)
|
|
else:
|
|
# For gevent >=1.1 and <1.3, IMap derives from IMapUnordered, so we derive
|
|
# from TracedIMapUnordered and get tracing that way
|
|
class TracedIMap(gpool.IMap, TracedIMapUnordered):
|
|
def __init__(self, *args, **kwargs):
|
|
super(TracedIMap, self).__init__(*args, **kwargs)
|