Make force_flush available on SDK's tracer provider (#594)

Co-authored-by: Yusuke Tsutsumi <[email protected]>

Author: mariojonke

opentelemetry-sdk/src/opentelemetry/sdk/trace/__init__.py

@@ -15,14 +15,25 @@
 
 import abc
 import atexit
+import concurrent.futures
 import json
 import logging
 import random
 import threading
 from collections import OrderedDict
 from contextlib import contextmanager
 from types import TracebackType
-from typing import Iterator, MutableSequence, Optional, Sequence, Tuple, Type
+from typing import (
+    Any,
+    Callable,
+    Iterator,
+    MutableSequence,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    Union,
+)
 
 from opentelemetry import context as context_api
 from opentelemetry import trace as trace_api
@@ -89,9 +100,12 @@ def force_flush(self, timeout_millis: int = 30000) -> bool:
         """
 
 
-class MultiSpanProcessor(SpanProcessor):
-    """Implementation of :class:`SpanProcessor` that forwards all received
-    events to a list of `SpanProcessor`.
+class SynchronousMultiSpanProcessor(SpanProcessor):
+    """Implementation of class:`SpanProcessor` that forwards all received
+    events to a list of span processors sequentially.
+
+    The underlying span processors are called in sequential order as they were
+    added.
     """
 
     def __init__(self):
@@ -114,9 +128,113 @@ def on_end(self, span: "Span") -> None:
             sp.on_end(span)
 
     def shutdown(self) -> None:
+        """Sequentially shuts down all underlying span processors.
+        """
         for sp in self._span_processors:
             sp.shutdown()
 
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Sequentially calls force_flush on all underlying
+        :class:`SpanProcessor`
+
+        Args:
+            timeout_millis: The maximum amount of time over all span processors
+                to wait for spans to be exported. In case the first n span
+                processors exceeded the timeout followup span processors will be
+                skipped.
+
+        Returns:
+            True if all span processors flushed their spans within the
+            given timeout, False otherwise.
+        """
+        deadline_ns = time_ns() + timeout_millis * 1000000
+        for sp in self._span_processors:
+            current_time_ns = time_ns()
+            if current_time_ns >= deadline_ns:
+                return False
+
+            if not sp.force_flush((deadline_ns - current_time_ns) // 1000000):
+                return False
+
+        return True
+
+
+class ConcurrentMultiSpanProcessor(SpanProcessor):
+    """Implementation of :class:`SpanProcessor` that forwards all received
+    events to a list of span processors in parallel.
+
+    Calls to the underlying span processors are forwarded in parallel by
+    submitting them to a thread pool executor and waiting until each span
+    processor finished its work.
+
+    Args:
+        num_threads: The number of threads managed by the thread pool executor
+            and thus defining how many span processors can work in parallel.
+    """
+
+    def __init__(self, num_threads: int = 2):
+        # use a tuple to avoid race conditions when adding a new span and
+        # iterating through it on "on_start" and "on_end".
+        self._span_processors = ()  # type: Tuple[SpanProcessor, ...]
+        self._lock = threading.Lock()
+        self._executor = concurrent.futures.ThreadPoolExecutor(
+            max_workers=num_threads
+        )
+
+    def add_span_processor(self, span_processor: SpanProcessor) -> None:
+        """Adds a SpanProcessor to the list handled by this instance."""
+        with self._lock:
+            self._span_processors = self._span_processors + (span_processor,)
+
+    def _submit_and_await(
+        self, func: Callable[[SpanProcessor], Callable[..., None]], *args: Any
+    ):
+        futures = []
+        for sp in self._span_processors:
+            future = self._executor.submit(func(sp), *args)
+            futures.append(future)
+        for future in futures:
+            future.result()
+
+    def on_start(self, span: "Span") -> None:
+        self._submit_and_await(lambda sp: sp.on_start, span)
+
+    def on_end(self, span: "Span") -> None:
+        self._submit_and_await(lambda sp: sp.on_end, span)
+
+    def shutdown(self) -> None:
+        """Shuts down all underlying span processors in parallel."""
+        self._submit_and_await(lambda sp: sp.shutdown)
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Calls force_flush on all underlying span processors in parallel.
+
+        Args:
+            timeout_millis: The maximum amount of time to wait for spans to be
+                exported.
+
+        Returns:
+            True if all span processors flushed their spans within the given
+            timeout, False otherwise.
+        """
+        futures = []
+        for sp in self._span_processors:  # type: SpanProcessor
+            future = self._executor.submit(sp.force_flush, timeout_millis)
+            futures.append(future)
+
+        timeout_sec = timeout_millis / 1e3
+        done_futures, not_done_futures = concurrent.futures.wait(
+            futures, timeout_sec
+        )
+        if not_done_futures:
+            return False
+
+        for future in done_futures:
+            if not future.result():
+                return False
+
+        return True
+
 
 class EventBase(abc.ABC):
     def __init__(self, name: str, timestamp: Optional[int] = None) -> None:
@@ -730,8 +848,13 @@ def __init__(
         sampler: sampling.Sampler = trace_api.sampling.ALWAYS_ON,
         resource: Resource = Resource.create_empty(),
         shutdown_on_exit: bool = True,
+        active_span_processor: Union[
+            SynchronousMultiSpanProcessor, ConcurrentMultiSpanProcessor
+        ] = None,
     ):
-        self._active_span_processor = MultiSpanProcessor()
+        self._active_span_processor = (
+            active_span_processor or SynchronousMultiSpanProcessor()
+        )
         self.resource = resource
         self.sampler = sampler
         self._atexit_handler = None
@@ -759,8 +882,8 @@ def add_span_processor(self, span_processor: SpanProcessor) -> None:
         The span processors are invoked in the same order they are registered.
         """
 
-        # no lock here because MultiSpanProcessor.add_span_processor is
-        # thread safe
+        # no lock here because add_span_processor is thread safe for both
+        # SynchronousMultiSpanProcessor and ConcurrentMultiSpanProcessor.
         self._active_span_processor.add_span_processor(span_processor)
 
     def shutdown(self):
@@ -769,3 +892,23 @@ def shutdown(self):
         if self._atexit_handler is not None:
             atexit.unregister(self._atexit_handler)
             self._atexit_handler = None
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        """Requests the active span processor to process all spans that have not
+        yet been processed.
+
+        By default force flush is called sequentially on all added span
+        processors. This means that span processors further back in the list
+        have less time to flush their spans.
+        To have span processors flush their spans in parallel it is possible to
+        initialize the tracer provider with an instance of
+        `ConcurrentMultiSpanProcessor` at the cost of using multiple threads.
+
+        Args:
+            timeout_millis: The maximum amount of time to wait for spans to be
+                processed.
+
+        Returns:
+            False if the timeout is exceeded, True otherwise.
+        """
+        return self._active_span_processor.force_flush(timeout_millis)