fix(metrics): remove obsolete gfe_enabled flag and refine AFE timing logic

Author: sinhasubham

packages/google-cloud-spanner/google/cloud/spanner_v1/_helpers.py

@@ -20,6 +20,7 @@
 import logging
 import math
 import operator
+import os
 import threading
 import time
 import uuid
@@ -69,6 +70,11 @@
 import random
 from typing import List, Tuple
 
+ENABLE_AFE_SERVER_TIMING = (
+    os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() != "true"
+    and os.environ.get("SPANNER_DISABLE_BUILTIN_METRICS", "").lower() != "true"
+)
+
 # Validation error messages
 NUMERIC_MAX_SCALE_ERR_MSG = (
     "Max scale for a numeric is 9. The requested numeric has scale {}"
@@ -707,6 +713,13 @@ def __init__(self, session):
         self._session = session
 
 
+def _append_routing_headers(metadata):
+    """Appends routing and backend-specific headers to the metadata."""
+    if ENABLE_AFE_SERVER_TIMING:
+        metadata.append(("x-goog-spanner-enable-afe-server-timing", "true"))
+    return metadata
+
+
 def _metadata_with_prefix(prefix, **kw):
     """Create RPC metadata containing a prefix.
 
@@ -716,7 +729,8 @@ def _metadata_with_prefix(prefix, **kw):
     Returns:
         List[Tuple[str, str]]: RPC metadata with supplied prefix
     """
-    return [("google-cloud-resource-prefix", prefix)]
+    metadata = [("google-cloud-resource-prefix", prefix)]
+    return _append_routing_headers(metadata)
 
 
 def _retry_on_aborted_exception(

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

@@ -16,7 +16,6 @@
 
 import inspect
 import logging
-import os
 import re
 from typing import Any, Dict
 
@@ -128,18 +127,11 @@ def intercept(self, invoked_method, request_or_iterator, call_details):
 
         ## Format method to be be spanner.<method name>
         method_str = call_details.method
-        if isinstance(method_str, bytes):
-            method_str = method_str.decode("utf-8")
         method_name = method_str.removeprefix(SPANNER_METHOD_PREFIX).replace("/", ".")
 
         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)
@@ -216,18 +208,11 @@ async def _async_intercept(
             MetricsInterceptor._set_metrics_tracer_attributes(resources)
 
         method_str = call_details.method
-        if isinstance(method_str, bytes):
-            method_str = method_str.decode("utf-8")
         method_name = method_str.removeprefix(SPANNER_METHOD_PREFIX).replace("/", ".")
 
         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 = await continuation(call_details, request_or_iterator)
         if hasattr(response, "__anext__"):
             return _AsyncStreamingResponseWrapper(response, tracer)

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

@@ -190,7 +190,6 @@ class should not have any knowledge about the observability framework used for m
     _instrument_afe_connectivity_error_count: "Counter"
     current_op: MetricOpTracer
     enabled: bool
-    gfe_enabled: bool
     method: str
 
     def __init__(
@@ -205,7 +204,6 @@ def __init__(
         instrument_gfe_connectivity_error_count: "Counter",
         instrument_afe_latency: "Histogram",
         instrument_afe_connectivity_error_count: "Counter",
-        gfe_enabled: bool = False,
     ):
         """
         Initialize a MetricsTracer instance with the given parameters.
@@ -221,7 +219,6 @@ def __init__(
             instrument_operation_latency (Histogram): Instrument for measuring operation latency.
             instrument_operation_counter (Counter): Instrument for counting operations.
             client_attributes (Dict[str, str]): Dictionary of client attributes used for metrics tracing.
-            gfe_enabled (bool, optional): Indicates if GFE metrics are enabled. Defaults to False.
             instrument_gfe_latency (Histogram): Instrument for measuring GFE latency.
             instrument_gfe_connectivity_error_count (Counter): Instrument for counting GFE connectivity errors.
             instrument_afe_latency (Histogram): Instrument for measuring AFE latency.
@@ -242,7 +239,9 @@ def __init__(
             instrument_afe_connectivity_error_count
         )
         self.enabled = enabled
-        self.gfe_enabled = True
+        self.afe_server_timing_enabled = (
+            os.environ.get("SPANNER_DISABLE_AFE_SERVER_TIMING", "").lower() != "true"
+        )
 
     @staticmethod
     def _get_ms_time_diff(start: datetime, end: datetime) -> float:
@@ -454,7 +453,7 @@ 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"
+            or not getattr(self, "afe_server_timing_enabled", True)
         ):
             return
         self._instrument_afe_latency.record(
@@ -469,7 +468,7 @@ 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"
+            or not getattr(self, "afe_server_timing_enabled", True)
         ):
             return
         self._instrument_afe_connectivity_error_count.add(
@@ -500,7 +499,7 @@ def extract_front_end_latencies(
         header_vals = []
         for key, val in items:
             key_str = key.decode("utf-8") if isinstance(key, bytes) else str(key)
-            if key_str and key_str.lower() in ("server-timing", "server_timing"):
+            if key_str and key_str.lower() == "server-timing":
                 if isinstance(val, (list, tuple)):
                     header_vals.extend(val)
                 else:
@@ -529,7 +528,7 @@ def extract_front_end_latencies(
                         pass
 
             if afe_latency is None:
-                match = re.search(r"afe(?:t4t7)?;\s*dur=([0-9.]+)", header_val)
+                match = re.search(r"afe;\s*dur=([0-9.]+)", header_val)
                 if match:
                     try:
                         afe_latency = int(float(match.group(1)))

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

@@ -52,7 +52,6 @@ class MetricsTracerFactory:
     """Factory class for creating MetricTracer instances. This class facilitates the creation of MetricTracer objects, which are responsible for collecting and tracing metrics."""
 
     enabled: bool
-    gfe_enabled: bool
     _instrument_attempt_latency: "Histogram"
     _instrument_attempt_counter: "Counter"
     _instrument_operation_latency: "Histogram"
@@ -89,7 +88,6 @@ def __init__(self, enabled: bool, service_name: str):
             project (str): The project ID for the monitored resource.
         """
         self.enabled = enabled
-        self.gfe_enabled = True
         self._create_metric_instruments(service_name)
         self._client_attributes = {}
 
@@ -273,7 +271,6 @@ def create_metrics_tracer(self) -> MetricsTracer:
             instrument_operation_latency=self._instrument_operation_latency,
             instrument_operation_counter=self._instrument_operation_counter,
             client_attributes=self._client_attributes.copy(),
-            gfe_enabled=True,
             instrument_gfe_latency=self._instrument_gfe_latency,
             instrument_gfe_connectivity_error_count=self._instrument_gfe_connectivity_error_count,
             instrument_afe_latency=self._instrument_afe_latency,

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

@@ -80,7 +80,6 @@ def __new__(cls, enabled: bool = True) -> "SpannerMetricsTracerFactory":
                 cls._generate_client_hash(client_uid)
             )
             cls._metrics_tracer_factory.set_location(_get_cloud_region())
-            cls._metrics_tracer_factory.gfe_enabled = True
 
             if cls._metrics_tracer_factory.enabled != enabled:
                 cls._metrics_tracer_factory.enabled = enabled

packages/google-cloud-spanner/tests/system/test_metrics.py

@@ -79,6 +79,7 @@ def test_builtin_metrics_with_default_otel(metrics_database):
         "spanner/operation_count",
         "spanner/attempt_count",
         "spanner/gfe_latencies",
+        "spanner/afe_latencies",
     }
     assert expected_metrics.issubset(collected_metrics)
 

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

@@ -41,7 +41,6 @@ def __init__(self):
         self.project = None
         self.instance = None
         self.database = None
-        self.gfe_enabled = True
         self.record_attempt_start = MagicMock()
         self.record_attempt_completion = MagicMock()
         self.set_method = MagicMock()
@@ -113,10 +112,9 @@ 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", replaced_call_details)
+    mock_invoked_method.assert_called_once_with("request", call_details)

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

@@ -235,7 +235,6 @@ def test_set_method(metrics_tracer):
 def test_record_gfe_latency(metrics_tracer):
     mock_gfe_latency = mock.create_autospec(Histogram, instance=True)
     metrics_tracer._instrument_gfe_latency = mock_gfe_latency
-    metrics_tracer.gfe_enabled = True  # Ensure GFE is enabled
 
     # Test when tracing is enabled
     metrics_tracer.record_gfe_latency(100)
@@ -258,7 +257,6 @@ def test_record_gfe_connectivity_error_count(metrics_tracer):
     metrics_tracer._instrument_gfe_connectivity_error_count = (
         mock_gfe_connectivity_error_count
     )
-    metrics_tracer.gfe_enabled = True  # Ensure GFE is enabled
 
     # Test when tracing is enabled
     metrics_tracer.record_gfe_connectivity_error_count()
@@ -305,7 +303,6 @@ def test_record_front_end_metrics(metrics_tracer):
     metrics_tracer._instrument_gfe_connectivity_error_count = mock_gfe_missing
     metrics_tracer._instrument_afe_latency = mock_afe_latency
     metrics_tracer._instrument_afe_connectivity_error_count = mock_afe_missing
-    metrics_tracer.gfe_enabled = True
 
     # With header
     metrics_tracer.record_front_end_metrics(
@@ -329,7 +326,6 @@ def test_record_front_end_metrics(metrics_tracer):
 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
@@ -339,8 +335,8 @@ def test_record_afe_latency(metrics_tracer):
         == metrics_tracer._create_attempt_otel_attributes()
     )
 
-    with mock.patch.dict("os.environ", {"SPANNER_DISABLE_AFE_SERVER_TIMING": "true"}):
-        metrics_tracer.record_afe_latency(300)
+    metrics_tracer.afe_server_timing_enabled = False
+    metrics_tracer.record_afe_latency(300)
     assert mock_afe_latency.record.call_count == 1
 
     metrics_tracer.enabled = False
@@ -352,7 +348,6 @@ def test_record_afe_latency(metrics_tracer):
 def test_record_afe_connectivity_error_count(metrics_tracer):
     mock_afe_missing = mock.create_autospec(Counter, instance=True)
     metrics_tracer._instrument_afe_connectivity_error_count = mock_afe_missing
-    metrics_tracer.gfe_enabled = True
 
     metrics_tracer.record_afe_connectivity_error_count()
     assert mock_afe_missing.add.call_count == 1
@@ -362,8 +357,8 @@ def test_record_afe_connectivity_error_count(metrics_tracer):
         == 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()
+    metrics_tracer.afe_server_timing_enabled = False
+    metrics_tracer.record_afe_connectivity_error_count()
     assert mock_afe_missing.add.call_count == 1
 
     metrics_tracer.enabled = False