fix(internal): add locks for thread-safety of trace encoder (#3022) (#3024)

* fix(internal): add buffered encoder test for thread safety

* fix(internal): wrap encoder with reentrant lock

* fix flake8

* add release note

* add missing locks

Co-authored-by: brettlangdon <[email protected]>
(cherry picked from commit dbcc631)

Co-authored-by: Tahir H. Butt <[email protected]>

Author: mergify[bot]

ddtrace/internal/_encoding.pyx

@@ -208,7 +208,7 @@ cdef class MsgpackStringTable(StringTable):
         self.max_size = max_size
         self.pk.length = MSGPACK_STRING_TABLE_LENGTH_PREFIX_SIZE
         self._sp_len = 0
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         super(MsgpackStringTable, self).__init__()
 
         assert self.index(ORIGIN_KEY) == 1
@@ -264,7 +264,8 @@ cdef class MsgpackStringTable(StringTable):
 
     @property
     def size(self):
-        return self.pk.length - MSGPACK_ARRAY_LENGTH_PREFIX_SIZE + array_prefix_size(self._next_id)
+        with self._lock:
+            return self.pk.length - MSGPACK_ARRAY_LENGTH_PREFIX_SIZE + array_prefix_size(self._next_id)
 
     cdef append_raw(self, long src, Py_ssize_t size):
         cdef int res
@@ -284,10 +285,11 @@ cdef class MsgpackStringTable(StringTable):
         self._sp_len = 0
 
     cpdef flush(self):
-        try:
-            return self.get_bytes()
-        finally:
-            self.reset()
+        with self._lock:
+            try:
+                return self.get_bytes()
+            finally:
+                self.reset()
 
 
 cdef class BufferedEncoder(object):
@@ -372,7 +374,7 @@ cdef class MsgpackEncoderBase(BufferedEncoder):
         self.max_size = max_size
         self.pk.buf_size = buf_size
         self.max_item_size = max_item_size if max_item_size < max_size else max_size
-        self._lock = threading.Lock()
+        self._lock = threading.RLock()
         self._reset_buffer()
 
     def __dealloc__(self):
@@ -393,10 +395,11 @@ cdef class MsgpackEncoderBase(BufferedEncoder):
         self.pk.length = MSGPACK_ARRAY_LENGTH_PREFIX_SIZE  # Leave room for array length prefix
 
     cpdef encode(self):
-        if not self._count:
-            return None
+        with self._lock:
+            if not self._count:
+                return None
 
-        return self.flush()
+            return self.flush()
 
     cdef inline int _update_array_len(self):
         """Update traces array size prefix"""
@@ -437,45 +440,46 @@ cdef class MsgpackEncoderBase(BufferedEncoder):
         if L > 0 and trace[0].context is not None and trace[0].context.dd_origin is not None:
             dd_origin = self.get_dd_origin_ref(trace[0].context.dd_origin)
 
-        with self._lock:
-            for span in trace:
-                ret = self.pack_span(span, dd_origin)
-                if ret != 0: raise RuntimeError("Couldn't pack span")
+        for span in trace:
+            ret = self.pack_span(span, dd_origin)
+            if ret != 0: raise RuntimeError("Couldn't pack span")
 
         return ret
 
     cpdef put(self, list trace):
         """Put a trace (i.e. a list of spans) in the buffer."""
         cdef int ret
 
-        len_before = self.pk.length
-        size_before = self.size
-        try:
-            ret = self._pack_trace(trace)
-            if ret:  # should not happen.
-                raise RuntimeError("internal error")
-
-            # DEV: msgpack avoids buffer overflows by calling PyMem_Realloc so
-            # we must check sizes manually.
-            # TODO: We should probably ensure that the buffer size doesn't
-            # grow arbitrarily because of the PyMem_Realloc and if it does then
-            # free and reallocate with the appropriate size.
-            if self.size - size_before > self.max_item_size:
-                raise BufferItemTooLarge(self.size - size_before)
-
-            if self.size > self.max_size:
-                raise BufferFull(self.size - size_before)
-
-            self._count += 1
-        except:
-            # rollback
-            self.pk.length = len_before
-            raise
+        with self._lock:
+            len_before = self.pk.length
+            size_before = self.size
+            try:
+                ret = self._pack_trace(trace)
+                if ret:  # should not happen.
+                    raise RuntimeError("internal error")
+
+                # DEV: msgpack avoids buffer overflows by calling PyMem_Realloc so
+                # we must check sizes manually.
+                # TODO: We should probably ensure that the buffer size doesn't
+                # grow arbitrarily because of the PyMem_Realloc and if it does then
+                # free and reallocate with the appropriate size.
+                if self.size - size_before > self.max_item_size:
+                    raise BufferItemTooLarge(self.size - size_before)
+
+                if self.size > self.max_size:
+                    raise BufferFull(self.size - size_before)
+
+                self._count += 1
+            except:
+                # rollback
+                self.pk.length = len_before
+                raise
 
     @property
     def size(self):
         """Return the size in bytes of the encoder buffer."""
-        return self.pk.length + array_prefix_size(self._count) - MSGPACK_ARRAY_LENGTH_PREFIX_SIZE
+        with self._lock:
+            return self.pk.length + array_prefix_size(self._count) - MSGPACK_ARRAY_LENGTH_PREFIX_SIZE
 
     # ---- Abstract methods ----
 
@@ -488,10 +492,11 @@ cdef class MsgpackEncoderBase(BufferedEncoder):
 
 cdef class MsgpackEncoderV03(MsgpackEncoderBase):
     cpdef flush(self):
-        try:
-            return self.get_bytes()
-        finally:
-            self._reset_buffer()
+        with self._lock:
+            try:
+                return self.get_bytes()
+            finally:
+                self._reset_buffer()
 
     cdef void * get_dd_origin_ref(self, str dd_origin):
         return string_to_buff(dd_origin)
@@ -637,24 +642,27 @@ cdef class MsgpackEncoderV05(MsgpackEncoderBase):
         self._st = MsgpackStringTable(max_size)
 
     cpdef flush(self):
-        try:
-            self._st.append_raw(PyLong_FromLong(<long> self.get_buffer()), <Py_ssize_t> super(MsgpackEncoderV05, self).size)
-            return self._st.flush()
-        finally:
-            self._reset_buffer()
+        with self._lock:
+            try:
+                self._st.append_raw(PyLong_FromLong(<long> self.get_buffer()), <Py_ssize_t> super(MsgpackEncoderV05, self).size)
+                return self._st.flush()
+            finally:
+                self._reset_buffer()
 
     @property
     def size(self):
         """Return the size in bytes of the encoder buffer."""
-        return self._st.size + super(MsgpackEncoderV05, self).size
+        with self._lock:
+            return self._st.size + super(MsgpackEncoderV05, self).size
 
     cpdef put(self, list trace):
-        try:
-            self._st.savepoint()
-            super(MsgpackEncoderV05, self).put(trace)
-        except Exception:
-            self._st.rollback()
-            raise
+        with self._lock:
+            try:
+                self._st.savepoint()
+                super(MsgpackEncoderV05, self).put(trace)
+            except Exception:
+                self._st.rollback()
+                raise
 
     cdef inline int _pack_string(self, object string):
         return msgpack_pack_uint32(&self.pk, self._st._index(string))

releasenotes/notes/fix-encoder-thread-safe-5c6924013c51d4f6.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    The thread safety of the custom buffered encoder was fixed in order to eliminate a potential cause of decoding errors of trace payloads (missing trace data) in the agent.

tests/tracer/test_encoders.py

@@ -2,6 +2,7 @@
 import json
 import random
 import string
+import threading
 from unittest import TestCase
 
 from hypothesis import given
@@ -595,3 +596,35 @@ def test_encoding_invalid_data(data):
         encoder.put(trace)
 
     assert encoder.encode() is None
+
+
+@allencodings
+def test_custom_msgpack_encode_thread_safe(encoding):
+    class TracingThread(threading.Thread):
+        def __init__(self, encoder, span_count, trace_count):
+            super(TracingThread, self).__init__()
+            trace = [
+                Span(tracer=None, name="span-{}-{}".format(self.name, _), service="threads", resource="TEST")
+                for _ in range(span_count)
+            ]
+            self._encoder = encoder
+            self._trace = trace
+            self._trace_count = trace_count
+
+        def run(self):
+            for _ in range(self._trace_count):
+                self._encoder.put(self._trace)
+
+    THREADS = 40
+    SPANS = 15
+    TRACES = 10
+    encoder = MSGPACK_ENCODERS[encoding](2 << 20, 2 << 20)
+
+    ts = [TracingThread(encoder, random.randint(1, SPANS), random.randint(1, TRACES)) for _ in range(THREADS)]
+    for t in ts:
+        t.start()
+    for t in ts:
+        t.join()
+
+    unpacked = decode(encoder.encode(), reconstruct=True)
+    assert unpacked is not None