Fix lossiness when submitting timestamped custom metrics (#527)

datadog_lambda/thread_stats_writer.py

@@ -27,8 +27,11 @@ def flush(self, tags=None):
         Modified based on `datadog.threadstats.base.ThreadStats.flush()`,
         to gain better control over exception handling.
         """
+        original_constant_tags = self.thread_stats.constant_tags.copy()
         if tags:
-            self.thread_stats.constant_tags = self.thread_stats.constant_tags + tags
+            # Temporarily add tags for this flush
+            self.thread_stats.constant_tags = original_constant_tags + tags
+
         _, dists = self.thread_stats._get_aggregate_metrics_and_dists(float("inf"))
         count_dists = len(dists)
         if not count_dists:
@@ -62,6 +65,9 @@ def flush(self, tags=None):
                 logger.debug(
                     "Flush #%s failed", self.thread_stats.flush_count, exc_info=True
                 )
+        finally:
+            # Reset constant_tags to its original state
+            self.thread_stats.constant_tags = original_constant_tags
 
     def stop(self):
         self.thread_stats.stop()

tests/test_metric.py

@@ -119,13 +119,69 @@ def test_retry_on_remote_disconnected(self):
 
     def test_flush_stats_with_tags(self):
         lambda_stats = ThreadStatsWriter(True)
+        original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()
         tags = ["tag1:value1", "tag2:value2"]
-        lambda_stats.flush(tags)
-        self.mock_threadstats_flush_distributions.assert_called_once_with(
-            lambda_stats.thread_stats._get_aggregate_metrics_and_dists(float("inf"))[1]
-        )
-        for tag in tags:
-            self.assertTrue(tag in lambda_stats.thread_stats.constant_tags)
+
+        # Add a metric to be flushed
+        lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])
+
+        with patch.object(
+            lambda_stats.thread_stats.reporter, "flush_distributions"
+        ) as mock_flush_distributions:
+            lambda_stats.flush(tags)
+            mock_flush_distributions.assert_called_once()
+            # Verify that after flush, constant_tags is reset to original
+            self.assertEqual(
+                lambda_stats.thread_stats.constant_tags, original_constant_tags
+            )
+
+    def test_flush_temp_constant_tags(self):
+        lambda_stats = ThreadStatsWriter(flush_in_thread=True)
+        lambda_stats.thread_stats.constant_tags = ["initial:tag"]
+        original_constant_tags = lambda_stats.thread_stats.constant_tags.copy()
+
+        lambda_stats.distribution("test.metric", 1, tags=["metric:tag"])
+        flush_tags = ["flush:tag1", "flush:tag2"]
+
+        with patch.object(
+            lambda_stats.thread_stats.reporter, "flush_distributions"
+        ) as mock_flush_distributions:
+            lambda_stats.flush(tags=flush_tags)
+            mock_flush_distributions.assert_called_once()
+            flushed_dists = mock_flush_distributions.call_args[0][0]
+
+            # Expected tags: original constant_tags + flush_tags + metric tags
+            expected_tags = original_constant_tags + flush_tags + ["metric:tag"]
+
+            # Verify the tags on the metric
+            self.assertEqual(len(flushed_dists), 1)
+            metric = flushed_dists[0]
+            self.assertEqual(sorted(metric["tags"]), sorted(expected_tags))
+
+            # Verify that constant_tags is reset after flush
+            self.assertEqual(
+                lambda_stats.thread_stats.constant_tags, original_constant_tags
+            )
+
+        # Repeat to ensure tags do not accumulate over multiple flushes
+        new_flush_tags = ["flush:tag3"]
+        lambda_stats.distribution("test.metric2", 2, tags=["metric2:tag"])
+
+        with patch.object(
+            lambda_stats.thread_stats.reporter, "flush_distributions"
+        ) as mock_flush_distributions:
+            lambda_stats.flush(tags=new_flush_tags)
+            mock_flush_distributions.assert_called_once()
+            flushed_dists = mock_flush_distributions.call_args[0][0]
+            # Expected tags for the new metric
+            expected_tags = original_constant_tags + new_flush_tags + ["metric2:tag"]
+
+            self.assertEqual(len(flushed_dists), 1)
+            metric = flushed_dists[0]
+            self.assertEqual(sorted(metric["tags"]), sorted(expected_tags))
+            self.assertEqual(
+                lambda_stats.thread_stats.constant_tags, original_constant_tags
+            )
 
     def test_flush_stats_without_context(self):
         flush_stats(lambda_context=None)