inline with dotnet implementation

Author: mladjan-gadzic

example.py

@@ -32,6 +32,11 @@
 tracer = trace.get_tracer(__name__)
 
 # Start a span (logs heartbeat and stop events)
-with tracer.start_as_current_span("partial_span_1"):
-  print("partial_span_1 is running")
-  sleep(8)
+with tracer.start_as_current_span("span 1"):
+  with tracer.start_as_current_span("span 2"):
+    print("sleeping inside span 2")
+    sleep(2)
+  print("sleeping inside span 1")
+  sleep(5)
+print("sleeping outside spans")
+sleep(5)

src/partial_span_processor/__init__.py

@@ -18,8 +18,7 @@
 import json
 import threading
 import time
-from queue import Queue
-from typing import TYPE_CHECKING, Deque
+from typing import TYPE_CHECKING
 
 from google.protobuf import json_format
 from opentelemetry._logs.severity import SeverityNumber
@@ -29,6 +28,8 @@
 from opentelemetry.sdk.trace import ReadableSpan, Span, SpanProcessor
 from opentelemetry.trace import TraceFlags
 
+from partial_span_processor.peekable_queue import PeekableQueue
+
 if TYPE_CHECKING:
   from opentelemetry import context as context_api
   from opentelemetry.sdk._logs.export import LogExporter
@@ -79,10 +80,11 @@ def __init__(
     self.resource = resource
 
     self.active_spans = {}
-    self.delayed_heartbeat_spans: Deque[tuple[int, datetime.datetime]] = Deque[
-      tuple[int, datetime.datetime]]()
+    self.delayed_heartbeat_spans: PeekableQueue[tuple[int, datetime.datetime]] = \
+      PeekableQueue()
     self.delayed_heartbeat_spans_lookup: set[int] = set()
-    self.ready_heartbeat_spans: Queue[tuple[int, datetime.datetime]] = Queue()
+    self.ready_heartbeat_spans: PeekableQueue[
+      tuple[int, datetime.datetime]] = PeekableQueue()
     self.lock = threading.Lock()
 
     self.done = False
@@ -110,18 +112,20 @@ def on_start(self, span: Span,
 
       next_heartbeat_time = datetime.datetime.now() + datetime.timedelta(
         milliseconds=self.initial_heartbeat_delay_millis)
-      self.delayed_heartbeat_spans.appendleft(
+      self.delayed_heartbeat_spans.put(
         (span.context.span_id, next_heartbeat_time))
 
   def on_end(self, span: ReadableSpan) -> None:
+    is_delayed_heartbeat_pending = False
     with self.lock:
       self.active_spans.pop(span.context.span_id)
 
       if span.context.span_id in self.delayed_heartbeat_spans_lookup:
+        is_delayed_heartbeat_pending = True
         self.delayed_heartbeat_spans_lookup.remove(span.context.span_id)
-        self.delayed_heartbeat_spans.remove(
-          (span.context.span_id, next(iter(self.delayed_heartbeat_spans))[1]))
-        return
+
+    if is_delayed_heartbeat_pending:
+      return
 
     self.export_log(span, get_stop_attributes())
 
@@ -188,28 +192,42 @@ def get_log_data(self, span: Span, attributes: dict[str, str]) -> LogData:
     )
 
   def process_delayed_heartbeat_spans(self) -> None:
-    with self.lock:
+    spans_to_be_logged = []
+    with (self.lock):
       now = datetime.datetime.now()
-      while self.delayed_heartbeat_spans and self.delayed_heartbeat_spans[-1][
-        1] <= now:
-        span_id, _ = self.delayed_heartbeat_spans.pop()
+      while True:
+        if self.delayed_heartbeat_spans.empty():
+          break
+
+        (span_id, next_heartbeat_time) = self.delayed_heartbeat_spans.peek()
+        if next_heartbeat_time > now:
+          break
+
         self.delayed_heartbeat_spans_lookup.discard(span_id)
+        self.delayed_heartbeat_spans.get()
 
         span = self.active_spans.get(span_id)
         if not span:
           continue
 
-        self.export_log(span, self.get_heartbeat_attributes())
+        spans_to_be_logged.append(span)
 
         next_heartbeat_time = now + datetime.timedelta(
           milliseconds=self.heartbeat_interval_millis)
         self.ready_heartbeat_spans.put((span_id, next_heartbeat_time))
 
+    for span in spans_to_be_logged:
+      self.export_log(span, self.get_heartbeat_attributes())
+
   def process_ready_heartbeat_spans(self) -> None:
+    spans_to_be_logged = []
     now = datetime.datetime.now()
     with self.lock:
-      while not self.ready_heartbeat_spans.empty():
-        span_id, next_heartbeat_time = self.ready_heartbeat_spans.queue[0]
+      while True:
+        if self.ready_heartbeat_spans.empty():
+          break
+
+        (span_id, next_heartbeat_time) = self.ready_heartbeat_spans.peek()
         if next_heartbeat_time > now:
           break
 
@@ -219,12 +237,15 @@ def process_ready_heartbeat_spans(self) -> None:
         if not span:
           continue
 
-        self.export_log(span, self.get_heartbeat_attributes())
+        spans_to_be_logged.append(span)
 
         next_heartbeat_time = now + datetime.timedelta(
           milliseconds=self.heartbeat_interval_millis)
         self.ready_heartbeat_spans.put((span_id, next_heartbeat_time))
 
+    for span in spans_to_be_logged:
+      self.export_log(span, self.get_heartbeat_attributes())
+
 
 def get_stop_attributes() -> dict[str, str]:
   return {

tests/partial_span_processor/test_partial_span_processor.py

@@ -105,30 +105,35 @@ def test_invalid_process_interval(self):
 
   def test_on_start(self):
     span = TestPartialSpanProcessor.create_mock_span()
-    span_id = span.get_span_context().span_id
+    expected_span_id = span.get_span_context().span_id
+    now = datetime.datetime.now()
     self.processor.on_start(span)
 
-    self.assertIn(span_id, self.processor.active_spans)
-    self.assertIn(span_id, self.processor.delayed_heartbeat_spans_lookup)
-    self.assertIn((span_id, unittest.mock.ANY),
-                  self.processor.delayed_heartbeat_spans)
+    self.assertIn(expected_span_id, self.processor.active_spans)
+    self.assertIn(expected_span_id,
+                  self.processor.delayed_heartbeat_spans_lookup)
+    self.assertEqual(self.processor.delayed_heartbeat_spans.qsize(), 1)
+    (
+      span_id,
+      next_heartbeat_time) = self.processor.delayed_heartbeat_spans.get()
+    self.assertEqual(expected_span_id, span_id)
+    self.assertGreater(next_heartbeat_time, now)
+    self.assertEqual(self.log_exporter.get_finished_logs(), ())
 
   def test_on_end_when_initial_heartbeat_not_sent(self):
     span = TestPartialSpanProcessor.create_mock_span()
     span_id = span.get_span_context().span_id
 
     self.processor.active_spans[span_id] = span
     self.processor.delayed_heartbeat_spans_lookup.add(span_id)
-    self.processor.delayed_heartbeat_spans.append(
-      (span_id, unittest.mock.ANY))
+    self.processor.delayed_heartbeat_spans.put((span_id, unittest.mock.ANY))
 
     self.processor.on_end(span)
 
     self.assertNotIn(span_id, self.processor.active_spans)
     self.assertNotIn(span_id,
                      self.processor.delayed_heartbeat_spans_lookup)
-    self.assertNotIn((span_id, unittest.mock.ANY),
-                     self.processor.delayed_heartbeat_spans)
+    self.assertFalse(self.processor.delayed_heartbeat_spans.empty())
     self.assertEqual(self.log_exporter.get_finished_logs(), ())
 
   def test_on_end_when_initial_heartbeat_sent(self):
@@ -155,15 +160,15 @@ def test_process_delayed_heartbeat_spans(self):
 
     self.processor.active_spans[span_id] = span
     now = datetime.datetime.now()
-    self.processor.delayed_heartbeat_spans.append((span_id, now))
+    self.processor.delayed_heartbeat_spans.put((span_id, now))
     self.processor.delayed_heartbeat_spans_lookup.add(span_id)
 
     with patch("datetime.datetime") as mock_datetime:
       mock_datetime.now.return_value = now
       self.processor.process_delayed_heartbeat_spans()
 
     self.assertNotIn(span_id, self.processor.delayed_heartbeat_spans_lookup)
-    self.assertNotIn((span_id, now), self.processor.delayed_heartbeat_spans)
+    self.assertTrue(self.processor.delayed_heartbeat_spans.empty())
 
     next_heartbeat_time = now + datetime.timedelta(
       milliseconds=self.processor.heartbeat_interval_millis)