From 7e9778017e6b91fc4c99ba4378b4fab02f4ef795 Mon Sep 17 00:00:00 2001
From: Vickenty Fesunov <vickenty@users.noreply.github.com>
Date: Mon, 18 Mar 2024 14:24:14 +0100
Subject: [PATCH] Fix potential metric loss when open_buffer is combined with
 disable_buffering=False (#820)

* Don't reset buffer in open_buffer

When disable_buffering=False, the buffer may contain still unsent
metrics, which are lost when we reset the buffer. When
disable_buffering=True, the reset is a noop anyway.

* Remove invalid todos

Even if disable_buffering defaults to False, the user may still opt to
disable it, in which case the code in these methods still must execute
as it is today.
---
 datadog/dogstatsd/base.py                      |  4 ----
 .../dogstatsd/test_statsd_sender.py            | 18 ++++++++++++++++++
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/datadog/dogstatsd/base.py b/datadog/dogstatsd/base.py
index c04c0f450..821c41801 100644
--- a/datadog/dogstatsd/base.py
+++ b/datadog/dogstatsd/base.py
@@ -734,14 +734,11 @@ def open_buffer(self, max_buffer_size=None):
 
         self._config_lock.acquire()
 
-        # XXX Remove if `disable_buffering` default is changed to False
         self._send = self._send_to_buffer
 
         if max_buffer_size is not None:
             log.warning("The parameter max_buffer_size is now deprecated and is not used anymore")
 
-        self._reset_buffer()
-
     def close_buffer(self):
         """
         Flush the buffer and switch back to single metric packets.
@@ -752,7 +749,6 @@ def close_buffer(self):
         try:
             self.flush()
         finally:
-            # XXX Remove if `disable_buffering` default is changed to False
             if self._disable_buffering:
                 self._send = self._send_to_server
 
diff --git a/tests/integration/dogstatsd/test_statsd_sender.py b/tests/integration/dogstatsd/test_statsd_sender.py
index 222897643..d3a56860c 100644
--- a/tests/integration/dogstatsd/test_statsd_sender.py
+++ b/tests/integration/dogstatsd/test_statsd_sender.py
@@ -83,3 +83,21 @@ def test_fork_hooks(disable_background_sender, disable_buffering):
 
     foo.close()
     bar.close()
+
+
+def test_buffering_with_context():
+    statsd = DogStatsd(
+        telemetry_min_flush_interval=0,
+        disable_buffering=False,
+    )
+
+    foo, bar = socket.socketpair(socket.AF_UNIX, socket.SOCK_DGRAM, 0)
+    statsd.socket = foo
+
+    statsd.increment("first")
+    with statsd: # should not erase previously buffered metrics
+        pass
+
+    bar.settimeout(5)
+    msg = bar.recv(8192)
+    assert msg == b"first:1|c\n"