Add a setter for the current log context.

Move the resource tracking inside that setter so that it is easier
to make sure that the resource tracking isn't double counting the
resource usage.
This commit is contained in:
Mark Haines 2015-12-07 10:51:18 +00:00
parent 3c85a317d6
commit 3dd1630848

View File

@ -87,13 +87,26 @@ class LoggingContext(object):
"""Get the current logging context from thread local storage""" """Get the current logging context from thread local storage"""
return getattr(cls.thread_local, "current_context", cls.sentinel) return getattr(cls.thread_local, "current_context", cls.sentinel)
@classmethod
def set_current_context(cls, context):
"""Set the current logging context in thread local storage
Args:
context(LoggingContext): The context to activate.
Returns:
The context that was previously active
"""
current = cls.current_context()
if current is not context:
current.stop()
cls.thread_local.current_context = context
context.start()
return current
def __enter__(self): def __enter__(self):
"""Enters this logging context into thread local storage""" """Enters this logging context into thread local storage"""
if self.parent_context is not None: if self.parent_context is not None:
raise Exception("Attempt to enter logging context multiple times") raise Exception("Attempt to enter logging context multiple times")
self.parent_context = self.current_context() self.parent_context = self.set_current_context(self)
self.thread_local.current_context = self
self.start()
return self return self
def __exit__(self, type, value, traceback): def __exit__(self, type, value, traceback):
@ -102,17 +115,16 @@ class LoggingContext(object):
Returns: Returns:
None to avoid suppressing any exeptions that were thrown. None to avoid suppressing any exeptions that were thrown.
""" """
if self.thread_local.current_context is not self: current = self.set_current_context(self.parent_context)
if self.thread_local.current_context is self.sentinel: if current is not self:
if current is self.sentinel:
logger.debug("Expected logging context %s has been lost", self) logger.debug("Expected logging context %s has been lost", self)
else: else:
logger.warn( logger.warn(
"Current logging context %s is not expected context %s", "Current logging context %s is not expected context %s",
self.thread_local.current_context, current,
self self
) )
self.thread_local.current_context = self.parent_context
self.stop()
self.parent_context = None self.parent_context = None
def __getattr__(self, name): def __getattr__(self, name):
@ -194,17 +206,13 @@ class PreserveLoggingContext(object):
def __enter__(self): def __enter__(self):
"""Captures the current logging context""" """Captures the current logging context"""
self.current_context = LoggingContext.current_context() self.current_context = LoggingContext.set_current_context(
if self.new_context is not self.current_context: self.new_context
self.current_context.stop() )
LoggingContext.thread_local.current_context = self.new_context
def __exit__(self, type, value, traceback): def __exit__(self, type, value, traceback):
"""Restores the current logging context""" """Restores the current logging context"""
context = LoggingContext.thread_local.current_context LoggingContext.set_current_context(self.current_context)
LoggingContext.thread_local.current_context = self.current_context
if context is not self.current_context:
self.current_context.start()
if self.current_context is not LoggingContext.sentinel: if self.current_context is not LoggingContext.sentinel:
if self.current_context.parent_context is None: if self.current_context.parent_context is None:
logger.warn( logger.warn(