Add service_rates to native transport responses

The `Writer` after_send callback calls `response.service_rates` to
update the priority sampler with agent-provided sampling rates.
Without this method, every successful flush would raise a
`NoMethodError` silently swallowed by the event publisher.

- Add `service_rates` to the pure-Ruby `Response` class, which
  parses the JSON `payload` and extracts `rate_by_service`
- Add `service_rates` returning `nil` to `InternalErrorResponse`
- Add RBS signatures for both

Author: lloeki

lib/datadog/tracing/transport/native.rb

@@ -1,5 +1,6 @@
 # frozen_string_literal: true
 
+require 'json'
 require_relative 'trace_formatter'
 require_relative 'statistics'
 
@@ -131,6 +132,7 @@ def not_found?;      false; end
           def unsupported?;    false; end
           def payload;         nil;   end
           def trace_count;     0;     end
+          def service_rates;   nil;   end
 
           def inspect
             "#<#{self.class} error=#{error.inspect}>"

lib/datadog/tracing/transport/native/response.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require 'json'
+
 module Datadog
   module Tracing
     module Transport
@@ -47,6 +49,21 @@ def not_found?
           def unsupported?
             @unsupported
           end
+
+          SERVICE_RATE_KEY = 'rate_by_service'
+
+          # Parse the agent's JSON response body and extract the
+          # +rate_by_service+ map.  Returns +nil+ when the payload
+          # is absent or does not contain sampling rates.
+          def service_rates
+            body = payload
+            return nil if body.nil? || body.empty?
+
+            parsed = JSON.parse(body)
+            parsed[SERVICE_RATE_KEY] if parsed.is_a?(Hash)
+          rescue JSON::ParserError
+            nil
+          end
         end
       end
     end

sig/datadog/tracing/transport/native.rbs

@@ -6,6 +6,12 @@ module Datadog
 
         def self?.supported?: () -> bool
 
+        class Response
+          SERVICE_RATE_KEY: String
+
+          def service_rates: () -> (Hash[String, Float] | nil)
+        end
+
         class Transport
           include Statistics
 
@@ -36,6 +42,7 @@ module Datadog
           def unsupported?: () -> false
           def payload: () -> nil
           def trace_count: () -> 0
+          def service_rates: () -> nil
           def inspect: () -> String
         end
       end

spec/datadog/tracing/transport/native/response_spec.rb

@@ -77,3 +77,63 @@ def make_response(ok:, internal_error: false, server_error: false,
     it { expect(response.payload).to be_nil }
   end
 end
+
+RSpec.describe 'Datadog::Tracing::Transport::Native::Response#service_rates' do
+  before do
+    skip_if_libdatadog_not_supported
+    require 'datadog/tracing/transport/native'
+  end
+
+  let(:response_class) { Datadog::Tracing::Transport::Native::Response }
+
+  def make_response(payload:)
+    resp = response_class.allocate
+    resp.instance_variable_set(:@ok, true)
+    resp.instance_variable_set(:@payload, payload)
+    resp
+  end
+
+  context 'with a valid rate_by_service payload' do
+    let(:payload) { '{"rate_by_service":{"service:web,env:prod":0.5}}' }
+
+    it 'returns the parsed rates hash' do
+      resp = make_response(payload: payload)
+      expect(resp.service_rates).to eq({'service:web,env:prod' => 0.5})
+    end
+  end
+
+  context 'with nil payload' do
+    it 'returns nil' do
+      resp = make_response(payload: nil)
+      expect(resp.service_rates).to be_nil
+    end
+  end
+
+  context 'with empty payload' do
+    it 'returns nil' do
+      resp = make_response(payload: '')
+      expect(resp.service_rates).to be_nil
+    end
+  end
+
+  context 'with invalid JSON' do
+    it 'returns nil' do
+      resp = make_response(payload: 'not json')
+      expect(resp.service_rates).to be_nil
+    end
+  end
+
+  context 'with JSON missing rate_by_service key' do
+    it 'returns nil' do
+      resp = make_response(payload: '{"other":"data"}')
+      expect(resp.service_rates).to be_nil
+    end
+  end
+end
+
+RSpec.describe 'Datadog::Tracing::Transport::Native::InternalErrorResponse#service_rates' do
+  it 'returns nil' do
+    resp = Datadog::Tracing::Transport::Native::InternalErrorResponse.new(RuntimeError.new('test'))
+    expect(resp.service_rates).to be_nil
+  end
+end