Use force_flush (#2853)

* Use force_flush

Fixes #2816

* Refactor time out passing

* Fix test cases

* Add documentation

* Add test case

* Fix docs

* Remove final decorator

* Check for timeout in async calls

* Remove timeout checking in metric reader

* Add force_flush to PeriodicExportingMetricReader

* Use TimeoutError

* Fix lint

* Address timeout comments

* Add MetricsTimeoutError

* Fix test case

Author: ocelotl

opentelemetry-sdk/src/opentelemetry/sdk/metrics/__init__.py

@@ -18,6 +18,9 @@
     Meter,
     MeterProvider,
 )
+from opentelemetry.sdk.metrics._internal.exceptions import (  # noqa: F401
+    MetricsTimeoutError,
+)
 from opentelemetry.sdk.metrics._internal.instrument import (  # noqa: F401
     Counter,
     Histogram,

opentelemetry-sdk/src/opentelemetry/sdk/metrics/_internal/__init__.py

@@ -30,6 +30,7 @@
     ObservableUpDownCounter as APIObservableUpDownCounter,
 )
 from opentelemetry.metrics import UpDownCounter as APIUpDownCounter
+from opentelemetry.sdk.metrics._internal.exceptions import MetricsTimeoutError
 from opentelemetry.sdk.metrics._internal.instrument import (
     _Counter,
     _Histogram,
@@ -384,8 +385,10 @@ def force_flush(self, timeout_millis: float = 10_000) -> bool:
             current_ts = _time_ns()
             try:
                 if current_ts >= deadline_ns:
-                    raise Exception("Timed out while flushing metric readers")
-                metric_reader.collect(
+                    raise MetricsTimeoutError(
+                        "Timed out while flushing metric readers"
+                    )
+                metric_reader.force_flush(
                     timeout_millis=(deadline_ns - current_ts) / 10**6
                 )
 

opentelemetry-sdk/src/opentelemetry/sdk/metrics/_internal/exceptions.py

@@ -0,0 +1,17 @@
+# Copyright The OpenTelemetry Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+class MetricsTimeoutError(Exception):
+    """Raised when a metrics function times out"""

opentelemetry-sdk/src/opentelemetry/sdk/metrics/_internal/export/__init__.py

@@ -294,14 +294,23 @@ def __init__(
     def collect(self, timeout_millis: float = 10_000) -> None:
         """Collects the metrics from the internal SDK state and
         invokes the `_receive_metrics` with the collection.
+
+        Args:
+            timeout_millis: Amount of time in milliseconds before this function
+              raises a timeout error.
+
+        If any of the underlying ``collect`` methods called by this method
+        fails by any reason (including timeout) an exception will be raised
+        detailing the individual errors that caused this function to fail.
         """
         if self._collect is None:
             _logger.warning(
                 "Cannot call collect on a MetricReader until it is registered on a MeterProvider"
             )
             return
+
         self._receive_metrics(
-            self._collect(self),
+            self._collect(self, timeout_millis=timeout_millis),
             timeout_millis=timeout_millis,
         )
 
@@ -328,6 +337,10 @@ def _receive_metrics(
     ) -> None:
         """Called by `MetricReader.collect` when it receives a batch of metrics"""
 
+    def force_flush(self, timeout_millis: float = 10_000) -> bool:
+        self.collect(timeout_millis=timeout_millis)
+        return True
+
     @abstractmethod
     def shutdown(self, timeout_millis: float = 30_000, **kwargs) -> None:
         """Shuts down the MetricReader. This method provides a way
@@ -485,3 +498,8 @@ def _shutdown():
         self._shutdown_event.set()
         self._daemon_thread.join(timeout=(deadline_ns - _time_ns()) / 10**9)
         self._exporter.shutdown(timeout=(deadline_ns - _time_ns()) / 10**6)
+
+    def force_flush(self, timeout_millis: float = 10_000) -> bool:
+        super().force_flush(timeout_millis=timeout_millis)
+        self._exporter.force_flush(timeout_millis=timeout_millis)
+        return True

opentelemetry-sdk/src/opentelemetry/sdk/metrics/_internal/measurement_consumer.py

@@ -23,11 +23,13 @@
 import opentelemetry.sdk.metrics._internal.instrument
 import opentelemetry.sdk.metrics._internal.sdk_configuration
 from opentelemetry.metrics._internal.instrument import CallbackOptions
+from opentelemetry.sdk.metrics._internal.exceptions import MetricsTimeoutError
 from opentelemetry.sdk.metrics._internal.measurement import Measurement
 from opentelemetry.sdk.metrics._internal.metric_reader_storage import (
     MetricReaderStorage,
 )
 from opentelemetry.sdk.metrics._internal.point import Metric
+from opentelemetry.util._time import _time_ns
 
 
 class MeasurementConsumer(ABC):
@@ -48,6 +50,7 @@ def register_asynchronous_instrument(
     def collect(
         self,
         metric_reader: "opentelemetry.sdk.metrics.MetricReader",
+        timeout_millis: float = 10_000,
     ) -> Iterable[Metric]:
         pass
 
@@ -90,12 +93,34 @@ def register_asynchronous_instrument(
     def collect(
         self,
         metric_reader: "opentelemetry.sdk.metrics.MetricReader",
+        timeout_millis: float = 10_000,
     ) -> Iterable[Metric]:
+
         with self._lock:
             metric_reader_storage = self._reader_storages[metric_reader]
             # for now, just use the defaults
             callback_options = CallbackOptions()
+            deadline_ns = _time_ns() + timeout_millis * 10**6
+
+            default_timeout_millis = 10000 * 10**6
+
             for async_instrument in self._async_instruments:
-                for measurement in async_instrument.callback(callback_options):
+
+                remaining_time = deadline_ns - _time_ns()
+
+                if remaining_time < default_timeout_millis:
+
+                    callback_options = CallbackOptions(
+                        timeout_millis=remaining_time
+                    )
+
+                measurements = async_instrument.callback(callback_options)
+                if _time_ns() >= deadline_ns:
+                    raise MetricsTimeoutError(
+                        "Timed out while executing callback"
+                    )
+
+                for measurement in measurements:
                     metric_reader_storage.consume_measurement(measurement)
+
         return self._reader_storages[metric_reader].collect()

opentelemetry-sdk/tests/metrics/test_measurement_consumer.py

@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from sys import version_info
+from time import sleep
 from unittest import TestCase
 from unittest.mock import MagicMock, Mock, patch
 
@@ -115,3 +117,74 @@ def test_collect_calls_async_instruments(self, MockMetricReaderStorage):
         self.assertEqual(
             len(reader_storage_mock.consume_measurement.mock_calls), 5
         )
+
+    def test_collect_timeout(self, MockMetricReaderStorage):
+        reader_mock = Mock()
+        reader_storage_mock = Mock()
+        MockMetricReaderStorage.return_value = reader_storage_mock
+        consumer = SynchronousMeasurementConsumer(
+            SdkConfiguration(
+                resource=Mock(),
+                metric_readers=[reader_mock],
+                views=Mock(),
+            )
+        )
+
+        def sleep_1(*args, **kwargs):
+            sleep(1)
+
+        consumer.register_asynchronous_instrument(
+            Mock(**{"callback.side_effect": sleep_1})
+        )
+
+        with self.assertRaises(Exception) as error:
+            consumer.collect(reader_mock, timeout_millis=10)
+
+        self.assertIn(
+            "Timed out while executing callback", error.exception.args[0]
+        )
+
+    @patch(
+        "opentelemetry.sdk.metrics._internal."
+        "measurement_consumer.CallbackOptions"
+    )
+    def test_collect_deadline(
+        self, mock_callback_options, MockMetricReaderStorage
+    ):
+        reader_mock = Mock()
+        reader_storage_mock = Mock()
+        MockMetricReaderStorage.return_value = reader_storage_mock
+        consumer = SynchronousMeasurementConsumer(
+            SdkConfiguration(
+                resource=Mock(),
+                metric_readers=[reader_mock],
+                views=Mock(),
+            )
+        )
+
+        def sleep_1(*args, **kwargs):
+            sleep(1)
+            return []
+
+        consumer.register_asynchronous_instrument(
+            Mock(**{"callback.side_effect": sleep_1})
+        )
+        consumer.register_asynchronous_instrument(
+            Mock(**{"callback.side_effect": sleep_1})
+        )
+
+        consumer.collect(reader_mock)
+
+        if version_info < (3, 8):
+            callback_options_time_call = mock_callback_options.mock_calls[-1][
+                2
+            ]["timeout_millis"]
+        else:
+            callback_options_time_call = mock_callback_options.mock_calls[
+                -1
+            ].kwargs["timeout_millis"]
+
+        self.assertLess(
+            callback_options_time_call,
+            10000 * 10**6,
+        )

opentelemetry-sdk/tests/metrics/test_metric_reader.py

@@ -14,6 +14,7 @@
 
 from typing import Dict, Iterable
 from unittest import TestCase
+from unittest.mock import patch
 
 from opentelemetry.sdk.metrics import Counter, Histogram, ObservableGauge
 from opentelemetry.sdk.metrics._internal.instrument import (
@@ -135,3 +136,9 @@ def test_configure_aggregation(self):
             dummy_metric_reader._instrument_class_aggregation[_Counter],
             LastValueAggregation,
         )
+
+    def test_force_flush(self):
+
+        with patch.object(DummyMetricReader, "collect") as mock_collect:
+            DummyMetricReader().force_flush(timeout_millis=10)
+            mock_collect.assert_called_with(timeout_millis=10)

opentelemetry-sdk/tests/metrics/test_metric_reader_storage.py

@@ -821,7 +821,7 @@ def test_view_instrument_match_conflict_8(self):
         # and also the temporality and monotonicity of the up down counter and
         # the histogram are the same.
 
-        observable_counter = _UpDownCounter(
+        up_down_counter = _UpDownCounter(
             "up_down_counter",
             Mock(),
             [Mock()],
@@ -859,7 +859,7 @@ def test_view_instrument_match_conflict_8(self):
         with self.assertRaises(AssertionError):
             with self.assertLogs(level=WARNING):
                 metric_reader_storage.consume_measurement(
-                    Measurement(1, observable_counter)
+                    Measurement(1, up_down_counter)
                 )
 
         with self.assertLogs(level=WARNING) as log:

opentelemetry-sdk/tests/metrics/test_periodic_exporting_metric_reader.py

@@ -118,9 +118,9 @@ def _create_periodic_reader(
             exporter, export_interval_millis=interval
         )
 
-        def _collect(reader):
+        def _collect(reader, timeout_millis):
             time.sleep(collect_wait)
-            pmr._receive_metrics(metrics)
+            pmr._receive_metrics(metrics, timeout_millis)
 
         pmr._set_collect_callback(_collect)
         return pmr