Context manager for Twisted

Twisted.py

from functools import wraps
from ddtrace import tracer, ext

from twisted.internet import defer
from twisted.python.context import get, call, theContextTracker


_SPAN_CONTEXT = 'span_context'
_CONTEXT_OBJECT = object()
_original_deferred_init = defer.Deferred.__init__
_original_deferred_runCallbacks = defer.Deferred._runCallbacks


def patch():
	if not getattr(defer, '__apm_patch', False):
		defer.__apm_patch = True
		defer.Deferred.__init__ = patched_deferred_init
		defer.Deferred._runCallbacks = patched_deferred_runCallbacks


def unpatch():
	if getattr(defer, '__apm_patch', False):
		defer.__apm_patch = False
		defer.Deferred.__init__ = _original_deferred_init
		defer.Deferred._runCallbacks = _original_deferred_runCallbacks


def get_contexts():
	return theContextTracker.currentContext().contexts


def patched_deferred_init(self, *args, **kwargs):
	_original_deferred_init(self, *args, **kwargs)
	self._context = get_contexts()[-1]


def patched_deferred_runCallbacks(self, *args, **kwargs):
	return call(
		self._context,
		_original_deferred_runCallbacks, self, *args, **kwargs
	)


def get_current_span():
	return get(_SPAN_CONTEXT)


def create_span(name, service, resource, span_type):
	return tracer.start_span(
		name,
		child_of=get_current_span(),
		service=service,
		resource=resource,
		span_type=span_type
	)


def run_in_context(span, fn, *args, **kwargs):
	return call(
		{_SPAN_CONTEXT: span},
		defer.maybeDeferred, fn, *args, **kwargs
	)