fix(fastapi): wrapper sync mode

- fastapi: Added wrapper sync mode
- Removed aiohttp support
- tornado: Fixed trace cleaning

Author: maorlx

README.md

@@ -195,7 +195,6 @@ The following frameworks are supported by Epsagon:
 |[Django](#django)                       |`>=1.11`                   |<ul><li>- [x] </li></ul>                             |
 |[Flask](#flask)                         |`>=0.5`                    |<ul><li>- [x] </li></ul>                             |
 |[Tornado](#tornado)                     |`>=4.0`                    |<ul><li>- [x] </li></ul>                             |
-|[aiohttp](#aiohttp)                     |`>=3.0.0`                  |<ul><li>- [x] </li></ul>                             |
 |[fastapi](#fastapi)                     |`>=0.62.0`                 |<ul><li>- [x] </li></ul>                             |
 |[Celery](#celery)                       |`>=4.0.0`                  |<ul><li>- [x] </li></ul>                             |
 |[Azure Functions](#azure-functions)     |`>=2.0.0`                  |<ul><li>- [ ] </li></ul>                             |
@@ -295,22 +294,6 @@ epsagon.init(
 )
 ```
 
-### aiohttp
-
-Tracing aiohttp application can be done in two methods:
-1. [Auto-tracing](#auto-tracing) using the environment variable.
-2. Calling the SDK.
-
-Calling the SDK is simple, and should be done in your main `py` file where the application is being initialized:
-```python
-import epsagon
-epsagon.init(
-    token='<epsagon-token>',
-    app_name='<app-name-stage>',
-    metadata_only=False,
-)
-```
-
 ### fastapi
 
 Tracing fastapi application can be done in two methods:

epsagon/modules/__init__.py

@@ -11,7 +11,6 @@
 IGNORE_MODULES = ('__init__',)
 PYTHON_EXTENSIONS = ('.py', '.pyc')
 VERSION_DEPENDENCIES = {
-    'aiohttp': (3, 5, 3),
     'fastapi': (3, 5, 3),
 }
 

epsagon/modules/aiohttp.py

@@ -94,6 +94,8 @@ def after_request(cls, wrapped, instance, args, kwargs):
             )
 
         res = wrapped(*args, **kwargs)
+        trace = None
+        is_trace_sent = False
         try:
             unique_id = getattr(instance, TORNADO_TRACE_ID, None)
             if not unique_id:
@@ -122,11 +124,11 @@ def after_request(cls, wrapped, instance, args, kwargs):
                 if not ignored:
                     tornado_runner.update_response(instance, response_body)
                     epsagon.trace.trace_factory.send_traces(trace)
-        except Exception as instrumentation_exception:  # pylint: disable=W0703
-            epsagon.trace.trace_factory.add_exception(
-                instrumentation_exception,
-                traceback.format_exc()
-            )
+                    is_trace_sent = True
+        except Exception:  # pylint: disable=W0703
+            if not is_trace_sent and trace:
+                epsagon.trace.trace_factory.pop_trace(trace=trace)
+
         return res
 
     @classmethod
@@ -205,9 +207,9 @@ def thread_pool_submit(cls, func, _, args, kwargs):
         unique_id = None
 
         try:
-            unique_id = (
-                epsagon.trace.trace_factory.get_or_create_trace().unique_id
-            )
+            trace = epsagon.trace.trace_factory.get_trace()
+            if trace:
+                unique_id = trace.unique_id
 
         except Exception as instrumentation_exception:  # pylint: disable=W0703
             epsagon.trace.trace_factory.add_exception(

epsagon/modules/tornado.py

@@ -22,7 +22,6 @@
 from epsagon.trace_transports import NoneTransport, HTTPTransport, LogTransport
 from .constants import (
     TIMEOUT_GRACE_TIME_MS,
-    EPSAGON_MARKER,
     MAX_LABEL_SIZE,
     DEFAULT_SAMPLE_RATE,
     TRACE_URL_PREFIX,
@@ -96,7 +95,6 @@ def __init__(self):
         self.keys_to_ignore = None
         self.keys_to_allow = None
         self.use_single_trace = True
-        self.use_async_tracer = False
         self.singleton_trace = None
         self.local_thread_to_unique_id = {}
         self.transport = NoneTransport()
@@ -209,13 +207,6 @@ def switch_to_multiple_traces(self):
         """
         self.use_single_trace = False
 
-    def switch_to_async_tracer(self):
-        """
-        Set the use_async_tracer flag to True.
-        :return: None
-        """
-        self.use_async_tracer = True
-
     def _create_new_trace(self, unique_id=None):
         """
         Creating new trace instance
@@ -264,66 +255,18 @@ def _get_thread_trace(self, should_create=False):
             self.traces[thread_id] = new_trace
         return self.traces[thread_id]
 
-    @staticmethod
-    def _get_current_task():
-        """
-        Gets the current asyncio task safely
-        :return: The task.
-        """
-        # Dynamic import since this is only valid in Python3+
-        asyncio = __import__('asyncio')
-        try:
-            return asyncio.Task.current_task()
-        except RuntimeError:
-            return None
-
-    def _get_tracer_async_mode(self, should_create):
-        """
-        Get trace assuming async tracer.
-        :return: The trace.
-        """
-        task = type(self)._get_current_task()
-        if not task:
-            return None
-
-        trace = getattr(task, EPSAGON_MARKER, None)
-        if not trace and should_create:
-            trace = self._create_new_trace()
-            setattr(task, EPSAGON_MARKER, trace)
-        return trace
-
-    def _pop_trace_async_mode(self):
-        """
-        Pops the trace from the current task, assuming async tracer
-        :return: The trace.
-        """
-        task = type(self)._get_current_task()
-        if not task:
-            return None
-
-        trace = getattr(task, EPSAGON_MARKER, None)
-        if trace: # can safely remove tracer from async task
-            delattr(task, EPSAGON_MARKER)
-        return trace
-
     def _get_trace(self, unique_id=None, should_create=False):
         """
         Get trace based on the tracing mode:
         single trace (use_single_trace=True) or
-        multiple threads (use_single_trace=False) or
-        async tasks (use_async_tracer=True)
+        multiple threads (use_single_trace=False)
 
         If should_create then creating a new trace if trace does not exist
         if use_single_trace is set to False, each thread will have
         it's own trace..
-        if use_async_tracer is set to True, each asyncio.Task will be
-        assigned with the trace.
         :return: The trace.
         """
         with TraceFactory.LOCK:
-            if self.use_async_tracer:
-                return self._get_tracer_async_mode(should_create=should_create)
-
             unique_id = self.get_thread_local_unique_id(unique_id)
             if unique_id:
                 trace = (
@@ -378,11 +321,6 @@ def pop_trace(self, trace=None):
         :return: unique id
         """
         with self.LOCK:
-            if self.use_async_tracer:
-                trace = self._pop_trace_async_mode()
-                if trace:
-                    # async tracer found
-                    return trace
             if self.traces:
                 trace = self.traces.pop(self.get_trace_identifier(trace), None)
                 if not self.traces: