feat(metrics): add support for AFE server timings

Author: sinhasubham

packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_interceptor.py

@@ -16,6 +16,7 @@
 
 import inspect
 import logging
+import os
 import re
 from typing import Any, Dict
 
@@ -133,6 +134,12 @@ def intercept(self, invoked_method, request_or_iterator, call_details):
 
         tracer.set_method(method_name)
         tracer.record_attempt_start()
+
+        if os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() != "true":
+            metadata = list(call_details.metadata or [])
+            metadata.append(("x-goog-spanner-enable-afe-server-timing", "true"))
+            call_details = call_details._replace(metadata=metadata)
+
         response = invoked_method(request_or_iterator, call_details)
 
         return _wrap_response(response, tracer)
@@ -215,8 +222,13 @@ async def _async_intercept(
 
         tracer.set_method(method_name)
         tracer.record_attempt_start()
-        response = await continuation(call_details, request_or_iterator)
 
+        if os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() != "true":
+            metadata = list(call_details.metadata or [])
+            metadata.append(("x-goog-spanner-enable-afe-server-timing", "true"))
+            call_details = call_details._replace(metadata=metadata)
+
+        response = await continuation(call_details, request_or_iterator)
         if hasattr(response, "__anext__"):
             return _AsyncStreamingResponseWrapper(response, tracer)
         else:

packages/google-cloud-spanner/google/cloud/spanner_v1/metrics/metrics_tracer.py

@@ -19,6 +19,7 @@
 while the helper classes provide additional functionality and context for the metrics being traced.
 """
 
+import os
 import re
 from datetime import datetime
 from typing import Any, Dict, Optional
@@ -425,7 +426,7 @@ def record_gfe_latency(self, latency: int) -> None:
         ):
             return
         self._instrument_gfe_latency.record(
-            amount=latency, attributes=self.client_attributes
+            amount=latency, attributes=self._create_attempt_otel_attributes()
         )
 
     def record_gfe_connectivity_error_count(self) -> None:
@@ -439,7 +440,7 @@ def record_gfe_connectivity_error_count(self) -> None:
         ):
             return
         self._instrument_gfe_connectivity_error_count.add(
-            amount=1, attributes=self.client_attributes
+            amount=1, attributes=self._create_attempt_otel_attributes()
         )
 
     def record_afe_latency(self, latency: int) -> None:
@@ -453,10 +454,11 @@ def record_afe_latency(self, latency: int) -> None:
             not self.enabled
             or not HAS_OPENTELEMETRY_INSTALLED
             or not getattr(self, "_instrument_afe_latency", None)
+            or os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() == "true"
         ):
             return
         self._instrument_afe_latency.record(
-            amount=latency, attributes=self.client_attributes
+            amount=latency, attributes=self._create_attempt_otel_attributes()
         )
 
     def record_afe_connectivity_error_count(self) -> None:
@@ -467,10 +469,11 @@ def record_afe_connectivity_error_count(self) -> None:
             not self.enabled
             or not HAS_OPENTELEMETRY_INSTALLED
             or not getattr(self, "_instrument_afe_connectivity_error_count", None)
+            or os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() == "true"
         ):
             return
         self._instrument_afe_connectivity_error_count.add(
-            amount=1, attributes=self.client_attributes
+            amount=1, attributes=self._create_attempt_otel_attributes()
         )
 
     @staticmethod

packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py

@@ -202,16 +202,20 @@ def test_gfe_missing_header_count_exported(self):
             }
 
             self.assertIn(
-                "gfe_missing_header_count", metrics, f"Metrics: {list(metrics.keys())}"
+                "gfe_connectivity_error_count",
+                metrics,
+                f"Metrics: {list(metrics.keys())}",
             )
-            missing_metric = metrics["gfe_missing_header_count"]
+            missing_metric = metrics["gfe_connectivity_error_count"]
             point = next(iter(missing_metric.data.data_points))
             self.assertGreaterEqual(point.value, 1)
 
             self.assertIn(
-                "afe_missing_header_count", metrics, f"Metrics: {list(metrics.keys())}"
+                "afe_connectivity_error_count",
+                metrics,
+                f"Metrics: {list(metrics.keys())}",
             )
-            afe_missing_metric = metrics["afe_missing_header_count"]
+            afe_missing_metric = metrics["afe_connectivity_error_count"]
             afe_point = next(iter(afe_missing_metric.data.data_points))
             self.assertGreaterEqual(afe_point.value, 1)
         finally:

packages/google-cloud-spanner/tests/unit/test_metrics_interceptor.py

@@ -113,9 +113,10 @@ def test_intercept_with_tracer(interceptor, mock_tracer_ctx):
         ],
     )
 
+    replaced_call_details = call_details._replace.return_value
     response = interceptor.intercept(mock_invoked_method, "request", call_details)
     assert response == invoked_response
     mock_tracer_ctx.record_attempt_start.assert_called()
     mock_tracer_ctx.record_attempt_completion.assert_called_once()
     mock_tracer_ctx.record_front_end_metrics.assert_called_once()
-    mock_invoked_method.assert_called_once_with("request", call_details)
+    mock_invoked_method.assert_called_once_with("request", replaced_call_details)

packages/google-cloud-spanner/tests/unit/test_metrics_tracer.py

@@ -243,7 +243,7 @@ def test_record_gfe_latency(metrics_tracer):
     assert mock_gfe_latency.record.call_args[1]["amount"] == 100
     assert (
         mock_gfe_latency.record.call_args[1]["attributes"]
-        == metrics_tracer.client_attributes
+        == metrics_tracer._create_attempt_otel_attributes()
     )
 
     # Test when tracing is disabled
@@ -266,7 +266,7 @@ def test_record_gfe_connectivity_error_count(metrics_tracer):
     assert mock_gfe_connectivity_error_count.add.call_args[1]["amount"] == 1
     assert (
         mock_gfe_connectivity_error_count.add.call_args[1]["attributes"]
-        == metrics_tracer.client_attributes
+        == metrics_tracer._create_attempt_otel_attributes()
     )
 
     # Test when tracing is disabled
@@ -336,9 +336,13 @@ def test_record_afe_latency(metrics_tracer):
     assert mock_afe_latency.record.call_args[1]["amount"] == 100
     assert (
         mock_afe_latency.record.call_args[1]["attributes"]
-        == metrics_tracer.client_attributes
+        == metrics_tracer._create_attempt_otel_attributes()
     )
 
+    with mock.patch.dict("os.environ", {"SPANNER_DISABLE_AFE_SERVER_TIMING": "true"}):
+        metrics_tracer.record_afe_latency(300)
+    assert mock_afe_latency.record.call_count == 1
+
     metrics_tracer.enabled = False
     metrics_tracer.record_afe_latency(200)
     assert mock_afe_latency.record.call_count == 1
@@ -355,9 +359,13 @@ def test_record_afe_connectivity_error_count(metrics_tracer):
     assert mock_afe_missing.add.call_args[1]["amount"] == 1
     assert (
         mock_afe_missing.add.call_args[1]["attributes"]
-        == metrics_tracer.client_attributes
+        == metrics_tracer._create_attempt_otel_attributes()
     )
 
+    with mock.patch.dict("os.environ", {"SPANNER_DISABLE_AFE_SERVER_TIMING": "true"}):
+        metrics_tracer.record_afe_connectivity_error_count()
+    assert mock_afe_missing.add.call_count == 1
+
     metrics_tracer.enabled = False
     metrics_tracer.record_afe_connectivity_error_count()
     assert mock_afe_missing.add.call_count == 1