Fix local span metrics observers registering child observers (#592)

This fixes an issue where the local span metrics observers (tagged and
untagged) would not register child observers when nesting spans beneath
local spans.

Author: spladug

baseplate/observers/metrics.py

@@ -124,6 +124,14 @@ def on_start(self) -> None:
     def on_incr_tag(self, key: str, delta: float) -> None:
         self.batch.counter(key).increment(delta, sample_rate=self.sample_rate)
 
+    def on_child_span_created(self, span: Span) -> None:
+        observer: SpanObserver
+        if isinstance(span, LocalSpan):
+            observer = MetricsLocalSpanObserver(self.batch, span, self.sample_rate)
+        else:
+            observer = MetricsClientSpanObserver(self.batch, span, self.sample_rate)
+        span.register(observer)
+
     def on_finish(self, exc_info: Optional[_ExcInfo]) -> None:
         self.timer.stop()
 

baseplate/observers/metrics_tagged.py

@@ -160,6 +160,18 @@ def on_incr_tag(self, key: str, delta: float) -> None:
     def on_set_tag(self, key: str, value: Any) -> None:
         self.tags[key] = value
 
+    def on_child_span_created(self, span: Span) -> None:
+        observer: SpanObserver
+        if isinstance(span, LocalSpan):
+            observer = TaggedMetricsLocalSpanObserver(
+                self.batch, span, self.allowlist, self.sample_rate
+            )
+        else:
+            observer = TaggedMetricsClientSpanObserver(
+                self.batch, span, self.allowlist, self.sample_rate
+            )
+        span.register(observer)
+
     def on_finish(self, exc_info: Optional[_ExcInfo]) -> None:
         filtered_tags = {k: v for (k, v) in self.tags.items() if k in self.allowlist}
 

tests/unit/observers/metrics_tagged_tests.py

@@ -9,6 +9,7 @@
 import pytest
 
 from baseplate import RequestContext
+from baseplate import ServerSpan
 from baseplate import Span
 from baseplate.lib.metrics import Counter
 from baseplate.lib.metrics import Gauge
@@ -156,3 +157,21 @@ def test_observer(observer_cls, name):
             assert counter["tags"].pop("endpoint") == "span"
         assert counter["tags"].pop("tag2") == "bar"
         assert counter["tags"] == {}
+
+
+def test_nested():
+    batch = FakeBatch()
+    span = ServerSpan(
+        trace_id=1234,
+        parent_id=2345,
+        span_id=3456,
+        sampled=None,
+        flags=None,
+        name="fancy.span",
+        context=RequestContext({}),
+    )
+    observer = TaggedMetricsLocalSpanObserver(batch, span, {"client", "endpoint"})
+    span.register(observer)
+
+    with span.make_child("foo", local=True, component_name="foo") as child_span:
+        assert len(child_span.observers) == 1

tests/unit/observers/metrics_tests.py

@@ -120,3 +120,19 @@ def test_metrics(self):
 
         observer.on_finish(exc_info=None)
         self.assertEqual(mock_timer.stop.call_count, 1)
+
+    def test_spans_under_local_spans(self):
+        mock_batch = mock.Mock(spec=Batch)
+
+        mock_local_span = mock.Mock(spec=LocalSpan)
+        mock_local_span.name = "example"
+        mock_local_span.component_name = "some_component"
+
+        mock_nested_span = mock.Mock(spec=LocalSpan)
+        mock_nested_span.name = "nested"
+        mock_nested_span.component_name = "some_component2"
+
+        observer = MetricsLocalSpanObserver(mock_batch, mock_local_span)
+        observer.on_child_span_created(mock_nested_span)
+
+        self.assertEqual(mock_nested_span.register.call_count, 1)