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):
+        raise NotImplementedError()
+
+    def flush(self):
+        raise NotImplementedError()
+
+    def stop(self):
+        raise NotImplementedError()
+
+
+class StatsDWriter(StatsWriter):
     """
-    Wraps StatsD calls, to give an identical interface to ThreadStats
+    Writes distribution metrics using StatsD protocol
     """
 
     def __init__(self):
@@ -33,21 +44,82 @@ def __init__(self):
     def distribution(self, metric_name, value, tags=[], timestamp=None):
         statsd.distribution(metric_name, value, tags=tags)
 
-    def flush(self, value):
+    def flush(self):
+        pass
+
+    def stop(self):
         pass
 
 
+class ThreadStatsWriter(StatsWriter):
+    """
+    Writes distribution metrics using the 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 +146,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 +162,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()
 

scripts/run_integration_tests.sh

@@ -197,6 +197,7 @@ for handler_name in "${LAMBDA_HANDLERS[@]}"; do
                 # Normalize package version so that these snapshots aren't broken on version bumps
                 sed -E "s/(dd_lambda_layer:datadog-python[0-9]+_)[0-9]+\.[0-9]+\.[0-9]+/\1X\.X\.X/g" |
                 sed -E "s/(datadog_lambda:v)([0-9]+\.[0-9]+\.[0-9])/\1XX/g" |
+                sed -E "s/(python )([0-9]\.[0-9]+\.[0-9]+)/\1XX/g" |
                 # Strip out run ID (from function name, resource, etc.)
                 sed -E "s/${!run_id}/XXXX/g" |
                 # Strip out trace/span/parent/timestamps