Fix issue with stats d method being called

datadog_lambda/metric.py

@@ -21,9 +21,20 @@
 logger = logging.getLogger(__name__)
 
 
-class StatsDWrapper:
+class StatsWriter:
+    def distribution(self, metric_name, value, tags=[], timestamp=None):
+        pass
+
+    def flush(self):
+        pass
+
+    def stop(self):
+        pass
+
+
+class StatsDWriter(StatsWriter):
     """
-    Wraps StatsD calls, to give an identical interface to ThreadStats
+    Wraps StatsD calls, creates a common
     """
 
     def __init__(self):
@@ -33,21 +44,76 @@ def __init__(self):
     def distribution(self, metric_name, value, tags=[], timestamp=None):
         statsd.distribution(metric_name, value, tags=tags)
 
-    def flush(self, value):
-        pass
+
+class ThreadStatsWriter(StatsWriter):
+    """
+    Writes distribution metrics using thre ThreadStats class
+    """
+
+    def __init__(self, flush_in_thread):
+        self.thread_stats = ThreadStats(compress_payload=True)
+        self.thread_stats.start(flush_in_thread=flush_in_thread)
+
+    def distribution(self, metric_name, value, tags=[], timestamp=None):
+        self.thread_stats.distribution(
+            metric_name, value, tags=tags, timestamp=timestamp
+        )
+
+    def flush(self):
+        """ "Flush distributions from ThreadStats to Datadog.
+        Modified based on `datadog.threadstats.base.ThreadStats.flush()`,
+        to gain better control over exception handling.
+        """
+        _, dists = self.thread_stats._get_aggregate_metrics_and_dists(float("inf"))
+        count_dists = len(dists)
+        if not count_dists:
+            logger.debug("No distributions to flush. Continuing.")
+
+        self.thread_stats.flush_count += 1
+        logger.debug(
+            "Flush #%s sending %s distributions",
+            self.thread_stats.flush_count,
+            count_dists,
+        )
+        try:
+            self.thread_stats.reporter.flush_distributions(dists)
+        except Exception as e:
+            # The nature of the root issue https://bugs.python.org/issue41345 is complex,
+            # but comprehensive tests suggest that it is safe to retry on this specific error.
+            if isinstance(
+                e, api.exceptions.ClientError
+            ) and "RemoteDisconnected" in str(e):
+                logger.debug(
+                    "Retry flush #%s due to RemoteDisconnected",
+                    self.thread_stats.flush_count,
+                )
+                try:
+                    self.thread_stats.reporter.flush_distributions(dists)
+                except Exception:
+                    logger.debug(
+                        "Flush #%s failed after retry",
+                        self.thread_stats.flush_count,
+                        exc_info=True,
+                    )
+            else:
+                logger.debug(
+                    "Flush #%s failed", self.thread_stats.flush_count, exc_info=True
+                )
+
+    def stop(self):
+        self.thread_stats.stop()
 
 
 lambda_stats = None
 if should_use_extension:
-    lambda_stats = StatsDWrapper()
+    lambda_stats = StatsDWriter()
 else:
     # Periodical flushing in a background thread is NOT guaranteed to succeed
     # and leads to data loss. When disabled, metrics are only flushed at the
     # end of invocation. To make metrics submitted from a long-running Lambda
     # function available sooner, consider using the Datadog Lambda extension.
     flush_in_thread = os.environ.get("DD_FLUSH_IN_THREAD", "").lower() == "true"
-    lambda_stats = ThreadStats(compress_payload=True)
-    lambda_stats.start(flush_in_thread=flush_in_thread)
+    lambda_stats = ThreadStatsWriter(flush_in_thread)
 
 
 def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=False):
@@ -74,8 +140,7 @@ def lambda_metric(metric_name, value, timestamp=None, tags=None, force_async=Fal
 
 
 def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
-    """Writes the specified metric point to standard output
-    """
+    """Writes the specified metric point to standard output"""
     logger.debug(
         "Sending metric %s value %s to Datadog via log forwarder", metric_name, value
     )
@@ -91,40 +156,8 @@ def write_metric_point_to_stdout(metric_name, value, timestamp=None, tags=[]):
     )
 
 
-def flush_thread_stats():
-    """"Flush distributions from ThreadStats to Datadog.
-
-    Modified based on `datadog.threadstats.base.ThreadStats.flush()`,
-    to gain better control over exception handling.
-    """
-    _, dists = lambda_stats._get_aggregate_metrics_and_dists(float("inf"))
-    count_dists = len(dists)
-    if not count_dists:
-        logger.debug("No distributions to flush. Continuing.")
-
-    lambda_stats.flush_count += 1
-    logger.debug(
-        "Flush #%s sending %s distributions", lambda_stats.flush_count, count_dists
-    )
-    try:
-        lambda_stats.reporter.flush_distributions(dists)
-    except Exception as e:
-        # The nature of the root issue https://bugs.python.org/issue41345 is complex,
-        # but comprehensive tests suggest that it is safe to retry on this specific error.
-        if isinstance(e, api.exceptions.ClientError) and "RemoteDisconnected" in str(e):
-            logger.debug(
-                "Retry flush #%s due to RemoteDisconnected", lambda_stats.flush_count
-            )
-            try:
-                lambda_stats.reporter.flush_distributions(dists)
-            except Exception:
-                logger.debug(
-                    "Flush #%s failed after retry",
-                    lambda_stats.flush_count,
-                    exc_info=True,
-                )
-        else:
-            logger.debug("Flush #%s failed", lambda_stats.flush_count, exc_info=True)
+def flush_stats():
+    lambda_stats.flush()
 
 
 def are_enhanced_metrics_enabled():

datadog_lambda/wrapper.py

@@ -12,7 +12,7 @@
 from datadog_lambda.cold_start import set_cold_start, is_cold_start
 from datadog_lambda.constants import XraySubsegment, TraceContextSource
 from datadog_lambda.metric import (
-    flush_thread_stats,
+    flush_stats,
     submit_invocations_metric,
     submit_errors_metric,
 )
@@ -177,7 +177,7 @@ def _after(self, event, context):
                 )
 
             if not self.flush_to_log or should_use_extension:
-                flush_thread_stats()
+                flush_stats()
             if should_use_extension:
                 flush_extension()
 

tests/test_metric.py

@@ -7,7 +7,7 @@
     from mock import patch, call
 
 from datadog.api.exceptions import ClientError
-from datadog_lambda.metric import lambda_metric, flush_thread_stats
+from datadog_lambda.metric import lambda_metric, ThreadStatsWriter
 from datadog_lambda.tags import _format_dd_lambda_layer_tag
 
 
@@ -49,10 +49,12 @@ def setUp(self):
 
     def test_retry_on_remote_disconnected(self):
         # Raise the RemoteDisconnected error
+        lambda_stats = ThreadStatsWriter(True)
+
         self.mock_threadstats_flush_distributions.side_effect = ClientError(
             "POST",
             "https://api.datadoghq.com/api/v1/distribution_points",
             "RemoteDisconnected('Remote end closed connection without response')",
         )
-        flush_thread_stats()
+        lambda_stats.flush()
         self.assertEqual(self.mock_threadstats_flush_distributions.call_count, 2)

tests/test_wrapper.py

@@ -7,7 +7,7 @@
     from mock import patch, call, ANY, MagicMock
 
 from datadog_lambda.wrapper import datadog_lambda_wrapper
-from datadog_lambda.metric import lambda_metric
+from datadog_lambda.metric import lambda_metric, ThreadStatsWriter
 
 
 def get_mock_context(
@@ -136,7 +136,7 @@ def test_datadog_lambda_wrapper_flush_in_thread(self):
         import datadog_lambda.metric as metric_module
 
         metric_module.lambda_stats.stop()
-        metric_module.lambda_stats.start(flush_in_thread=True)
+        metric_module.lambda_stats = ThreadStatsWriter(True)
 
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -156,14 +156,14 @@ def lambda_handler(event, context):
 
         # reset ThreadStats
         metric_module.lambda_stats.stop()
-        metric_module.lambda_stats.start(flush_in_thread=False)
+        metric_module.lambda_stats = ThreadStatsWriter(False)
 
     def test_datadog_lambda_wrapper_not_flush_in_thread(self):
         # force ThreadStats to not flush in thread
         import datadog_lambda.metric as metric_module
 
         metric_module.lambda_stats.stop()
-        metric_module.lambda_stats.start(flush_in_thread=False)
+        metric_module.lambda_stats = ThreadStatsWriter(False)
 
         @datadog_lambda_wrapper
         def lambda_handler(event, context):
@@ -183,7 +183,7 @@ def lambda_handler(event, context):
 
         # reset ThreadStats
         metric_module.lambda_stats.stop()
-        metric_module.lambda_stats.start(flush_in_thread=False)
+        metric_module.lambda_stats = ThreadStatsWriter(False)
 
     def test_datadog_lambda_wrapper_inject_correlation_ids(self):
         os.environ["DD_LOGS_INJECTION"] = "True"