diff --git a/skywalking/agent/protocol/grpc.py b/skywalking/agent/protocol/grpc.py index ddee75aa..0a7da16f 100644 --- a/skywalking/agent/protocol/grpc.py +++ b/skywalking/agent/protocol/grpc.py @@ -15,6 +15,7 @@ # limitations under the License. # +import logging from skywalking.loggings import logger import traceback from queue import Queue @@ -62,7 +63,7 @@ def connected(self): return self.state == grpc.ChannelConnectivity.READY def on_error(self): - traceback.print_exc() + traceback.print_exc() if logger.isEnabledFor(logging.DEBUG) else None self.channel.unsubscribe(self._cb) self.channel.subscribe(self._cb, try_to_connect=True) diff --git a/skywalking/trace/context.py b/skywalking/trace/context.py index 61b06359..9f73bf78 100644 --- a/skywalking/trace/context.py +++ b/skywalking/trace/context.py @@ -15,7 +15,6 @@ # limitations under the License. # -import threading from typing import List from skywalking import agent, config @@ -111,12 +110,13 @@ def start(self, span: Span): self.spans.append(span) def stop(self, span: Span) -> bool: - assert span is self.spans[-1] + idx = self.spans.index(span) # span SHOULD always be at end but in async-world it MAY not be, so don't crash - span.finish(self.segment) and self.spans.pop() + if span.finish(self.segment): + del self.spans[idx] if len(self.spans) == 0: - _thread_local.context = None + _local().context = None agent.archive(self.segment) return len(self.spans) == 0 @@ -212,13 +212,38 @@ def continued(self, snapshot: 'Snapshot'): self._correlation.update(snapshot.correlation) -_thread_local = threading.local() -_thread_local.context = None +try: # attempt to use async-local instead of thread-local context + import contextvars + + __local = contextvars.ContextVar('local') + + class AsyncLocal: + pass + + def _local(): + try: + return __local.get() + + except LookupError: + local = AsyncLocal() + __local.set(local) + + return local + +except ImportError: + import threading + + __local = threading.local() + + def _local(): + return __local def get_context() -> SpanContext: - if not hasattr(_thread_local, 'context'): - _thread_local.context = None - _thread_local.context = _thread_local.context or (SpanContext() if agent.connected() else NoopContext()) + local = _local() + context = getattr(local, 'context', False) + + if not context: + context = local.context = (SpanContext() if agent.connected() else NoopContext()) - return _thread_local.context + return context