locks :)

DataDog · Jan 17, 2025 · 60143c6 · 60143c6
1 parent f13b48c
commit 60143c6
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 13 deletions.
diff --git a/datadog/dogstatsd/max_sample_metric.py b/datadog/dogstatsd/max_sample_metric.py
@@ -25,18 +25,17 @@ def sample(self, value):
         self.total_metric_samples += 1
 
     def maybe_keep_sample(self, value):
-        with self.lock:
-            if self.max_metric_samples > 0:
-                self.total_metric_samples += 1
-                if self.stored_metric_samples < self.max_metric_samples:
-                    self.data[self.stored_metric_samples] = value
-                    self.stored_metric_samples += 1
-                else:
-                    i = random.randint(0, self.total_metric_samples - 1)
-                    if i < self.max_metric_samples:
-                        self.data[i] = value
+        if self.max_metric_samples > 0:
+            self.total_metric_samples += 1
+            if self.stored_metric_samples < self.max_metric_samples:
+                self.data[self.stored_metric_samples] = value
+                self.stored_metric_samples += 1
             else:
-                self.sample(value)
+                i = random.randint(0, self.total_metric_samples - 1)
+                if i < self.max_metric_samples:
+                    self.data[i] = value
+        else:
+            self.sample(value)
 
     def skip_sample(self):
         self.total_metric_samples += 1

diff --git a/datadog/dogstatsd/max_sample_metric_context.py b/datadog/dogstatsd/max_sample_metric_context.py
@@ -11,9 +11,11 @@ def __init__(self, max_sample_metric_type):
     def flush(self):
         metrics = []
         """Flush the metrics and reset the stored values."""
-        for _, metric in self.values.items():
+        with self.lock:
+            temp = self.values
+            self.values = {}
+        for _, metric in temp.items():
             metrics.append(metric.flush())
-        self.values = {}
         return metrics
 
     def sample(self, name, value, tags, rate, context_key, max_samples_per_context):
@@ -24,10 +26,12 @@ def sample(self, name, value, tags, rate, context_key, max_samples_per_context):
                 # Create a new metric if it doesn't exist
                 self.values[context_key] = self.max_sample_metric_type(name, tags, rate, max_samples_per_context)
             metric = self.values[context_key]
+            metric.lock.acquire()
         if keeping_sample:
             metric.maybe_keep_sample(value)
         else:
             metric.skip_sample()
+        metric.lock.release()
 
     def should_sample(self, rate):
         """Determine if a sample should be kept based on the specified rate."""