feat(metrics): add AFE latency metrics and simplify metadata extraction

Author: sinhasubham

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

@@ -60,11 +60,17 @@
 METRIC_NAME_ATTEMPT_COUNT = "attempt_count"
 METRIC_NAME_GFE_LATENCY = "gfe_latency"
 METRIC_NAME_GFE_MISSING_HEADER_COUNT = "gfe_missing_header_count"
+METRIC_NAME_AFE_LATENCY = "afe_latency"
+METRIC_NAME_AFE_MISSING_HEADER_COUNT = "afe_missing_header_count"
 METRIC_NAMES = [
     METRIC_NAME_OPERATION_LATENCIES,
     METRIC_NAME_ATTEMPT_LATENCIES,
     METRIC_NAME_OPERATION_COUNT,
     METRIC_NAME_ATTEMPT_COUNT,
+    METRIC_NAME_GFE_LATENCY,
+    METRIC_NAME_GFE_MISSING_HEADER_COUNT,
+    METRIC_NAME_AFE_LATENCY,
+    METRIC_NAME_AFE_MISSING_HEADER_COUNT,
 ]
 
 METRIC_EXPORT_INTERVAL_MS = 60000  # 1 Minute

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

@@ -147,12 +147,8 @@ def _wrap_response(response: Any, tracer: Any) -> Any:
                     metadata.extend(response.initial_metadata() or [])
                 except Exception as e:
                     logger.warning(f"Failed to retrieve initial metadata: {e}")
-            if hasattr(response, "trailing_metadata"):
-                try:
-                    metadata.extend(response.trailing_metadata() or [])
-                except Exception as e:
-                    logger.warning(f"Failed to retrieve trailing metadata: {e}")
             tracer.record_gfe_metrics(metadata)
+            tracer.record_afe_metrics(metadata)
         except Exception as e:
             logger.warning(f"Failed to record metrics: {e}")
         return response
@@ -260,12 +256,8 @@ def _record_metrics(self):
                     metadata.extend(self._response.initial_metadata() or [])
                 except Exception as e:
                     logger.warning(f"Failed to retrieve initial metadata: {e}")
-            if hasattr(self._response, "trailing_metadata"):
-                try:
-                    metadata.extend(self._response.trailing_metadata() or [])
-                except Exception as e:
-                    logger.warning(f"Failed to retrieve trailing metadata: {e}")
             self._tracer.record_gfe_metrics(metadata)
+            self._tracer.record_afe_metrics(metadata)
         except Exception as e:
             logger.warning(f"Failed to record metrics: {e}")
 
@@ -341,15 +333,8 @@ async def _record_metrics(self):
                     metadata.extend(res or [])
                 except Exception as e:
                     logger.warning(f"Failed to retrieve initial metadata: {e}")
-            if hasattr(self._response, "trailing_metadata"):
-                try:
-                    res = self._response.trailing_metadata()
-                    if inspect.isawaitable(res):
-                        res = await res
-                    metadata.extend(res or [])
-                except Exception as e:
-                    logger.warning(f"Failed to retrieve trailing metadata: {e}")
             self._tracer.record_gfe_metrics(metadata)
+            self._tracer.record_afe_metrics(metadata)
         except Exception as e:
             logger.warning(f"Failed to record metrics: {e}")
 
@@ -454,15 +439,8 @@ async def _record_metrics(self):
                     metadata.extend(res or [])
                 except Exception as e:
                     logger.warning(f"Failed to retrieve initial metadata: {e}")
-            if hasattr(self._response, "trailing_metadata"):
-                try:
-                    res = self._response.trailing_metadata()
-                    if inspect.isawaitable(res):
-                        res = await res
-                    metadata.extend(res or [])
-                except Exception as e:
-                    logger.warning(f"Failed to retrieve trailing metadata: {e}")
             self._tracer.record_gfe_metrics(metadata)
+            self._tracer.record_afe_metrics(metadata)
         except Exception as e:
             logger.warning(f"Failed to record metrics: {e}")
 

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

@@ -185,6 +185,8 @@ class should not have any knowledge about the observability framework used for m
     _instrument_operation_latency: "Histogram"
     _instrument_gfe_latency: "Histogram"
     _instrument_gfe_missing_header_count: "Counter"
+    _instrument_afe_latency: "Histogram"
+    _instrument_afe_missing_header_count: "Counter"
     current_op: MetricOpTracer
     enabled: bool
     gfe_enabled: bool
@@ -201,6 +203,8 @@ def __init__(
         gfe_enabled: bool = False,
         instrument_gfe_latency: Optional["Histogram"] = None,
         instrument_gfe_missing_header_count: Optional["Counter"] = None,
+        instrument_afe_latency: Optional["Histogram"] = None,
+        instrument_afe_missing_header_count: Optional["Counter"] = None,
     ):
         """
         Initialize a MetricsTracer instance with the given parameters.
@@ -219,6 +223,8 @@ def __init__(
             gfe_enabled (bool, optional): Indicates if GFE metrics are enabled. Defaults to False.
             instrument_gfe_latency (Histogram, optional): Instrument for measuring GFE latency.
             instrument_gfe_missing_header_count (Counter, optional): Instrument for counting missing GFE headers.
+            instrument_afe_latency (Histogram, optional): Instrument for measuring AFE latency.
+            instrument_afe_missing_header_count (Counter, optional): Instrument for counting missing AFE headers.
         """
         self.current_op = MetricOpTracer()
         self._client_attributes = client_attributes
@@ -228,6 +234,8 @@ def __init__(
         self._instrument_operation_counter = instrument_operation_counter
         self._instrument_gfe_latency = instrument_gfe_latency
         self._instrument_gfe_missing_header_count = instrument_gfe_missing_header_count
+        self._instrument_afe_latency = instrument_afe_latency
+        self._instrument_afe_missing_header_count = instrument_afe_missing_header_count
         self.enabled = enabled
         self.gfe_enabled = True
 
@@ -430,6 +438,37 @@ def record_gfe_missing_header_count(self) -> None:
             amount=1, attributes=self.client_attributes
         )
 
+    def record_afe_latency(self, latency: int) -> None:
+        """
+        Records the AFE latency using the Histogram instrument.
+
+        Args:
+            latency (int): The latency duration to be recorded.
+        """
+        if (
+            not self.enabled
+            or not HAS_OPENTELEMETRY_INSTALLED
+            or not getattr(self, "_instrument_afe_latency", None)
+        ):
+            return
+        self._instrument_afe_latency.record(
+            amount=latency, attributes=self.client_attributes
+        )
+
+    def record_afe_missing_header_count(self) -> None:
+        """
+        Increments the counter for missing AFE headers.
+        """
+        if (
+            not self.enabled
+            or not HAS_OPENTELEMETRY_INSTALLED
+            or not getattr(self, "_instrument_afe_missing_header_count", None)
+        ):
+            return
+        self._instrument_afe_missing_header_count.add(
+            amount=1, attributes=self.client_attributes
+        )
+
     @staticmethod
     def extract_gfe_latency(metadata: Any) -> Optional[int]:
         """
@@ -486,6 +525,62 @@ def record_gfe_metrics(self, metadata: Any) -> None:
         else:
             self.record_gfe_missing_header_count()
 
+    @staticmethod
+    def extract_afe_latency(metadata: Any) -> Optional[int]:
+        """
+        Extracts the AFE latency value (in milliseconds) from response metadata.
+        """
+        if not metadata:
+            return None
+
+        header_vals = []
+        if isinstance(metadata, dict):
+            for key, val in metadata.items():
+                if key and str(key).lower() in ("server-timing", "server_timing"):
+                    if isinstance(val, (list, tuple)):
+                        header_vals.extend(val)
+                    else:
+                        header_vals.append(val)
+        elif isinstance(metadata, (list, tuple)):
+            for item in metadata:
+                if isinstance(item, (list, tuple)) and len(item) == 2:
+                    key, val = item
+                    if key and str(key).lower() in ("server-timing", "server_timing"):
+                        if isinstance(val, (list, tuple)):
+                            header_vals.extend(val)
+                        else:
+                            header_vals.append(val)
+
+        for header_val in header_vals:
+            if not header_val:
+                continue
+            if isinstance(header_val, bytes):
+                try:
+                    header_val = header_val.decode("utf-8")
+                except Exception:
+                    header_val = str(header_val)
+            elif not isinstance(header_val, str):
+                header_val = str(header_val)
+            match = re.search(r"afe(?:t4t7)?;\s*dur=([0-9.]+)", header_val)
+            if match:
+                try:
+                    return int(float(match.group(1)))
+                except ValueError:
+                    pass
+        return None
+
+    def record_afe_metrics(self, metadata: Any) -> None:
+        """
+        Extracts and records AFE metrics from the RPC response metadata.
+        """
+        if not self.enabled or not HAS_OPENTELEMETRY_INSTALLED:
+            return
+        latency = self.extract_afe_latency(metadata)
+        if latency is not None:
+            self.record_afe_latency(latency)
+        else:
+            self.record_afe_missing_header_count()
+
     def _create_operation_otel_attributes(self) -> dict:
         """
         Create additional attributes for operation metrics tracing.

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

@@ -26,6 +26,8 @@
     METRIC_NAME_ATTEMPT_LATENCIES,
     METRIC_NAME_GFE_LATENCY,
     METRIC_NAME_GFE_MISSING_HEADER_COUNT,
+    METRIC_NAME_AFE_LATENCY,
+    METRIC_NAME_AFE_MISSING_HEADER_COUNT,
     METRIC_NAME_OPERATION_COUNT,
     METRIC_NAME_OPERATION_LATENCIES,
     MONITORED_RES_LABEL_KEY_CLIENT_HASH,
@@ -57,6 +59,8 @@ class MetricsTracerFactory:
     _instrument_operation_counter: "Counter"
     _instrument_gfe_latency: "Histogram"
     _instrument_gfe_missing_header_count: "Counter"
+    _instrument_afe_latency: "Histogram"
+    _instrument_afe_missing_header_count: "Counter"
     _client_attributes: Dict[str, str]
 
     @property
@@ -274,6 +278,10 @@ def create_metrics_tracer(self) -> MetricsTracer:
             instrument_gfe_missing_header_count=getattr(
                 self, "_instrument_gfe_missing_header_count", None
             ),
+            instrument_afe_latency=getattr(self, "_instrument_afe_latency", None),
+            instrument_afe_missing_header_count=getattr(
+                self, "_instrument_afe_missing_header_count", None
+            ),
         )
         return metrics_tracer
 
@@ -331,3 +339,15 @@ def _create_metric_instruments(self, service_name: str) -> None:
             unit="1",
             description="GFE missing header count.",
         )
+
+        self._instrument_afe_latency = meter.create_histogram(
+            name=METRIC_NAME_AFE_LATENCY,
+            unit="ms",
+            description="AFE Latency.",
+        )
+
+        self._instrument_afe_missing_header_count = meter.create_counter(
+            name=METRIC_NAME_AFE_MISSING_HEADER_COUNT,
+            unit="1",
+            description="AFE missing header count.",
+        )

…sts/mockserver_tests/test_gfe_metrics.py …ockserver_tests/test_frontend_metrics.pypackages/google-cloud-spanner/tests/mockserver_tests/test_gfe_metrics.py renamed to packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py packages/google-cloud-spanner/tests/mockserver_tests/test_gfe_metrics.py renamed to packages/google-cloud-spanner/tests/mockserver_tests/test_frontend_metrics.py

@@ -34,7 +34,7 @@
 )
 
 
-class TestGFEMetricsIntegration(MockServerTestBase):
+class TestFrontendMetricsIntegration(MockServerTestBase):
     def setUp(self):
         super().setUp()
         os.environ["SPANNER_DISABLE_BUILTIN_METRICS"] = "false"
@@ -61,13 +61,13 @@ def test_gfe_metrics_exported(self):
         def custom_initial_metadata(self):
             mocked = getattr(self, "_is_execute_streaming_sql_mock", False)
             if mocked:
-                return (("server-timing", "gfet4t7; dur=55"),)
+                return (("server-timing", "gfet4t7; dur=55, afe; dur=23"),)
             return orig_initial_metadata(self)
 
         def custom_trailing_metadata(self):
             mocked = getattr(self, "_is_execute_streaming_sql_mock", False)
             if mocked:
-                return (("server-timing", "gfet4t7; dur=55"),)
+                return (("server-timing", "gfet4t7; dur=55, afe; dur=23"),)
             return orig_trailing_metadata(self)
 
         def custom_call(self_callable, request, *args, **kwargs):
@@ -142,6 +142,11 @@ def custom_call(self_callable, request, *args, **kwargs):
             point = next(iter(gfe_metric.data.data_points))
             self.assertEqual(point.sum, 55)
 
+            self.assertIn("afe_latency", metrics, f"Metrics: {list(metrics.keys())}")
+            afe_metric = metrics["afe_latency"]
+            point = next(iter(afe_metric.data.data_points))
+            self.assertEqual(point.sum, 23)
+
         finally:
             pass
 
@@ -202,5 +207,12 @@ def test_gfe_missing_header_count_exported(self):
             missing_metric = metrics["gfe_missing_header_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_missing_metric = metrics["afe_missing_header_count"]
+            afe_point = next(iter(afe_missing_metric.data.data_points))
+            self.assertGreaterEqual(afe_point.value, 1)
         finally:
             pass

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

@@ -46,6 +46,7 @@ def __init__(self):
         self.record_attempt_completion = MagicMock()
         self.set_method = MagicMock()
         self.record_gfe_metrics = MagicMock()
+        self.record_afe_metrics = MagicMock()
         self.set_project = MagicMock()
         self.set_instance = MagicMock()
         self.set_database = MagicMock()
@@ -118,4 +119,5 @@ def test_intercept_with_tracer(interceptor, mock_tracer_ctx):
     mock_tracer_ctx.record_attempt_start.assert_called()
     mock_tracer_ctx.record_attempt_completion.assert_called_once()
     mock_tracer_ctx.record_gfe_metrics.assert_called_once()
+    mock_tracer_ctx.record_afe_metrics.assert_called_once()
     mock_invoked_method.assert_called_once_with("request", call_details)

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

@@ -297,3 +297,69 @@ def test_record_gfe_metrics(metrics_tracer):
     metrics_tracer.record_gfe_metrics([("other", "1")])
     assert mock_gfe_latency.record.call_count == 1
     assert mock_gfe_missing.add.call_count == 1
+
+
+def test_record_afe_latency(metrics_tracer):
+    mock_afe_latency = mock.create_autospec(Histogram, instance=True)
+    metrics_tracer._instrument_afe_latency = mock_afe_latency
+    metrics_tracer.gfe_enabled = True
+
+    metrics_tracer.record_afe_latency(100)
+    assert mock_afe_latency.record.call_count == 1
+    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.enabled = False
+    metrics_tracer.record_afe_latency(200)
+    assert mock_afe_latency.record.call_count == 1
+    metrics_tracer.enabled = True
+
+
+def test_record_afe_missing_header_count(metrics_tracer):
+    mock_afe_missing = mock.create_autospec(Counter, instance=True)
+    metrics_tracer._instrument_afe_missing_header_count = mock_afe_missing
+    metrics_tracer.gfe_enabled = True
+
+    metrics_tracer.record_afe_missing_header_count()
+    assert mock_afe_missing.add.call_count == 1
+    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.enabled = False
+    metrics_tracer.record_afe_missing_header_count()
+    assert mock_afe_missing.add.call_count == 1
+    metrics_tracer.enabled = True
+
+
+def test_extract_afe_latency():
+    # Valid trailing metadata list of tuples
+    metadata_list = [("server-timing", "afe; dur=123")]
+    assert MetricsTracer.extract_afe_latency(metadata_list) == 123
+
+    # Valid metadata dict
+    metadata_dict = {"server-timing": "afet4t7; dur=456"}
+    assert MetricsTracer.extract_afe_latency(metadata_dict) == 456
+
+    # Missing header
+    assert MetricsTracer.extract_afe_latency([("other-header", "val")]) is None
+    assert MetricsTracer.extract_afe_latency(None) is None
+
+
+def test_record_afe_metrics(metrics_tracer):
+    mock_afe_latency = mock.create_autospec(Histogram, instance=True)
+    mock_afe_missing = mock.create_autospec(Counter, instance=True)
+    metrics_tracer._instrument_afe_latency = mock_afe_latency
+    metrics_tracer._instrument_afe_missing_header_count = mock_afe_missing
+    metrics_tracer.gfe_enabled = True
+
+    # With header
+    metrics_tracer.record_afe_metrics([("server-timing", "afe; dur=88")])
+    assert mock_afe_latency.record.call_count == 1
+    assert mock_afe_latency.record.call_args[1]["amount"] == 88
+    assert mock_afe_missing.add.call_count == 0
+
+    # Without header
+    metrics_tracer.record_afe_metrics([("other", "1")])
+    assert mock_afe_latency.record.call_count == 1
+    assert mock_afe_missing.add.call_count == 1
+