[core] Add Payload class helper (#834)

brettlangdon · web-flow · commit af1706360fc7 · 2019-03-12T10:59:17.000-04:00
* [core] Add Payload class

* update encoder tests

* Add payload tests

* use byte string for python 3.x
diff --git a/ddtrace/api.py b/ddtrace/api.py
@@ -7,6 +7,7 @@
 from .encoding import get_encoder, JSONEncoder
 from .compat import httplib, PYTHON_VERSION, PYTHON_INTERPRETER, get_connection_response
 from .internal.logger import get_logger
+from .payload import Payload
 from .utils.deprecation import deprecated
 
 
@@ -147,9 +148,13 @@ def _downgrade(self):
     def send_traces(self, traces):
         if not traces:
             return
+
         start = time.time()
-        data = self._encoder.encode_traces(traces)
-        response = self._put(self._traces, data, len(traces))
+        payload = Payload(encoder=self._encoder)
+        for trace in traces:
+            payload.add_trace(trace)
+
+        response = self._put(self._traces, payload.get_payload(), payload.length)
 
         # the API endpoint is not available so we should downgrade the connection and re-try the call
         if response.status in [404, 415] and self._fallback:
diff --git a/ddtrace/encoding.py b/ddtrace/encoding.py
@@ -1,4 +1,5 @@
 import json
+import struct
 
 from .internal.logger import get_logger
 
@@ -45,41 +46,78 @@ def encode_traces(self, traces):
         :param traces: A list of traces that should be serialized
         """
         normalized_traces = [[span.to_dict() for span in trace] for trace in traces]
-        return self._encode(normalized_traces)
+        return self.encode(normalized_traces)
 
-    def encode_services(self, services):
+    def encode_trace(self, trace):
         """
-        Encodes a dictionary of services.
+        Encodes a trace, expecting a list of spans. Before dump the string in a
+        serialized format all traces are normalized, calling the ``to_dict()`` method.
+        The traces nesting is not changed.
 
-        :param services: A dictionary that contains one or more services
+        :param trace: A list of traces that should be serialized
         """
-        return self._encode(services)
+        return self.encode([span.to_dict() for span in trace])
 
-    def _encode(self, obj):
+    def encode(self, obj):
         """
         Defines the underlying format used during traces or services encoding.
         This method must be implemented and should only be used by the internal functions.
         """
         raise NotImplementedError
 
+    def decode(self, data):
+        """
+        Defines the underlying format used during traces or services encoding.
+        This method must be implemented and should only be used by the internal functions.
+        """
+        raise NotImplementedError
+
+    def join_encoded(self, objs):
+        """Helper used to join a list of encoded objects into an encoded list of objects"""
+        raise NotImplementedError
+
 
 class JSONEncoder(Encoder):
     def __init__(self):
         # TODO[manu]: add instructions about how users can switch to Msgpack
         log.debug('using JSON encoder; application performance may be degraded')
         self.content_type = 'application/json'
 
-    def _encode(self, obj):
+    def encode(self, obj):
         return json.dumps(obj)
 
+    def decode(self, data):
+        return json.loads(data)
+
+    def join_encoded(self, objs):
+        """Join a list of encoded objects together as a json array"""
+        return '[' + ','.join(objs) + ']'
+
 
 class MsgpackEncoder(Encoder):
     def __init__(self):
         log.debug('using Msgpack encoder')
         self.content_type = 'application/msgpack'
 
-    def _encode(self, obj):
-        return msgpack.packb(obj, **MSGPACK_PARAMS)
+    def encode(self, obj):
+        return msgpack.packb(obj)
+
+    def decode(self, data):
+        return msgpack.unpackb(data)
+
+    def join_encoded(self, objs):
+        """Join a list of encoded objects together as a msgpack array"""
+        buf = b''.join(objs)
+
+        # Prepend array header to buffer
+        # https://github.com/msgpack/msgpack-python/blob/f46523b1af7ff2d408da8500ea36a4f9f2abe915/msgpack/fallback.py#L948-L955
+        count = len(objs)
+        if count <= 0xf:
+            return struct.pack('B', 0x90 + count) + buf
+        elif count <= 0xffff:
+            return struct.pack('>BH', 0xdc, count) + buf
+        else:
+            return struct.pack('>BI', 0xdd, count) + buf
 
 
 def get_encoder():
diff --git a/ddtrace/payload.py b/ddtrace/payload.py
@@ -0,0 +1,96 @@
+import logging
+
+from .encoding import get_encoder
+
+log = logging.getLogger(__name__)
+
+
+class Payload:
+    """
+    Trace agent API payload buffer class
+
+    This class is used to encoded and store traces to build the payload we send to
+    the trace agent.
+
+    DEV: We encoded and buffer traces so that we can reliable determine the size of
+         the payload easily so we can flush based on the payload size.
+    """
+    __slots__ = ('traces', 'size', 'encoder', 'max_payload_size')
+
+    # Default max payload size of 5mb
+    # DEV: Trace agent limit is 10mb, cutoff at 5mb to ensure we don't hit 10mb
+    DEFAULT_MAX_PAYLOAD_SIZE = 5 * 1000000
+
+    def __init__(self, encoder=None, max_payload_size=DEFAULT_MAX_PAYLOAD_SIZE):
+        """
+        Constructor for Payload
+
+        :param encoder: The encoded to use, default is the default encoder
+        :type encoder: ``ddtrace.encoding.Encoder``
+        :param max_payload_size: The max number of bytes a payload should be before
+            being considered full (default: 5mb)
+        """
+        self.max_payload_size = max_payload_size
+        self.encoder = encoder or get_encoder()
+        self.traces = []
+        self.size = 0
+
+    def add_trace(self, trace):
+        """
+        Encode and append a trace to this payload
+
+        :param trace: A trace to append
+        :type trace: A list of ``ddtrace.span.Span``s
+        """
+        # No trace or empty trace was given, ignore
+        if not trace:
+            return
+
+        # Encode the trace, append, and add it's length to the size
+        encoded = self.encoder.encode_trace(trace)
+        self.traces.append(encoded)
+        self.size += len(encoded)
+
+    @property
+    def length(self):
+        """
+        Get the number of traces in this payload
+
+        :returns: The number of traces in the payload
+        :rtype: int
+        """
+        return len(self.traces)
+
+    @property
+    def empty(self):
+        """
+        Whether this payload is empty or not
+
+        :returns: Whether this payload is empty or not
+        :rtype: bool
+        """
+        return self.length == 0
+
+    @property
+    def full(self):
+        """
+        Whether this payload is at or over the max allowed payload size
+
+        :returns: Whether we have reached the max payload size yet or not
+        :rtype: bool
+        """
+        return self.size >= self.max_payload_size
+
+    def get_payload(self):
+        """
+        Get the fully encoded payload
+
+        :returns: The fully encoded payload
+        :rtype: str | bytes
+        """
+        # DEV: `self.traces` is an array of encoded traces, `join_encoded` joins them together
+        return self.encoder.join_encoded(self.traces)
+
+    def __repr__(self):
+        """Get the string representation of this payload"""
+        return '{0}(length={1}, size={2}b, full={3})'.format(self.__class__.__name__, self.length, self.size, self.full)
diff --git a/tests/test_encoders.py b/tests/test_encoders.py
@@ -39,6 +39,42 @@ def test_encode_traces_json(self):
             for j in range(2):
                 eq_('client.testing', items[i][j]['name'])
 
+    def test_join_encoded_json(self):
+        # test encoding for JSON format
+        traces = []
+        traces.append([
+            Span(name='client.testing', tracer=None),
+            Span(name='client.testing', tracer=None),
+        ])
+        traces.append([
+            Span(name='client.testing', tracer=None),
+            Span(name='client.testing', tracer=None),
+        ])
+
+        encoder = JSONEncoder()
+
+        # Encode each trace on it's own
+        encoded_traces = [
+            encoder.encode_trace(trace)
+            for trace in traces
+        ]
+
+        # Join the encoded traces together
+        data = encoder.join_encoded(encoded_traces)
+
+        # Parse the resulting data
+        items = json.loads(data)
+
+        # test the encoded output that should be a string
+        # and the output must be flatten
+        ok_(isinstance(data, string_type))
+        eq_(len(items), 2)
+        eq_(len(items[0]), 2)
+        eq_(len(items[1]), 2)
+        for i in range(2):
+            for j in range(2):
+                eq_('client.testing', items[i][j]['name'])
+
     def test_encode_traces_msgpack(self):
         # test encoding for MsgPack format
         traces = []
@@ -64,3 +100,38 @@ def test_encode_traces_msgpack(self):
         for i in range(2):
             for j in range(2):
                 eq_(b'client.testing', items[i][j][b'name'])
+
+    def test_join_encoded_msgpack(self):
+        # test encoding for MsgPack format
+        traces = []
+        traces.append([
+            Span(name='client.testing', tracer=None),
+            Span(name='client.testing', tracer=None),
+        ])
+        traces.append([
+            Span(name='client.testing', tracer=None),
+            Span(name='client.testing', tracer=None),
+        ])
+
+        encoder = MsgpackEncoder()
+
+        # Encode each individual trace on it's own
+        encoded_traces = [
+            encoder.encode_trace(trace)
+            for trace in traces
+        ]
+        # Join the encoded traces together
+        data = encoder.join_encoded(encoded_traces)
+
+        # Parse the encoded data
+        items = msgpack.unpackb(data)
+
+        # test the encoded output that should be a string
+        # and the output must be flatten
+        ok_(isinstance(data, msgpack_type))
+        eq_(len(items), 2)
+        eq_(len(items[0]), 2)
+        eq_(len(items[1]), 2)
+        for i in range(2):
+            for j in range(2):
+                eq_(b'client.testing', items[i][j][b'name'])
diff --git a/tests/test_payload.py b/tests/test_payload.py
