Reuse HTTP connection between trace flushes

Author: marcotc

lib/datadog/tracing/writer.rb

@@ -95,6 +95,8 @@ def stop
       def stop_worker
         @stopped = true
 
+        @transport.stop
+
         return if @worker.nil?
 
         @worker.stop

lib/ddtrace/transport/http/adapters/net.rb

@@ -7,17 +7,24 @@ module Datadog
   module Transport
     module HTTP
       module Adapters
-        # Adapter for Net::HTTP
+        # Adapter for Net::HTTP.
+        # This adapter reuses a long-running HTTP connection with the target URL
+        # when possible.
+        # It is required to invoke {#close} to ensure the connection is closed.
         class Net
           attr_reader \
             :hostname,
             :port,
             :timeout,
             :ssl
 
-          # in seconds
+          # In seconds.
           DEFAULT_TIMEOUT = 30
 
+          # If this many seconds have passed since the last request, create a new TCP connection.
+          # In seconds.
+          KEEP_ALIVE_TIMEOUT = 60
+
           # @deprecated Positional parameters are deprecated. Use named parameters instead.
           def initialize(hostname = nil, port = nil, **options)
             @hostname = hostname || options.fetch(:hostname)
@@ -39,12 +46,17 @@ def open(&block)
             # DEV Initializing +Net::HTTP+ directly help us avoid expensive
             # options processing done in +Net::HTTP.start+:
             # https://github.com/ruby/ruby/blob/b2d96abb42abbe2e01f010ffc9ac51f0f9a50002/lib/net/http.rb#L614-L618
-            req = ::Net::HTTP.new(hostname, port, nil)
+            @req ||= begin
+              req = ::Net::HTTP.new(hostname, port, nil)
+
+              req.use_ssl = ssl
+              req.open_timeout = req.read_timeout = timeout
+              req.keep_alive_timeout = KEEP_ALIVE_TIMEOUT
 
-            req.use_ssl = ssl
-            req.open_timeout = req.read_timeout = timeout
+              req.start # Calling #start without a block creates a long-lived connection
+            end
 
-            req.start(&block)
+            yield @req
           end
 
           def call(env)
@@ -82,6 +94,14 @@ def url
             "http://#{hostname}:#{port}?timeout=#{timeout}"
           end
 
+          def close
+            # Clean up HTTP socket
+            if @req && @req.started?
+              @req.finish
+              @req = nil
+            end
+          end
+
           # Raised when called with an unknown HTTP method
           class UnknownHTTPMethod < StandardError
             attr_reader :verb

lib/ddtrace/transport/http/client.rb

@@ -48,6 +48,10 @@ def send_request(request, &block)
         def build_env(request)
           Env.new(request)
         end
+
+        def close
+          @api.adapter.close if @api.adapter.respond_to?(:close)
+        end
       end
     end
   end

lib/ddtrace/transport/io/traces.rb

@@ -43,6 +43,8 @@ def send_traces(traces)
               Traces::Response.new(result)
             end]
           end
+
+          def stop; end
         end
 
         # Encoder for IO-specific trace encoding

lib/ddtrace/transport/traces.rb

@@ -165,6 +165,10 @@ def current_api
           apis[@current_api_id]
         end
 
+        def stop
+          @client.close if @client
+        end
+
         private
 
         def downgrade?(response)
@@ -183,6 +187,8 @@ def downgrade!
         def change_api!(api_id)
           raise UnknownApiVersionError, api_id unless apis.key?(api_id)
 
+          @client.close if @client
+
           @current_api_id = api_id
           @client = HTTP::Client.new(current_api)
         end

spec/datadog/integration_spec.rb

@@ -4,6 +4,8 @@
 require 'datadog/opentracer'
 require 'datadog/statsd'
 
+require_relative '../support/system_helper'
+
 RSpec.describe 'Datadog integration' do
   context 'graceful shutdown', :integration do
     before do
@@ -42,27 +44,15 @@ def thread_count
     end
 
     context 'for file descriptors' do
-      def open_file_descriptors
-        # Unix-specific way to get the current process' open file descriptors and the files (if any) they correspond to
-        Dir['/dev/fd/*'].each_with_object({}) do |fd, hash|
-          hash[fd] =
-            begin
-              File.realpath(fd)
-            rescue SystemCallError # This can fail due to... reasons, and we only want it for debugging so let's ignore
-              nil
-            end
-        end
-      end
-
       it 'closes tracer file descriptors (known flaky test)' do
-        before_open_file_descriptors = open_file_descriptors
+        before_open_file_descriptors = SystemHelper.open_fds
 
         start_tracer
         wait_for_tracer_sent
 
         shutdown
 
-        after_open_file_descriptors = open_file_descriptors
+        after_open_file_descriptors = SystemHelper.open_fds
 
         expect(after_open_file_descriptors.size)
           .to(

spec/datadog/tracing/integration_spec.rb

@@ -1,6 +1,7 @@
 # typed: ignore
 
 require 'spec_helper'
+require_relative '../../support/system_helper'
 
 require 'datadog/statsd'
 
@@ -20,6 +21,10 @@
 require 'ddtrace/transport/io/client'
 require 'ddtrace/transport/traces'
 
+require_relative '../../support/test_access_refinement'
+
+using TestAccessRefinement
+
 RSpec.describe 'Tracer integration tests' do
   shared_context 'agent-based test' do
     before do
@@ -45,6 +50,11 @@
     it { expect(stats).to include(traces_flushed: 0) }
   end
 
+  before do
+    WebMock.disable!
+    Datadog.configuration.reset!
+  end
+
   after { tracer.shutdown! }
 
   describe 'agent receives span' do
@@ -57,14 +67,6 @@ def create_trace
       end
     end
 
-    def wait_for_flush(stat, num = 1)
-      test_repeat.times do
-        break if tracer.writer.stats[stat] >= num
-
-        sleep(0.1)
-      end
-    end
-
     def agent_receives_span_step1
       stats = tracer.writer.stats
       expect(stats[:traces_flushed]).to eq(0)
@@ -78,7 +80,7 @@ def agent_receives_span_step2
       create_trace
 
       # Timeout after 3 seconds, waiting for 1 flush
-      wait_for_flush(:traces_flushed)
+      tracer.writer.force_flush
 
       stats = tracer.writer.stats
       expect(stats[:traces_flushed]).to eq(1)
@@ -95,8 +97,8 @@ def agent_receives_span_step2
     def agent_receives_span_step3(previous_success)
       create_trace
 
-      # Timeout after 3 seconds, waiting for another flush
-      wait_for_flush(:traces_flushed, 2)
+      # Timeout after 3 seconds, waiting for second flush
+      tracer.writer.force_flush
 
       stats = tracer.writer.stats
       expect(stats[:traces_flushed]).to eq(2)
@@ -143,24 +145,32 @@ def agent_receives_span_step3(previous_success)
   describe 'agent receives short span' do
     include_context 'agent-based test'
 
-    before do
+    def trace
       tracer.trace('my.short.op') do |span|
         @span = span
         span.service = 'my.service'
       end
 
-      @first_shutdown = tracer.shutdown!
+      tracer.writer.force_flush
     end
 
     let(:stats) { tracer.writer.stats }
 
     it do
-      expect(@first_shutdown).to be true
+      trace
+
       expect(@span.finished?).to be true
       expect(stats[:services_flushed]).to be_nil
     end
 
-    it_behaves_like 'flushed trace'
+    it 'reuses the same HTTP connection' do
+      trace
+      expect { trace }.to_not(change { SystemHelper.open_fds })
+    end
+
+    it_behaves_like 'flushed trace' do
+      before { trace }
+    end
   end
 
   describe 'rule sampler' do
@@ -319,6 +329,8 @@ def agent_receives_span_step3(previous_success)
           function.call(traces)
         end
 
+      WebMock.enable!
+
       trace # Run test subject
       tracer.shutdown! # Ensure trace is flushed, so we can read writer statistics
     end
@@ -347,6 +359,7 @@ def agent_receives_span_step3(previous_success)
     end
 
     after do
+      WebMock.disable!
       Datadog.configuration.tracing.sampling.reset!
     end
 
@@ -468,6 +481,8 @@ def sample!(trace)
           span.service = 'my.service'
         end
 
+        tracer.writer.force_flush
+
         threads = Array.new(10) do
           Thread.new { tracer.shutdown! }
         end
@@ -849,7 +864,7 @@ def sample!(trace)
 
           tracer.trace('name') {}
 
-          try_wait_until(attempts: 20) { tracer.writer.stats[:traces_flushed] >= 1 }
+          tracer.writer.force_flush
 
           stats = tracer.writer.stats
           expect(stats[:traces_flushed]).to eq(1)

spec/datadog/tracing/writer_spec.rb

@@ -22,7 +22,7 @@
     subject(:writer) { described_class.new(options) }
 
     let(:options) { { transport: transport } }
-    let(:transport) { instance_double(Datadog::Transport::Traces::Transport) }
+    let(:transport) { instance_double(Datadog::Transport::Traces::Transport, stop: nil) }
 
     describe 'behavior' do
       describe '#initialize' do

spec/ddtrace/transport/http/adapters/net_spec.rb

@@ -26,9 +26,10 @@
 
       allow(http_connection).to receive(:open_timeout=).with(adapter.timeout)
       allow(http_connection).to receive(:read_timeout=).with(adapter.timeout)
+      allow(http_connection).to receive(:keep_alive_timeout=).with(60)
       allow(http_connection).to receive(:use_ssl=).with(adapter.ssl)
 
-      allow(http_connection).to receive(:start).and_yield(http_connection)
+      allow(http_connection).to receive(:start).and_return(http_connection)
     end
   end
 
@@ -186,6 +187,29 @@
 
     it { is_expected.to eq('http://local.test:345?timeout=7') }
   end
+
+  describe '#close' do
+    include_context 'HTTP connection stub'
+
+    before do
+      allow(http_connection).to receive(:started?).and_return(true)
+
+      adapter.open {}
+    end
+
+    it 'closes connection' do
+      expect(http_connection).to receive(:finish).once
+
+      adapter.close
+    end
+
+    it 'close can be called multiple times' do
+      expect(http_connection).to receive(:finish).once
+
+      adapter.close
+      adapter.close
+    end
+  end
 end
 
 RSpec.describe Datadog::Transport::HTTP::Adapters::Net::Response do

spec/ddtrace/transport/traces_spec.rb

@@ -132,8 +132,8 @@
       ].with_fallbacks(v2: :v1)
     end
 
-    let(:api_v1) { instance_double(Datadog::Transport::HTTP::API::Instance, 'v1', encoder: encoder_v1) }
-    let(:api_v2) { instance_double(Datadog::Transport::HTTP::API::Instance, 'v2', encoder: encoder_v2) }
+    let(:api_v1) { instance_double(Datadog::Transport::HTTP::API::Instance, 'v1', encoder: encoder_v1, adapter: nil) }
+    let(:api_v2) { instance_double(Datadog::Transport::HTTP::API::Instance, 'v2', encoder: encoder_v2, adapter: nil) }
     let(:encoder_v1) { instance_double(Datadog::Core::Encoding::Encoder, content_type: 'text/plain') }
     let(:encoder_v2) { instance_double(Datadog::Core::Encoding::Encoder, content_type: 'text/csv') }
   end
@@ -162,8 +162,8 @@
     let(:lazy_chunks) { chunks.lazy }
 
     let(:request) { instance_double(Datadog::Transport::Traces::Request) }
-    let(:client_v2) { instance_double(Datadog::Transport::HTTP::Client) }
-    let(:client_v1) { instance_double(Datadog::Transport::HTTP::Client) }
+    let(:client_v2) { instance_double(Datadog::Transport::HTTP::Client, close: nil) }
+    let(:client_v1) { instance_double(Datadog::Transport::HTTP::Client, close: nil) }
 
     let(:chunker) { instance_double(Datadog::Transport::Traces::Chunker, max_size: 1) }