From 6eca5a3cf96d494c8b0cffcb08d15c25b399de0c Mon Sep 17 00:00:00 2001
From: Mark Rada <markrada26@gmail.com>
Date: Sat, 26 Mar 2016 19:06:12 -0400
Subject: [PATCH] Add IPv6 support to the UDP sender

And set the socket host and port at init time so that it does not have
to make a DNS call for every send.
---
 lib/gelf/transport/udp.rb | 55 +++++++++++++++++++++++++++------------
 test/test_ruby_sender.rb  | 25 +++++++++++++++---
 2 files changed, 59 insertions(+), 21 deletions(-)

diff --git a/lib/gelf/transport/udp.rb b/lib/gelf/transport/udp.rb
index 8b7f30c..4ff4791 100644
--- a/lib/gelf/transport/udp.rb
+++ b/lib/gelf/transport/udp.rb
@@ -1,40 +1,61 @@
 module GELF
   module Transport
     class UDP
-      attr_accessor :addresses
+      attr_reader :addresses
 
-      def initialize(addresses)
-        @addresses = addresses
+      def initialize(initial_addresses)
+        self.addresses = initial_addresses
       end
 
+      def addresses=(new_addresses)
+        @addresses = new_addresses
+        reset_sockets
+      end
+      
       def send_datagrams(datagrams)
-        socket = get_socket
-        idx = get_address_index
-
-        host, port = @addresses[idx]
-        set_address_index((idx + 1) % @addresses.length)
+        sock = socket
         datagrams.each do |datagram|
-          socket.send(datagram, 0, host, port)
+          sock.send(datagram, 0)
         end
       end
 
       def close
-        socket = get_socket
-        socket.close if socket
+        reset_sockets
       end
 
       private
 
-      def get_socket
-        Thread.current[:gelf_udp_socket] ||= UDPSocket.open
+      def socket
+        idx = socket_index
+        sock = sockets[idx]
+        set_socket_index((idx + 1) % @addresses.length)
+        sock
       end
 
-      def get_address_index
-        Thread.current[:gelf_udp_address_idx] ||= 0
+      def sockets
+        Thread.current[:gelf_udp_sockets] ||= configure_sockets
       end
 
-      def set_address_index(value)
-        Thread.current[:gelf_udp_address_idx] = value
+      def reset_sockets
+        return unless Thread.current.key?(:gelf_udp_sockets)
+        Thread.current[:gelf_udp_sockets].each(&:close)
+        Thread.current[:gelf_udp_sockets] = nil
+      end
+
+      def socket_index
+        Thread.current[:gelf_udp_socket_idx] ||= 0
+      end
+
+      def set_socket_index(value)
+        Thread.current[:gelf_udp_socket_idx] = value
+      end
+
+      def configure_sockets
+        @addresses.map do |host, port|
+          UDPSocket.new(Addrinfo.ip(host).afamily).tap do |socket|
+            socket.connect(host, port)
+          end
+        end
       end
     end
   end
diff --git a/test/test_ruby_sender.rb b/test/test_ruby_sender.rb
index fc7c41b..a5f10b2 100644
--- a/test/test_ruby_sender.rb
+++ b/test/test_ruby_sender.rb
@@ -9,6 +9,23 @@ class TestRubyUdpSender < Test::Unit::TestCase
       @datagrams2 = %w(e1 e2 e3)
     end
 
+    context "setup_sockets" do
+      setup do
+        @sender.send_datagrams(%w(a1))
+        @sender.send_datagrams(%w(b1))
+      end
+
+      before_should "be configured with a socket for each address" do
+        UDPSocket.any_instance.expects(:connect).with do |host,port|
+          host == 'localhost' && port == 12201
+        end
+        UDPSocket.any_instance.expects(:connect).with do |host,port|
+          host == 'localhost' && port == 12202
+        end
+        UDPSocket.any_instance.expects(:send).times(2).returns(nil)
+      end
+    end
+
     context "send_datagrams" do
       setup do
         @sender.send_datagrams(@datagrams1)
@@ -16,11 +33,11 @@ class TestRubyUdpSender < Test::Unit::TestCase
       end
 
       before_should "be called 3 times with 1st and 2nd address" do
-        UDPSocket.any_instance.expects(:send).times(3).with do |datagram, _, host, port|
-          datagram.start_with?('d') && host == 'localhost' && port == 12201
+        UDPSocket.any_instance.expects(:send).times(3).with do |datagram,_|
+          datagram.start_with?('d')
         end
-        UDPSocket.any_instance.expects(:send).times(3).with do |datagram, _, host, port|
-          datagram.start_with?('e') && host == 'localhost' && port == 12202
+        UDPSocket.any_instance.expects(:send).times(3).with do |datagram,_|
+          datagram.start_with?('e')
         end
       end
     end