Merge pull request #591 from DataDog/fix/priority_sampling_response_p…

…arsing Fix priority sampling mishandling agent response
DataDog · Oct 17, 2018 · 3be61a8 · 3be61a8
2 parents d8d6777 + 8ae1836
commit 3be61a8
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 3 deletions.
diff --git a/lib/ddtrace/transport.rb b/lib/ddtrace/transport.rb
@@ -221,6 +221,7 @@ def process_callback(action, response)
       @response_callback.call(action, response, @api)
     rescue => e
       Tracer.log.debug("Error processing callback: #{e}")
+      @mutex.synchronize { @count_internal_error += 1 }
     end
   end
 end
diff --git a/lib/ddtrace/writer.rb b/lib/ddtrace/writer.rb
@@ -120,8 +120,10 @@ def sampling_updater(action, response, api)
       return unless action == :traces && response.is_a?(Net::HTTPOK)
 
       if api[:version] == HTTPTransport::V4
-        service_rates = JSON.parse(response.body)
-        @priority_sampler.update(service_rates)
+        body = JSON.parse(response.body)
+        if body.is_a?(Hash) && body.key?('rate_by_service')
+          @priority_sampler.update(body['rate_by_service'])
+        end
         true
       else
         false

diff --git a/spec/ddtrace/writer_spec.rb b/spec/ddtrace/writer_spec.rb
@@ -167,8 +167,9 @@
             context 'and is a :traces action' do
               context 'and is API v4' do
                 let(:api) { { version: Datadog::HTTPTransport::V4 } }
+                let(:body) { sampling_response.to_json }
+                let(:sampling_response) { { 'rate_by_service' => service_rates } }
                 let(:service_rates) { { 'service:a,env:test' => 0.1, 'service:b,env:test' => 0.5 } }
-                let(:body) { service_rates.to_json }
 
                 it do
                   expect(sampler).to receive(:update).with(service_rates)