From 45540196237c3731ddf7382ef97b902cae3c8955 Mon Sep 17 00:00:00 2001 From: Reiley Yang Date: Fri, 22 Mar 2019 12:36:57 -0700 Subject: [PATCH] Refactor execution_context to use RuntimeContext (#573) Migrate execution_context to RuntimeContext --- opencensus/stats/execution_context.py | 13 +++--- opencensus/tags/execution_context.py | 12 +++--- opencensus/trace/execution_context.py | 62 +++++++++++---------------- setup.py | 1 + 4 files changed, 40 insertions(+), 48 deletions(-) diff --git a/opencensus/stats/execution_context.py b/opencensus/stats/execution_context.py index d5d34399e..2b06c9abd 100644 --- a/opencensus/stats/execution_context.py +++ b/opencensus/stats/execution_context.py @@ -12,18 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -import threading +from opencensus.common.runtime_context import RuntimeContext -_thread_local = threading.local() +_measure_to_view_map_slot = RuntimeContext.register_slot( + 'measure_to_view_map', + lambda: {}) def get_measure_to_view_map(): - return getattr(_thread_local, 'measure_to_view_map', {}) + return RuntimeContext.measure_to_view_map def set_measure_to_view_map(measure_to_view_map): - setattr(_thread_local, 'measure_to_view_map', measure_to_view_map) + RuntimeContext.measure_to_view_map = measure_to_view_map def clear(): - _thread_local.__dict__.clear() + """Clear the context, used in test.""" + _measure_to_view_map_slot.clear() diff --git a/opencensus/tags/execution_context.py b/opencensus/tags/execution_context.py index df774c9c9..a9403d444 100644 --- a/opencensus/tags/execution_context.py +++ b/opencensus/tags/execution_context.py @@ -12,19 +12,19 @@ # See the License for the specific language governing permissions and # limitations under the License. -import threading +from opencensus.common.runtime_context import RuntimeContext -_thread_local = threading.local() +_current_tag_map_slot = RuntimeContext.register_slot('current_tag_map', None) def get_current_tag_map(): - return getattr(_thread_local, 'current_tag_map', None) + return RuntimeContext.current_tag_map def set_current_tag_map(current_tag_map): - setattr(_thread_local, 'current_tag_map', current_tag_map) + RuntimeContext.current_tag_map = current_tag_map def clear(): - """Clear the thread local, used in test.""" - _thread_local.__dict__.clear() + """Clear the context, used in test.""" + _current_tag_map_slot.clear() diff --git a/opencensus/trace/execution_context.py b/opencensus/trace/execution_context.py index bfe96e8bf..89344c6f4 100644 --- a/opencensus/trace/execution_context.py +++ b/opencensus/trace/execution_context.py @@ -12,81 +12,69 @@ # See the License for the specific language governing permissions and # limitations under the License. -import threading - +from opencensus.common.runtime_context import RuntimeContext from opencensus.trace.tracers import noop_tracer -_thread_local = threading.local() +_attrs_slot = RuntimeContext.register_slot('attrs', lambda: {}) +_current_span_slot = RuntimeContext.register_slot('current_span', None) +_tracer_slot = RuntimeContext.register_slot('tracer', noop_tracer.NoopTracer()) def get_opencensus_tracer(): - """Get the opencensus tracer from thread local.""" - return getattr(_thread_local, 'tracer', noop_tracer.NoopTracer()) + """Get the opencensus tracer from runtime context.""" + return RuntimeContext.tracer def set_opencensus_tracer(tracer): - """Add the tracer to thread local.""" - setattr(_thread_local, 'tracer', tracer) + """Add the tracer to runtime context.""" + RuntimeContext.tracer = tracer def set_opencensus_attr(attr_key, attr_value): - # If there is no attrs, initialize it to empty dict. - attrs = getattr(_thread_local, 'attrs', {}) - + attrs = RuntimeContext.attrs.copy() attrs[attr_key] = attr_value - - setattr(_thread_local, 'attrs', attrs) + RuntimeContext.attrs = attrs def set_opencensus_attrs(attrs): - setattr(_thread_local, 'attrs', attrs) + RuntimeContext.attrs = attrs def get_opencensus_attr(attr_key): - attrs = getattr(_thread_local, 'attrs', None) - - if attrs is not None: - return attrs.get(attr_key) - - return None + return RuntimeContext.attrs.get(attr_key) def get_opencensus_attrs(): - return getattr(_thread_local, 'attrs', None) + return RuntimeContext.attrs def get_current_span(): - return getattr(_thread_local, 'current_span', None) + return RuntimeContext.current_span def set_current_span(current_span): - setattr(_thread_local, 'current_span', current_span) + RuntimeContext.current_span = current_span def get_opencensus_full_context(): - _tracer = get_opencensus_tracer() - _span = get_current_span() - _attrs = get_opencensus_attrs() - return _tracer, _span, _attrs + attrs = RuntimeContext.attrs + current_span = RuntimeContext.current_span + tracer = RuntimeContext.tracer + return tracer, current_span, attrs def set_opencensus_full_context(tracer, span, attrs): set_opencensus_tracer(tracer) set_current_span(span) - if not attrs: - set_opencensus_attrs({}) - else: - set_opencensus_attrs(attrs) + set_opencensus_attrs(attrs or {}) def clean(): - setattr(_thread_local, 'attrs', {}) - if hasattr(_thread_local, 'current_span'): - delattr(_thread_local, 'current_span') - if hasattr(_thread_local, 'tracer'): - delattr(_thread_local, 'tracer') + _attrs_slot.clear() + _current_span_slot.clear() + _tracer_slot.clear() def clear(): - """Clear the thread local, used in test.""" - _thread_local.__dict__.clear() + """Clear the context, used in test.""" + clean() diff --git a/setup.py b/setup.py index fb6846585..0013e3cd3 100644 --- a/setup.py +++ b/setup.py @@ -40,6 +40,7 @@ include_package_data=True, long_description=open('README.rst').read(), install_requires=[ + 'opencensus-context == 0.2.dev0', 'google-api-core >= 1.0.0, < 2.0.0', ], extras_require={},