Support statsbeat in EU regions (census-instrumentation#1105)

Author: lzchen

contrib/opencensus-ext-azure/CHANGELOG.md

@@ -4,6 +4,8 @@
 
 - Statsbeat bug fixes, shorten host in network stats
 ([#1100](https://github.com/census-instrumentation/opencensus-python/pull/1100))
+- Support statsbeat in EU regions
+([#1105](https://github.com/census-instrumentation/opencensus-python/pull/1105))
 
 ## 1.1.1
 Released 2022-01-19

contrib/opencensus-ext-azure/opencensus/ext/azure/metrics_exporter/statsbeat_metrics/__init__.py

@@ -16,25 +16,29 @@
 
 from opencensus.ext.azure.metrics_exporter import MetricsExporter
 from opencensus.ext.azure.metrics_exporter.statsbeat_metrics.statsbeat import (
-    _STATS_CONNECTION_STRING,
     _STATS_SHORT_EXPORT_INTERVAL,
+    _get_stats_connection_string,
     _StatsbeatMetrics,
 )
 from opencensus.metrics import transport
 from opencensus.metrics.export.metric_producer import MetricProducer
 
 _STATSBEAT_METRICS = None
+_STATSBEAT_EXPORTER = None
 _STATSBEAT_LOCK = threading.Lock()
 
 
 def collect_statsbeat_metrics(options):
-    with _STATSBEAT_LOCK:
-        # Only start statsbeat if did not exist before
-        global _STATSBEAT_METRICS  # pylint: disable=global-statement
-        if _STATSBEAT_METRICS is None:
+    # pylint: disable=global-statement
+    global _STATSBEAT_METRICS
+    global _STATSBEAT_EXPORTER
+    # Only start statsbeat if did not exist before
+    if _STATSBEAT_METRICS is None and _STATSBEAT_EXPORTER is None:
+        with _STATSBEAT_LOCK:
+            # Only start statsbeat if did not exist before
             exporter = MetricsExporter(
                 is_stats=True,
-                connection_string=_STATS_CONNECTION_STRING,
+                connection_string=_get_stats_connection_string(options.endpoint),  # noqa: E501
                 enable_standard_metrics=False,
                 export_interval=_STATS_SHORT_EXPORT_INTERVAL,  # 15m by default
             )
@@ -47,6 +51,7 @@ def collect_statsbeat_metrics(options):
                 transport.get_exporter_thread([_STATSBEAT_METRICS],
                                               exporter,
                                               exporter.options.export_interval)
+            _STATSBEAT_EXPORTER = exporter
 
 
 class _AzureStatsbeatMetricsProducer(MetricProducer):

contrib/opencensus-ext-azure/opencensus/ext/azure/metrics_exporter/statsbeat_metrics/statsbeat.py

@@ -37,9 +37,22 @@
 _AIMS_API_VERSION = "api-version=2017-12-01"
 _AIMS_FORMAT = "format=json"
 
-_DEFAULT_STATS_CONNECTION_STRING = "InstrumentationKey=c4a29126-a7cb-47e5-b348-11414998b11e;IngestionEndpoint=https://westus-0.in.applicationinsights.azure.com/"  # noqa: E501
+_DEFAULT_NON_EU_STATS_CONNECTION_STRING = "InstrumentationKey=c4a29126-a7cb-47e5-b348-11414998b11e;IngestionEndpoint=https://westus-0.in.applicationinsights.azure.com/"  # noqa: E501
+_DEFAULT_EU_STATS_CONNECTION_STRING = "InstrumentationKey=7dc56bab-3c0c-4e9f-9ebb-d1acadee8d0f;IngestionEndpoint=https://westeurope-5.in.applicationinsights.azure.com/"  # noqa: E501
 _DEFAULT_STATS_SHORT_EXPORT_INTERVAL = 900  # 15 minutes
 _DEFAULT_STATS_LONG_EXPORT_INTERVAL = 86400  # 24 hours
+_EU_ENDPOINTS = [
+    "westeurope",
+    "northeurope",
+    "francecentral",
+    "francesouth",
+    "germanywestcentral",
+    "norwayeast",
+    "norwaywest",
+    "swedencentral",
+    "switzerlandnorth",
+    "switzerlandwest",
+]
 
 _ATTACH_METRIC_NAME = "Attach"
 _FEATURE_METRIC_NAME = "Feature"
@@ -69,12 +82,16 @@ class _StatsbeatFeature:
     AAD = 2
 
 
-def _get_stats_connection_string():
+def _get_stats_connection_string(endpoint):
     cs_env = os.environ.get("APPLICATION_INSIGHTS_STATS_CONNECTION_STRING")
     if cs_env:
         return cs_env
     else:
-        return _DEFAULT_STATS_CONNECTION_STRING
+        for ep in _EU_ENDPOINTS:
+            if ep in endpoint:
+                # Use statsbeat EU endpoint if user is in EU region
+                return _DEFAULT_EU_STATS_CONNECTION_STRING
+        return _DEFAULT_NON_EU_STATS_CONNECTION_STRING
 
 
 def _get_stats_short_export_interval():
@@ -93,7 +110,6 @@ def _get_stats_long_export_interval():
         return _DEFAULT_STATS_LONG_EXPORT_INTERVAL
 
 
-_STATS_CONNECTION_STRING = _get_stats_connection_string()
 _STATS_SHORT_EXPORT_INTERVAL = _get_stats_short_export_interval()
 _STATS_LONG_EXPORT_INTERVAL = _get_stats_long_export_interval()
 _STATS_LONG_INTERVAL_THRESHOLD = _STATS_LONG_EXPORT_INTERVAL / _STATS_SHORT_EXPORT_INTERVAL  # noqa: E501

contrib/opencensus-ext-azure/tests/test_azure_statsbeat_metrics.py

@@ -23,8 +23,13 @@
 from opencensus.ext.azure.common import Options
 from opencensus.ext.azure.common.transport import _requests_map
 from opencensus.ext.azure.common.version import __version__ as ext_version
-from opencensus.ext.azure.metrics_exporter import statsbeat_metrics
+from opencensus.ext.azure.metrics_exporter import (
+    MetricsExporter,
+    statsbeat_metrics,
+)
 from opencensus.ext.azure.metrics_exporter.statsbeat_metrics.statsbeat import (
+    _DEFAULT_EU_STATS_CONNECTION_STRING,
+    _DEFAULT_NON_EU_STATS_CONNECTION_STRING,
     _ENDPOINT_TYPES,
     _FEATURE_TYPES,
     _RP_NAMES,
@@ -37,6 +42,7 @@
     _get_feature_properties,
     _get_network_properties,
     _get_retry_count_value,
+    _get_stats_connection_string,
     _get_success_count_value,
     _get_throttle_count_value,
     _shorten_host,
@@ -78,6 +84,7 @@ class TestStatsbeatMetrics(unittest.TestCase):
     def setUp(self):
         # pylint: disable=protected-access
         statsbeat_metrics._STATSBEAT_METRICS = None
+        statsbeat_metrics._STATSBEAT_EXPORTER = None
 
     def test_producer_ctor(self):
         # pylint: disable=protected-access
@@ -121,6 +128,12 @@ def test_collect_statsbeat_metrics(self, thread_mock, stats_mock):
                 statsbeat_metrics._AzureStatsbeatMetricsProducer
             )
         )
+        self.assertTrue(
+            isinstance(
+                statsbeat_metrics._STATSBEAT_EXPORTER,
+                MetricsExporter,
+            )
+        )
         self.assertEqual(
             statsbeat_metrics._STATSBEAT_METRICS._statsbeat._instrumentation_key, "ikey")  # noqa: E501
         thread_mock.assert_called_once()
@@ -137,6 +150,74 @@ def test_collect_statsbeat_metrics_exists(self, thread_mock, stats_mock):
         thread_mock.assert_not_called()
         stats_mock.assert_not_called()
 
+    @mock.patch.object(_StatsbeatMetrics, 'get_initial_metrics')
+    @mock.patch('opencensus.metrics.transport.get_exporter_thread')
+    def test_collect_statsbeat_metrics_non_eu(self, thread_mock, stats_mock):
+        # pylint: disable=protected-access
+        cs = "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://westus-0.in.applicationinsights.azure.com/"  # noqa: E501
+        non_eu = Options(
+            connection_string=cs
+        )
+        with mock.patch.dict(
+                os.environ, {
+                    "APPLICATION_INSIGHTS_STATS_CONNECTION_STRING": "",
+                }):
+            statsbeat_metrics.collect_statsbeat_metrics(non_eu)
+            self.assertTrue(
+                isinstance(
+                    statsbeat_metrics._STATSBEAT_METRICS,
+                    statsbeat_metrics._AzureStatsbeatMetricsProducer
+                )
+            )
+            self.assertTrue(
+                isinstance(
+                    statsbeat_metrics._STATSBEAT_EXPORTER,
+                    MetricsExporter,
+                )
+            )
+            self.assertEqual(
+                statsbeat_metrics._STATSBEAT_EXPORTER.options.instrumentation_key,  # noqa: E501
+                _DEFAULT_NON_EU_STATS_CONNECTION_STRING.split(";")[0].split("=")[1]   # noqa: E501
+            )
+            self.assertEqual(
+                statsbeat_metrics._STATSBEAT_EXPORTER.options.endpoint,
+                _DEFAULT_NON_EU_STATS_CONNECTION_STRING.split(";")[1].split("=")[1]   # noqa: E501
+            )
+
+    @mock.patch.object(_StatsbeatMetrics, 'get_initial_metrics')
+    @mock.patch('opencensus.metrics.transport.get_exporter_thread')
+    def test_collect_statsbeat_metrics_eu(self, thread_mock, stats_mock):
+        # pylint: disable=protected-access
+        cs = "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://northeurope-0.in.applicationinsights.azure.com/"  # noqa: E501
+        eu = Options(
+            connection_string=cs
+        )
+        with mock.patch.dict(
+                os.environ, {
+                    "APPLICATION_INSIGHTS_STATS_CONNECTION_STRING": "",
+                }):
+            statsbeat_metrics.collect_statsbeat_metrics(eu)
+            self.assertTrue(
+                isinstance(
+                    statsbeat_metrics._STATSBEAT_METRICS,
+                    statsbeat_metrics._AzureStatsbeatMetricsProducer
+                )
+            )
+            self.assertTrue(
+                isinstance(
+                    statsbeat_metrics._STATSBEAT_EXPORTER,
+                    MetricsExporter,
+                )
+            )
+            self.assertEqual(
+                statsbeat_metrics._STATSBEAT_EXPORTER.options.instrumentation_key,  # noqa: E501
+                _DEFAULT_EU_STATS_CONNECTION_STRING.split(";")[0].split("=")[1]   # noqa: E501
+            )
+            self.assertEqual(
+                statsbeat_metrics._STATSBEAT_EXPORTER.options.endpoint,
+                _DEFAULT_EU_STATS_CONNECTION_STRING.split(";")[1].split("=")[1]   # noqa: E501
+            )
+
     @mock.patch(
         'opencensus.ext.azure.metrics_exporter.statsbeat_metrics.statsbeat._get_feature_properties')  # noqa: E501
     @mock.patch(
@@ -149,6 +230,7 @@ def test_statsbeat_metric_init(self, attach_mock, network_mock, feature_mock):
         self.assertEqual(len(metric._vm_data), 0)
         self.assertTrue(metric._vm_retry)
         self.assertEqual(metric._instrumentation_key, "ikey")
+        self.assertEqual(metric._feature, 1)
         self.assertTrue(
             isinstance(
                 metric._attach_metric,
@@ -327,7 +409,7 @@ def test_get_feature_metric(self):
         self.assertEqual(
             properties[8].value, ext_version)  # noqa: E501
 
-    def test_get_feature_metric_wtih_aad(self):
+    def test_get_feature_metric_with_aad(self):
         aad_options = Options(
             instrumentation_key="ikey",
             enable_local_storage=True,
@@ -584,7 +666,7 @@ def test_get_azure_compute_metadata_not_vm_timeout(self):
             self.assertEqual(len(stats._vm_data), 0)
             self.assertFalse(stats._vm_retry)
 
-    def test_get_azure_compute_metadata__vm_retry(self):
+    def test_get_azure_compute_metadata_vm_retry(self):
         with mock.patch(
             'requests.get',
             throw(requests.exceptions.RequestException)
@@ -612,3 +694,39 @@ def test_shorten_host(self):
         self.assertEqual(_shorten_host(url), "fakehost")
         url = "http://fakehost-5/"
         self.assertEqual(_shorten_host(url), "fakehost-5")
+
+    def test_get_stats_connection_string_env(self):
+        cs = "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://westus-0.in.applicationinsights.azure.com/"  # noqa: E501
+        with mock.patch.dict(
+            os.environ, {
+                "APPLICATION_INSIGHTS_STATS_CONNECTION_STRING": cs
+            }
+        ):
+            stats_cs = _get_stats_connection_string(_OPTIONS.endpoint)
+            self.assertEqual(stats_cs, cs)
+
+    def test_get_stats_connection_string_non_eu(self):
+        with mock.patch.dict(
+            os.environ, {
+                "APPLICATION_INSIGHTS_STATS_CONNECTION_STRING": ""
+            }
+        ):
+            cs = "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://westus-0.in.applicationinsights.azure.com/"  # noqa: E501
+            non_eu = Options(
+                connection_string=cs,
+            )
+            stats_cs = _get_stats_connection_string(non_eu.endpoint)
+            self.assertEqual(stats_cs, _DEFAULT_NON_EU_STATS_CONNECTION_STRING)
+
+    def test_get_stats_connection_string_eu(self):
+        with mock.patch.dict(
+            os.environ, {
+                "APPLICATION_INSIGHTS_STATS_CONNECTION_STRING": ""
+            }
+        ):
+            cs = "InstrumentationKey=1aa11111-bbbb-1ccc-8ddd-eeeeffff3333;IngestionEndpoint=https://northeurope-0.in.applicationinsights.azure.com/"  # noqa: E501
+            eu = Options(
+                connection_string=cs,
+            )
+            stats_cs = _get_stats_connection_string(eu.endpoint)
+            self.assertEqual(stats_cs, _DEFAULT_EU_STATS_CONNECTION_STRING)

tox.ini

@@ -18,6 +18,9 @@ deps =
   unit,lint: pytest-cov
   unit,lint: retrying
   unit,lint: unittest2
+  py27-unit: markupsafe==1.1.1
+  py35-unit: markupsafe==1.1.1
+  py3{6,7,8,9}-unit: markupsafe==2.0.1  #  https://github.com/pallets/markupsafe/issues/282
   bandit: bandit
   unit,lint,setup,docs,bandit: -e context/opencensus-context
   unit,lint,docs,bandit: -e contrib/opencensus-correlation