Introduce method for send buffer size and extract module

Signed-off-by: Pedro Tanaka <[email protected]>

lib/statsd/instrument/connection_behavior.rb

@@ -0,0 +1,37 @@
+# frozen_string_literal: true
+
+module StatsD
+  module Instrument
+    module ConnectionBehavior
+      def close
+        @socket&.close
+      rescue IOError, SystemCallError => e
+        StatsD.logger.debug do
+          "[#{self.class.name}] Error closing socket: #{e.class}: #{e.message}"
+        end
+      ensure
+        @socket = nil
+      end
+
+      def send_buffer_size
+        if socket
+          socket.getsockopt(Socket::SOL_SOCKET, Socket::SO_SNDBUF).int
+        else
+          @max_packet_size
+        end
+      end
+
+      private
+
+      def setup_socket(socket)
+        socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_SNDBUF, @max_packet_size.to_i)
+        socket
+      rescue IOError => e
+        StatsD.logger.debug do
+          "[#{self.class.name}] Failed to create socket: #{e.class}: #{e.message}"
+        end
+        nil
+      end
+    end
+  end
+end

lib/statsd/instrument/udp_connection.rb

@@ -3,6 +3,8 @@
 module StatsD
   module Instrument
     class UdpConnection
+      include ConnectionBehavior
+
       DEFAULT_MAX_PACKET_SIZE = 1_472
 
       attr_reader :host, :port
@@ -17,11 +19,6 @@ def send_datagram(message)
         socket.send(message, 0)
       end
 
-      def close
-        @socket&.close
-        @socket = nil
-      end
-
       def type
         :udp
       end
@@ -31,9 +28,9 @@ def type
       def socket
         @socket ||= begin
           socket = UDPSocket.new
-          socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_SNDBUF, @max_packet_size)
-          socket.connect(@host, @port)
-          socket
+          setup_socket(socket)&.tap do |s|
+            s.connect(@host, @port)
+          end
         end
       end
     end

lib/statsd/instrument/uds_connection.rb

@@ -3,6 +3,8 @@
 module StatsD
   module Instrument
     class UdsConnection
+      include ConnectionBehavior
+
       DEFAULT_MAX_PACKET_SIZE = 8_192
 
       def initialize(socket_path, max_packet_size: DEFAULT_MAX_PACKET_SIZE)
@@ -17,17 +19,7 @@ def initialize(socket_path, max_packet_size: DEFAULT_MAX_PACKET_SIZE)
       end
 
       def send_datagram(message)
-        socket.sendmsg(message, 0)
-      end
-
-      def close
-        @socket&.close
-      rescue IOError, SystemCallError => e
-        StatsD.logger.debug do
-          "[#{self.class.name}] Error closing socket: #{e.class}: #{e.message}"
-        end
-      ensure
-        @socket = nil
+        socket&.sendmsg(message, 0)
       end
 
       def host
@@ -47,14 +39,9 @@ def type
       def socket
         @socket ||= begin
           socket = Socket.new(Socket::AF_UNIX, Socket::SOCK_DGRAM)
-          socket.setsockopt(Socket::SOL_SOCKET, Socket::SO_SNDBUF, @max_packet_size.to_i)
-          socket.connect(Socket.pack_sockaddr_un(@socket_path))
-          socket
-        rescue IOError => e
-          StatsD.logger.debug do
-            "[#{self.class.name}] Failed to create socket: #{e.class}: #{e.message}"
+          setup_socket(socket)&.tap do |s|
+            s.connect(Socket.pack_sockaddr_un(@socket_path))
           end
-          nil
         end
       end
     end