Trace payload chunking

lib/ddtrace/encoding.rb

@@ -10,19 +10,84 @@ def content_type
         raise NotImplementedError
       end
 
-      # Encodes a list of traces, expecting a list of items where each items
-      # is a list of spans. Before dump the string in a serialized format all
-      # traces are normalized. The traces nesting is not changed.
-      def encode_traces(traces)
-        to_send = []
+      # Trace agent limit payload size of 10 MiB (as of agent v6.14.1):
+      # https://github.com/DataDog/datadog-agent/blob/6.14.1/pkg/trace/api/api.go#L46
+      #
+      # This value is set lower than the hard limit, in case transport overhead pushes
+      # the payload size over the limit.
+      DEFAULT_MAX_PAYLOAD_SIZE = 9 * 1024 * 1024
+
+      # Encodes a list of traces in batches, expecting a list of items where each items
+      # is a list of spans.
+      # A serialized batch payload will not exceed +max_size+.
+      # Single traces larger than +max_size+ will be discarded.
+      # Before serializing, all traces are normalized. Trace nesting is not changed.
+      #
+      # @param traces [Array<Trace>] list of traces
+      # @param max_size [String] maximum acceptable payload size
+      # @yield [encoded_batch, batch_size] block invoked for every serialized batch of traces
+      # @yieldparam encoded_batch [String] serialized batch of traces, ready to be transmitted
+      # @yieldparam batch_size [Integer] number of traces serialized in this batch
+      # @return concatenated list of return values from the provided block
+      def encode_traces(traces, max_size: DEFAULT_MAX_PAYLOAD_SIZE)
+        # Captures all return values from the provided block
+        returns = []
+
+        encoded_batch = []
+        batch_size = 0
         traces.each do |trace|
-          to_send << trace.map(&:to_hash)
+          encoded_trace = encode_one(trace, max_size)
+
+          next unless encoded_trace
+
+          if encoded_trace.size + batch_size > max_size
+            # Can't fit trace in current batch
+            # TODO Datadog::Debug::HealthMetrics.increment('tracer.encoder.batch.chunked')
+
+            # Flush current batch
+            returns << yield(join(encoded_batch), encoded_batch.size)
+            # TODO: Datadog::Debug::HealthMetrics.increment('tracer.encoder.batch.yield')
+
+            # Create new batch
+            encoded_batch = []
+            batch_size = 0
+          end
+
+          encoded_batch << encoded_trace
+          batch_size += encoded_trace.size
+        end
+
+        unless encoded_batch.empty?
+          returns << yield(join(encoded_batch), encoded_batch.size)
+          # TODO: Datadog::Debug::HealthMetrics.increment('tracer.encoder.batch.yield')
+        end
+
+        returns
+      end
+
+      private
+
+      def encode_one(trace, max_size)
+        encoded = encode(trace.map(&:to_hash))
+
+        # TODO: Datadog::Debug::HealthMetrics.increment('tracer.encoder.trace.encode')
+        if encoded.size > max_size
+          # This single trace is too large, we can't flush it
+          # TODO should I add `trace` to the message? It will definitely be very large!
+          # TODO Datadog::Debug::HealthMetrics.increment('tracer.encoder.trace.max_size_exceed')
+          Datadog::Tracer.log.debug('Trace payload too large')
+          return nil
         end
-        encode(to_send)
+
+        encoded
+      end
+
+      # Concatenates a list of traces previously encoded by +#encode+.
+      def join(encoded_traces)
+        raise NotImplementedError
       end
 
-      # Defines the underlying format used during traces or services encoding.
-      # This method must be implemented and should only be used by the internal functions.
+      # Serializes a single trace into a String suitable for network transmission.
       def encode(_)
         raise NotImplementedError
       end
@@ -40,8 +105,12 @@ def content_type
         CONTENT_TYPE
       end
 
-      def encode(obj)
-        JSON.dump(obj)
+      def encode(trace)
+        JSON.dump(trace)
+      end
+
+      def join(encoded_traces)
+        "[#{encoded_traces.join(',')}]"
       end
     end
 
@@ -57,8 +126,15 @@ def content_type
         CONTENT_TYPE
       end
 
-      def encode(obj)
-        MessagePack.pack(obj)
+      def encode(trace)
+        MessagePack.pack(trace)
+      end
+
+      def join(encoded_traces)
+        packer = MessagePack::Packer.new
+        packer.write_array_header(encoded_traces.size)
+
+        (packer.to_a + encoded_traces).join
       end
     end
   end

lib/ddtrace/transport/http/client.rb

@@ -23,19 +23,19 @@ def send_request(request, &block)
           # Build request into env
           env = build_env(request)
 
-          # Get response from API
-          response = yield(current_api, env)
+          # Get responses from API
+          responses = yield(current_api, env)
 
           # Update statistics
-          update_stats_from_response!(response)
+          responses.each { |r| update_stats_from_response!(r) }
 
           # If API should be downgraded, downgrade and try again.
-          if downgrade?(response)
+          if responses.find { |r| downgrade?(r) }
             downgrade!
-            response = send_request(request, &block)
+            responses = send_request(request, &block)
           end
 
-          response
+          responses
         rescue StandardError => e
           message = "Internal error during HTTP transport request. Cause: #{e.message} Location: #{e.backtrace.first}"
 

lib/ddtrace/transport/http/traces.rb

@@ -17,6 +17,7 @@ class Response
           def initialize(http_response, options = {})
             super(http_response)
             @service_rates = options.fetch(:service_rates, nil)
+            @trace_count = options.fetch(:trace_count, 0)
           end
         end
 
@@ -103,18 +104,30 @@ def service_rates?
             end
 
             def call(env, &block)
+              encoder.encode_traces(env.request.parcel.data) do |encoded_data, count|
+                # Ensure no data is leaked between each request.
+                # We have perform this copy before we start modifying headers and body.
+                new_env = env.dup
+
+                process_batch(new_env, encoded_data, count) { |e| super(e, &block) }
+              end
+            end
+
+            private
+
+            def process_batch(env, encoded_data, count)
               # Add trace count header
-              env.headers[HEADER_TRACE_COUNT] = env.request.parcel.count.to_s
+              env.headers[HEADER_TRACE_COUNT] = count.to_s
 
               # Encode body & type
               env.headers[HEADER_CONTENT_TYPE] = encoder.content_type
-              env.body = env.request.parcel.encode_with(encoder)
+              env.body = encoded_data
 
               # Query for response
-              http_response = super(env, &block)
+              http_response = yield env
 
               # Process the response
-              response_options = {}.tap do |options|
+              response_options = { trace_count: count }.tap do |options|
                 # Parse service rates, if configured to do so.
                 if service_rates? && !http_response.payload.to_s.empty?
                   body = JSON.parse(http_response.payload)

lib/ddtrace/transport/traces.rb

@@ -11,10 +11,6 @@ class Parcel
         def count
           data.length
         end
-
-        def encode_with(encoder)
-          encoder.encode_traces(data)
-        end
       end
 
       # Traces request
@@ -26,7 +22,7 @@ def initialize(traces)
 
       # Traces response
       module Response
-        attr_reader :service_rates
+        attr_reader :service_rates, :trace_count
       end
     end
   end

lib/ddtrace/writer.rb

@@ -78,20 +78,19 @@ def send_spans(traces, transport)
       # Inject hostname if configured to do so
       inject_hostname!(traces) if Datadog.configuration.report_hostname
 
-      # Send traces an get a response.
-      response = transport.send_traces(traces)
+      # Send traces and get responses
+      responses = transport.send_traces(traces)
 
-      unless response.internal_error?
-        @traces_flushed += traces.length unless response.server_error?
-
-        # Update priority sampler
-        unless priority_sampler.nil? || response.service_rates.nil?
-          priority_sampler.update(response.service_rates)
-        end
+      # Tally up successful flushes
+      responses.reject { |x| x.internal_error? || x.server_error? }.each do |response|
+        @traces_flushed += response.trace_count
       end
 
+      # Update priority sampler
+      update_priority_sampler(responses.last)
+
       # Return if server error occurred.
-      !response.server_error?
+      !responses.find(&:server_error?)
     end
 
     def send_runtime_metrics
@@ -153,5 +152,11 @@ def inject_hostname!(traces)
         end
       end
     end
+
+    def update_priority_sampler(response)
+      return unless response && !response.internal_error? && priority_sampler && response.service_rates
+
+      priority_sampler.update(response.service_rates)
+    end
   end
 end

spec/ddtrace/encoding_spec.rb

@@ -0,0 +1,87 @@
+require 'spec_helper'
+require 'spec/support/language_helpers'
+
+require 'ddtrace/encoding'
+
+RSpec.describe Datadog::Encoding do
+  subject(:encode) { encoder.method(:encode_traces) }
+
+  let(:block) { proc { block_response } }
+  let(:block_response) { double('response') }
+
+  context 'Base encoder' do
+    let(:encoder) { Class.new { include Datadog::Encoding::Encoder }.new }
+
+    let(:traces) { get_test_traces(3) }
+
+    before do
+      allow(encoder).to receive(:encode).with(traces[0].map(&:to_hash)).and_return('blob1')
+      allow(encoder).to receive(:encode).with(traces[1].map(&:to_hash)).and_return('blob2')
+      allow(encoder).to receive(:encode).with(traces[2].map(&:to_hash)).and_return('blob3')
+      allow(encoder).to receive(:join) { |arr| arr.join(',') }
+    end
+
+    it do
+      expect { |b| encode.call(traces, &b) }.to yield_with_args('blob1,blob2,blob3', 3)
+    end
+
+    it 'returns yielded block returns' do
+      expect(encode.call(traces, &block)).to eq([block_response])
+    end
+
+    context 'with large batch of traces' do
+      let(:max_size) { 10 }
+
+      it do
+        expect { |b| encode.call(traces, max_size: max_size, &b) }
+          .to yield_successive_args(['blob1,blob2', 2], ['blob3', 1])
+      end
+
+      it 'returns yielded block returns' do
+        expect(encode.call(traces, max_size: max_size, &block)).to eq([block_response, block_response])
+      end
+    end
+
+    context 'with individual traces too large' do
+      let(:max_size) { 4 }
+
+      it do
+        expect { |b| encode.call(traces, max_size: max_size, &b) }.not_to yield_control
+      end
+    end
+  end
+
+  context 'Msgpack encoding' do
+    let(:encoder) { Datadog::Encoding::MsgpackEncoder }
+    let(:traces) { get_test_traces(2) }
+
+    it do
+      expect(encode.call(traces) do |encoded, size|
+        expect(size).to eq(2)
+
+        items = MessagePack.unpack(encoded)
+        expect(items.size).to eq(2)
+        expect(items.first).to eq(traces.first.map(&:to_hash).map(&:stringify_keys))
+
+        block_response
+      end).to eq([block_response])
+    end
+  end
+
+  context 'JSON encoding' do
+    let(:encoder) { Datadog::Encoding::JSONEncoder }
+    let(:traces) { get_test_traces(2) }
+
+    it do
+      expect(encode.call(traces) do |encoded, size|
+        expect(size).to eq(2)
+
+        items = JSON.parse(encoded)
+        expect(items.size).to eq(2)
+        expect(items.first).to eq(traces.first.map(&:to_hash).map(&:stringify_keys))
+
+        block_response
+      end).to eq([block_response])
+    end
+  end
+end