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

38 lines
1.4 KiB
Python

import dogpile
from ddtrace.pin import Pin, _DD_PIN_NAME, _DD_PIN_PROXY_NAME
from ddtrace.vendor.wrapt import wrap_function_wrapper as _w
from .lock import _wrap_lock_ctor
from .region import _wrap_get_create, _wrap_get_create_multi
_get_or_create = dogpile.cache.region.CacheRegion.get_or_create
_get_or_create_multi = dogpile.cache.region.CacheRegion.get_or_create_multi
_lock_ctor = dogpile.lock.Lock.__init__
def patch():
if getattr(dogpile.cache, '_datadog_patch', False):
return
setattr(dogpile.cache, '_datadog_patch', True)
_w('dogpile.cache.region', 'CacheRegion.get_or_create', _wrap_get_create)
_w('dogpile.cache.region', 'CacheRegion.get_or_create_multi', _wrap_get_create_multi)
_w('dogpile.lock', 'Lock.__init__', _wrap_lock_ctor)
Pin(app='dogpile.cache', service='dogpile.cache').onto(dogpile.cache)
def unpatch():
if not getattr(dogpile.cache, '_datadog_patch', False):
return
setattr(dogpile.cache, '_datadog_patch', False)
# This looks silly but the unwrap util doesn't support class instance methods, even
# though wrapt does. This was causing the patches to stack on top of each other
# during testing.
dogpile.cache.region.CacheRegion.get_or_create = _get_or_create
dogpile.cache.region.CacheRegion.get_or_create_multi = _get_or_create_multi
dogpile.lock.Lock.__init__ = _lock_ctor
setattr(dogpile.cache, _DD_PIN_NAME, None)
setattr(dogpile.cache, _DD_PIN_PROXY_NAME, None)