Files
2020-04-08 10:39:44 -07:00

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)