Merge pull request #204 from DataDog/remeh/connection

[connection] clearer connection design and retry mechanism.

lib/datadog/statsd/connection.rb

@@ -8,16 +8,6 @@ def initialize(telemetry: nil, logger: nil)
         @logger = logger
       end
 
-      # Close the underlying socket
-      def close
-        begin
-          @socket && @socket.close if instance_variable_defined?(:@socket)
-        rescue StandardError => boom
-          logger.error { "Statsd: #{boom.class} #{boom}" } if logger
-        end
-        @socket = nil
-      end
-
       def write(payload)
         logger.debug { "Statsd: #{payload}" } if logger
 
@@ -36,6 +26,7 @@ def write(payload)
           retries += 1
           begin
             close
+            connect
             retry
           rescue StandardError => e
             boom = e
@@ -48,11 +39,16 @@ def write(payload)
       end
 
       private
+
       attr_reader :telemetry
       attr_reader :logger
 
-      def socket
-        @socket ||= connect
+      def connect
+        raise 'Should be implemented by subclass'
+      end
+
+      def close
+        raise 'Should be implemented by subclass'
       end
     end
   end

lib/datadog/statsd/udp_connection.rb

@@ -19,18 +19,26 @@ def initialize(host, port, **kwargs)
 
         @host = host || ENV.fetch('DD_AGENT_HOST', DEFAULT_HOST)
         @port = port || ENV.fetch('DD_DOGSTATSD_PORT', DEFAULT_PORT).to_i
+        @socket = nil
+        connect
+      end
+
+      def close
+        @socket.close if @socket
+        @socket = nil
       end
 
       private
 
       def connect
-        UDPSocket.new.tap do |socket|
-          socket.connect(host, port)
-        end
+        close if @socket
+
+        @socket = UDPSocket.new
+        @socket.connect(host, port)
       end
 
       def send_message(message)
-        socket.send(message, 0)
+        @socket.send(message, 0)
       end
     end
   end

lib/datadog/statsd/uds_connection.rb

@@ -14,20 +14,27 @@ def initialize(socket_path, **kwargs)
         super(**kwargs)
 
         @socket_path = socket_path
+        @socket = nil
+        connect
+      end
+
+      def close
+        @socket.close if @socket
+        @socket = nil
       end
 
       private
 
       def connect
-        socket = Socket.new(Socket::AF_UNIX, Socket::SOCK_DGRAM)
-        socket.connect(Socket.pack_sockaddr_un(@socket_path))
-        socket
+        close if @socket
+
+        @socket = Socket.new(Socket::AF_UNIX, Socket::SOCK_DGRAM)
+        @socket.connect(Socket.pack_sockaddr_un(@socket_path))
       end
 
       def send_message(message)
-        socket.sendmsg_nonblock(message)
+        @socket.sendmsg_nonblock(message)
       rescue Errno::ECONNREFUSED, Errno::ECONNRESET, Errno::ENOENT => e
-        @socket = nil
         # TODO: FIXME: This error should be considered as a retryable error in the
         # Connection class. An even better solution would be to make BadSocketError inherit
         # from a specific retryable error class in the Connection class.

spec/statsd_spec.rb

@@ -160,11 +160,8 @@
           )
         end
 
-        before do
-          allow(Socket).to receive(:new).and_call_original
-        end
-
         it 'uses an UDS socket' do
+          expect(Socket).to receive(:new).and_return(fake_socket)
           expect(subject.transport_type).to eq :uds
         end