Repository: jina-ai/serve
Pull request: feat: add deployment label in grpc stub metrics
<!---Decomposing the complex issue into subtasks can help you build it step-by-step. Thanks for your pull request! :rocket: --->
<!---We know that dev life is hectic, but **please provide a (brief) description** of what your PR does, and how it does it. **Otherwise, your PR cannot be reviewed!** --->
<!---This policy was agreed upon in a past company retro, and makes everyone's life a little easier. Thanks for your collaboration!--->

**Goals:**
<!---https://docs.github.com/en/issues/tracking-your-work-with-issues/linking-a-pull-request-to-an-issue#linking-a-pull-request-to-an-issue-using-a-keyword--->

- Resolves #5154
- [x] Handle connection reset case
- [x] Add regression tests



Problem statement:
Metrics on the gateway should always contained a label based on the Executor target
# Context

At the moment the gateway average all of its metrics and does not allow users to make distinction between different Executor.

To give an example:

Lets say we have a flow with two executors, one takes a lot of time to process the other one take few times.

At the moment the `jina_sending_request_seconds` metrics for instance will just average both of them

Actually this was not an issue initially since you can actually do the separation at the at the Executor level but the issue is with external Executor where you will have metrics coming from all the different flow.

# Proposition

Therefore we should rather measure the metrics on gateway level depending on which Executor the request is send.

Note: telemetry and tracing will also fix this problem

Hints from issue discussion:
I think then we should duplicate metrics and have `gateway` measure Executor latency regardless of anything
gotcha makes sense

Commit message: test: add tests for gateway labels

Patch to review:
diff --git a/jina/serve/networking.py b/jina/serve/networking.py
index 5f18a15d44762..ada96400a1a89 100644
--- a/jina/serve/networking.py
+++ b/jina/serve/networking.py
@@ -64,19 +64,31 @@ class _NetworkingHistograms:
     send_requests_bytes_metrics: Optional['Histogram'] = None
     histogram_metric_labels: Dict[str, str] = None
 
-    def record_sending_requests_time_metrics(self, value: int):
+    def get_labels(self, additional_labels: Optional[Dict[str, str]] = None):
+        # TODO: This is quite ugly. Surely there's a better way
+        if self.histogram_metric_labels is None:
+            return None
+        if additional_labels is None:
+            return self.histogram_metric_labels
+        return {**self.histogram_metric_labels, **additional_labels}
+
+    def record_sending_requests_time_metrics(self, value: int, additional_labels: Optional[Dict[str, str]] = None):
+        labels = self.get_labels(additional_labels)
+            
         if self.sending_requests_time_metrics:
-            self.sending_requests_time_metrics.record(
-                value, self.histogram_metric_labels
-            )
+            self.sending_requests_time_metrics.record(value, labels)
+
+    def record_received_response_bytes(self, value: int, additional_labels: Optional[Dict[str, str]] = None):
+        labels = self.get_labels(additional_labels)
 
-    def record_received_response_bytes(self, value: int):
         if self.received_response_bytes:
-            self.received_response_bytes.record(value, self.histogram_metric_labels)
+            self.received_response_bytes.record(value, labels)
+
+    def record_send_requests_bytes_metrics(self, value: int, additional_labels: Optional[Dict[str, str]] = None):
+        labels = self.get_labels(additional_labels)
 
-    def record_send_requests_bytes_metrics(self, value: int):
         if self.send_requests_bytes_metrics:
-            self.send_requests_bytes_metrics.record(value, self.histogram_metric_labels)
+            self.send_requests_bytes_metrics.record(value, labels)
 
 
 class ReplicaList:
@@ -105,13 +117,14 @@ def __init__(
         self.aio_tracing_client_interceptors = aio_tracing_client_interceptors
         self.tracing_client_interceptors = tracing_client_interceptor
 
-    async def reset_connection(self, address: str) -> Union[grpc.aio.Channel, None]:
+    async def reset_connection(self, address: str, deployment_name: str) -> Union[grpc.aio.Channel, None]:
         """
         Removes and then re-adds a connection.
         Result is the same as calling :meth:`remove_connection` and then :meth:`add_connection`, but this allows for
         handling of race condition if multiple callers reset a connection at the same time.
 
         :param address: Target address of this connection
+        :param deployment_name: Target deployment of this connection
         :returns: The reset connection or None if there was no connection for the given address
         """
         self._logger.debug(f'resetting connection to {address}')
@@ -134,21 +147,22 @@ async def reset_connection(self, address: str) -> Union[grpc.aio.Channel, None]:
             self._destroyed_event.set()
             # re-add connection:
             self._address_to_connection_idx[address] = id_to_reset
-            stubs, channel = self._create_connection(address)
+            stubs, channel = self._create_connection(address, deployment_name)
             self._address_to_channel[address] = channel
             self._connections[id_to_reset] = stubs
 
             return connection_to_reset
         return None
 
-    def add_connection(self, address: str):
+    def add_connection(self, address: str, deployment_name: str):
         """
         Add connection with address to the connection list
         :param address: Target address of this connection
+        :param deployment_name: Target deployment of this connection
         """
         if address not in self._address_to_connection_idx:
             self._address_to_connection_idx[address] = len(self._connections)
-            stubs, channel = self._create_connection(address)
+            stubs, channel = self._create_connection(address, deployment_name)
             self._address_to_channel[address] = channel
             self._connections.append(stubs)
 
@@ -187,13 +201,14 @@ async def remove_connection(self, address: str) -> Union[grpc.aio.Channel, None]
 
         return None
 
-    def _create_connection(self, address):
+    def _create_connection(self, address, deployment_name: str):
         parsed_address = urlparse(address)
         address = parsed_address.netloc if parsed_address.netloc else address
         use_tls = parsed_address.scheme in TLS_PROTOCOL_SCHEMES
 
         stubs, channel = GrpcConnectionPool.create_async_channel_stub(
             address,
+            deployment_name=deployment_name,
             metrics=self._metrics,
             histograms=self._histograms,
             tls=use_tls,
@@ -315,15 +330,23 @@ def __init__(
             self,
             address,
             channel,
+            deployment_name: str,
             metrics: _NetworkingMetrics,
             histograms: _NetworkingHistograms,
         ):
             self.address = address
             self.channel = channel
+            self.deployment_name = deployment_name
             self._metrics = metrics
             self._histograms = histograms
             self._initialized = False
 
+            if self._histograms:
+                self.stub_specific_labels = {
+                    'deployment': deployment_name,
+                    'address': address,
+                }
+
         # This has to be done lazily, because the target endpoint may not be available
         # when a connection is added
         async def _init_stubs(self):
@@ -363,6 +386,31 @@ async def send_discover_endpoint(
             )
             return response, metadata
 
+        def _get_metric_timer(self):
+            if self._histograms.histogram_metric_labels is None:
+                labels = None
+            else:
+                labels = {
+                    **self._histograms.histogram_metric_labels,
+                    **self.stub_specific_labels,
+                }
+
+            return MetricsTimer(
+                self._metrics.sending_requests_time_metrics,
+                self._histograms.sending_requests_time_metrics,
+                labels,
+            )
+
+        def _record_request_bytes_metric(self, nbytes: int):
+            if self._metrics.send_requests_bytes_metrics:
+                self._metrics.send_requests_bytes_metrics.observe(nbytes)
+            self._histograms.record_send_requests_bytes_metrics(nbytes, self.stub_specific_labels)
+
+        def _record_received_bytes_metric(self, nbytes: int):
+            if self._metrics.received_response_bytes:
+                self._metrics.received_response_bytes.observe(nbytes)
+            self._histograms.record_received_response_bytes(nbytes, self.stub_specific_labels)
+
         async def send_requests(
             self,
             requests: List[Request],
@@ -385,6 +433,7 @@ async def send_requests(
                 await self._init_stubs()
             request_type = type(requests[0])
 
+            timer = self._get_metric_timer()
             if request_type == DataRequest and len(requests) == 1:
                 if self.single_data_stub:
                     call_result = self.single_data_stub.process_single_data(
@@ -393,61 +442,32 @@ async def send_requests(
                         compression=compression,
                         timeout=timeout,
                     )
-                    if self._metrics.send_requests_bytes_metrics:
-                        self._metrics.send_requests_bytes_metrics.observe(
-                            requests[0].nbytes
-                        )
-                    self._histograms.record_send_requests_bytes_metrics(
-                        requests[0].nbytes
-                    )
-                    with MetricsTimer(
-                        self._metrics.sending_requests_time_metrics,
-                        self._histograms.sending_requests_time_metrics,
-                        self._histograms.histogram_metric_labels,
-                    ):
+                    self._record_request_bytes_metric(requests[0].nbytes)
+                    with timer:
                         metadata, response = (
                             await call_result.trailing_metadata(),
                             await call_result,
                         )
-
-                        if self._metrics.received_response_bytes:
-                            self._metrics.received_response_bytes.observe(
-                                response.nbytes
-                            )
-                        self._histograms.record_received_response_bytes(response.nbytes)
+                        self._record_received_bytes_metric(response.nbytes)
                     return response, metadata
 
                 elif self.stream_stub:
                     for response in requests:
-                        if self._metrics.send_requests_bytes_metrics:
-                            self._metrics.send_requests_bytes_metrics.observe(
-                                response.nbytes
-                            )
-                        self._histograms.record_send_requests_bytes_metrics(
-                            response.nbytes
-                        )
+                        self._record_request_bytes_metric(response.nbytes)
 
-                    with MetricsTimer(
-                        self._metrics.sending_requests_time_metrics,
-                        self._histograms.sending_requests_time_metrics,
-                        self._histograms.histogram_metric_labels,
-                    ):
+                    with timer:
                         async for response in self.stream_stub.Call(
                             iter(requests),
                             compression=compression,
                             timeout=timeout,
                             metadata=metadata,
                         ):
-                            if self._metrics.received_response_bytes:
-                                self._metrics.received_response_bytes.observe(
-                                    response.nbytes
-                                )
-                            self._histograms.record_received_response_bytes(
-                                response.nbytes
-                            )
-
+                            self._record_received_bytes_metric(response.nbytes)
                             return response, None
+
             if request_type == DataRequest and len(requests) > 1:
+                # TODO: This section of the code has no observability
+                # TODO: Does this section not need to deal with streaming stubs?
                 if self.data_list_stub:
                     call_result = self.data_list_stub.process_data(
                         requests,
@@ -455,11 +475,7 @@ async def send_requests(
                         compression=compression,
                         timeout=timeout,
                     )
-                    with MetricsTimer(
-                        self._metrics.sending_requests_time_metrics,
-                        self._histograms.sending_requests_time_metrics,
-                        self._histograms.histogram_metric_labels,
-                    ):
+                    with timer:
                         metadata, response = (
                             await call_result.trailing_metadata(),
                             await call_result,
@@ -614,7 +630,7 @@ def _add_connection(
                 self._logger.debug(
                     f'adding connection for deployment {deployment}/{type}/{entity_id} to {address}'
                 )
-                self._deployments[deployment][type][entity_id].add_connection(address)
+                self._deployments[deployment][type][entity_id].add_connection(address, deployment_name=deployment)
             else:
                 self._logger.debug(
                     f'ignoring activation of pod, {address} already known'
@@ -987,6 +1003,7 @@ async def _handle_aiorpcerror(
         },  # same deployment can have multiple addresses (replicas)
         total_num_tries: int = 1,  # number of retries + 1
         current_address: str = '',  # the specific address that was contacted during this attempt
+        current_deployment: str = '',  # the specific deployment that was contacted during this attempt
         connection_list: Optional[ReplicaList] = None,
     ):
         # connection failures, cancelled requests, and timed out requests should be retried
@@ -1012,7 +1029,7 @@ async def _handle_aiorpcerror(
             # after connection failure the gRPC `channel` gets stuck in a failure state for a few seconds
             # removing and re-adding the connection (stub) is faster & more reliable than just waiting
             if connection_list:
-                await connection_list.reset_connection(current_address)
+                await connection_list.reset_connection(current_address, current_deployment)
 
             raise InternalNetworkError(
                 og_exception=error,
@@ -1074,6 +1091,7 @@ async def task_wrapper():
                         tried_addresses=tried_addresses,
                         total_num_tries=total_num_tries,
                         current_address=current_connection.address,
+                        current_deployment=current_connection.deployment_name,
                         connection_list=connections,
                     )
 
@@ -1115,6 +1133,7 @@ async def task_wrapper():
                         retry_i=i,
                         tried_addresses=tried_addresses,
                         current_address=connection.address,
+                        current_deployment=connection.deployment_name,
                         connection_list=connection_list,
                         total_num_tries=total_num_tries,
                     )
@@ -1379,6 +1398,7 @@ async def send_request_async(
     @staticmethod
     def create_async_channel_stub(
         address,
+        deployment_name: str,
         metrics: _NetworkingMetrics,
         histograms: _NetworkingHistograms,
         tls=False,
@@ -1389,6 +1409,7 @@ def create_async_channel_stub(
         Creates an async GRPC Channel. This channel has to be closed eventually!
 
         :param address: the address to create the connection to, like 127.0.0.0.1:8080
+        :param deployment_name: the name of the deployment (e.g. executor0)
         :param tls: if True, use tls for the grpc channel
         :param root_certificates: the path to the root certificates for tls, only u
         :param metrics: NetworkingMetrics object that contain optional metrics
@@ -1405,7 +1426,7 @@ def create_async_channel_stub(
         )
 
         return (
-            GrpcConnectionPool.ConnectionStubs(address, channel, metrics, histograms),
+            GrpcConnectionPool.ConnectionStubs(address, channel, deployment_name, metrics, histograms),
             channel,
         )
 
diff --git a/tests/unit/serve/instrumentation/conftest.py b/tests/unit/serve/instrumentation/conftest.py
new file mode 100644
index 0000000000000..c06261eabb134
--- /dev/null
+++ b/tests/unit/serve/instrumentation/conftest.py
@@ -0,0 +1,119 @@
+import random
+import pytest
+from pathlib import Path
+from typing import Dict, Tuple, Callable
+import opentelemetry.sdk.metrics.export
+import opentelemetry.sdk.metrics.view
+from opentelemetry.sdk.metrics.export import (
+    AggregationTemporality,
+    MetricExporter,
+    MetricExportResult,
+    MetricsData,
+)
+
+
+class DirMetricExporter(MetricExporter):
+    """Implementation of :class:`MetricExporter` that prints metrics to a file in a given directory.
+
+    This class can be used for diagnostic or testing purposes.
+    """
+
+    def __init__(
+        self,
+        metric_dir: str,
+        preferred_temporality: Dict[type, AggregationTemporality] = None,
+        preferred_aggregation: Dict[
+            type, "opentelemetry.sdk.metrics.view.Aggregation"
+        ] = None,
+    ):
+        super().__init__(
+            preferred_temporality=preferred_temporality,
+            preferred_aggregation=preferred_aggregation,
+        )
+        self.metric_filename: Path = Path(metric_dir) / str(random.randint(0, 1048575))
+        self.f = open(self.metric_filename, 'a')
+
+    def export(
+        self,
+        metrics_data: MetricsData,
+        timeout_millis: float = 10_000,
+        **kwargs,
+    ) -> MetricExportResult:
+        self.f.write(metrics_data.to_json())
+        self.f.write('\n')
+        self.f.flush()
+        return MetricExportResult.SUCCESS
+
+    def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None:
+        pass
+
+    def force_flush(self, timeout_millis: float = 10_000) -> bool:
+        return True
+
+    def __del__(self):
+        self.f.close()
+
+
+@pytest.fixture(scope='session')
+def monkeypatch_metric_exporter(
+    tmpdir_factory: pytest.TempdirFactory,
+) -> Tuple[Callable, Callable]:
+    import opentelemetry.sdk.metrics.export
+    from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader
+    from pathlib import Path
+    import time
+    import os
+    import json
+
+    collect_path = Path(tmpdir_factory.mktemp('otel-collector'))
+    metrics_path = collect_path / 'metrics'
+    os.mkdir(metrics_path)
+
+    tick_counter_filename = collect_path / 'tick_counter'
+    with open(tick_counter_filename, 'w') as f:
+        f.write('0')
+
+    def collect_metrics():
+        with open(tick_counter_filename, 'r') as f:
+            tick_counter = int(f.read())
+        with open(tick_counter_filename, 'w') as f:
+            f.write(str(tick_counter + 1))
+        time.sleep(1)
+
+    def _get_service_name(otel_measurement):
+        return otel_measurement[0]['resource_metrics'][0]['resource']['attributes'][
+            'service.name'
+        ]
+
+    def read_metrics():
+        def read_metric_file(filename):
+            with open(filename, 'r') as f:
+                return list(map(json.loads, f.readlines()))
+
+        return {
+            _get_service_name(i): i
+            for i in map(read_metric_file, metrics_path.glob('*'))
+        }
+
+    class PatchedTextReader(PeriodicExportingMetricReader):
+        def __init__(self, *args, **kwargs) -> None:
+            self.exporter = DirMetricExporter(metrics_path)
+            self.tick_counter = 0
+
+            super().__init__(
+                exporter=self.exporter,
+                export_interval_millis=1_000,
+            )
+
+        def _ticker(self) -> None:
+            interval_secs = self._export_interval_millis / 1e3
+            while not self._shutdown_event.wait(interval_secs):
+                with open(tick_counter_filename, 'r') as f:
+                    tick_counter = int(f.read())
+                if tick_counter != self.tick_counter:
+                    self.tick_counter = tick_counter
+                    self.collect(timeout_millis=self._export_timeout_millis)
+            self.collect(timeout_millis=self._export_interval_millis)
+
+    opentelemetry.sdk.metrics.export.PeriodicExportingMetricReader = PatchedTextReader
+    return collect_metrics, read_metrics
diff --git a/tests/unit/serve/instrumentation/test_gateway_metric_labels.py b/tests/unit/serve/instrumentation/test_gateway_metric_labels.py
new file mode 100644
index 0000000000000..1167e50abe95f
--- /dev/null
+++ b/tests/unit/serve/instrumentation/test_gateway_metric_labels.py
@@ -0,0 +1,115 @@
+from jina import Executor, Flow, requests, DocumentArray
+
+
+def test_gateway_metric_labels(monkeypatch_metric_exporter):
+    collect_metrics, read_metrics = monkeypatch_metric_exporter
+
+    class FirstExec(Executor):
+        @requests()
+        def meow(self, docs, **kwargs):
+            return DocumentArray.empty(3)
+
+    class SecondExec(Executor):
+        @requests()
+        def meow(self, docs, **kwargs):
+            return DocumentArray.empty(3)
+
+    with Flow(
+        tracing=False,
+        metrics=True,
+        metrics_exporter_host='localhost',
+        metrics_exporter_port=4317,
+        port=12345,
+    ).add(name='first_exec', uses=FirstExec).add(
+        name="second_exec", uses=SecondExec
+    ) as f:
+        f.post('/')
+        f.plot('flow.png')
+
+        collect_metrics()
+        metrics = read_metrics()
+        gateway_metrics = metrics['gateway/rep-0'][0]['resource_metrics'][0][
+            'scope_metrics'
+        ][0]['metrics']
+        gateway_metric_data_point = {
+            i['name']: i['data']['data_points'] for i in gateway_metrics
+        }
+
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_sending_request_seconds'][0][
+                'attributes'
+            ]
+        )
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_sent_request_bytes'][0]['attributes']
+        )
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_received_response_bytes'][0][
+                'attributes'
+            ]
+        )
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_sending_request_seconds'][1][
+                'attributes'
+            ]
+        )
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_sent_request_bytes'][1]['attributes']
+        )
+        assert (
+            'address'
+            in gateway_metric_data_point['jina_received_response_bytes'][1][
+                'attributes'
+            ]
+        )
+
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_sending_request_seconds'][0][
+                'attributes'
+            ]
+        )
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_sent_request_bytes'][0]['attributes']
+        )
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_received_response_bytes'][0][
+                'attributes'
+            ]
+        )
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_sending_request_seconds'][1][
+                'attributes'
+            ]
+        )
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_sent_request_bytes'][1]['attributes']
+        )
+        assert (
+            'deployment'
+            in gateway_metric_data_point['jina_received_response_bytes'][1][
+                'attributes'
+            ]
+        )
+
+        assert (
+            gateway_metric_data_point['jina_received_response_bytes'][0]['attributes'][
+                'deployment'
+            ]
+            == 'first_exec'
+        )
+        assert (
+            gateway_metric_data_point['jina_received_response_bytes'][1]['attributes'][
+                'deployment'
+            ]
+            == 'second_exec'
+        )
diff --git a/tests/unit/serve/instrumentation/test_networking_histograms.py b/tests/unit/serve/instrumentation/test_networking_histograms.py
new file mode 100644
index 0000000000000..6d11d511b1f72
--- /dev/null
+++ b/tests/unit/serve/instrumentation/test_networking_histograms.py
@@ -0,0 +1,68 @@
+import pytest
+from typing import Tuple, List, Dict
+from opentelemetry.sdk.metrics import MeterProvider
+from opentelemetry.sdk.metrics.export import (
+    InMemoryMetricReader,
+    Metric,
+    HistogramDataPoint,
+)
+from opentelemetry.metrics import Meter
+
+from jina.serve.networking import _NetworkingHistograms
+
+
+@pytest.fixture
+def metrics_setup() -> Tuple[InMemoryMetricReader, MeterProvider]:
+    metric_reader = InMemoryMetricReader()
+    meter_provider = MeterProvider(metric_readers=[metric_reader])
+    meter = meter_provider.get_meter('test')
+    yield metric_reader, meter
+    if hasattr(meter_provider, 'force_flush'):
+        metric_reader.force_flush()
+    if hasattr(meter_provider, 'shutdown'):
+        meter_provider.shutdown()
+
+
+def test_get_labels():
+    a: _NetworkingHistograms = _NetworkingHistograms()
+    assert a.get_labels() == None
+
+    HIST_LABELS = {
+        'a': 1,
+        'b': 2,
+    }
+    a.histogram_metric_labels = HIST_LABELS
+    assert a.get_labels() == HIST_LABELS
+
+    ADD_LABELS = {
+        'b': 3,
+        'c': 4,
+    }
+    assert a.get_labels(ADD_LABELS) == {**HIST_LABELS, **ADD_LABELS}
+
+
+def test_recording_methods(metrics_setup: Tuple[InMemoryMetricReader, Meter]):
+    metric_reader, meter = metrics_setup
+
+    a: _NetworkingHistograms = _NetworkingHistograms(
+        sending_requests_time_metrics=meter.create_histogram("request_time"),
+        send_requests_bytes_metrics=meter.create_histogram("request_bytes"),
+        received_response_bytes=meter.create_histogram("response_bytes"),
+        histogram_metric_labels=None,
+    )
+
+    a.record_sending_requests_time_metrics(10)
+    a.record_send_requests_bytes_metrics(20)
+    a.record_received_response_bytes(30)
+
+    histogram_metrics: List[Metric] = (
+        metric_reader.get_metrics_data().resource_metrics[0].scope_metrics[0].metrics
+    )
+    data_points_sums: Dict[str, HistogramDataPoint] = {
+        hist.name: next(hist.data.data_points).sum for hist in histogram_metrics
+    }
+    assert data_points_sums == {
+        'request_time': 10,
+        'request_bytes': 20,
+        'response_bytes': 30,
+    }


Code review comments:
diff_hunk: @@ -64,19 +64,31 @@ class _NetworkingHistograms:
     send_requests_bytes_metrics: Optional['Histogram'] = None
     histogram_metric_labels: Dict[str, str] = None
 
-    def record_sending_requests_time_metrics(self, value: int):
+    def get_labels(self, additional_labels: Optional[Dict[str, str]] = None):
path: jina/serve/networking.py
text: @user1:
seems like it should be a private function