fix server-mode shutdown bug, add baseline specs

Author: yaauie

lib/logstash/outputs/tcp.rb

@@ -169,19 +169,19 @@ def register
         end
       end
     else
-      client_socket = nil
+      @client_socket = nil
       peer_info = nil
       @codec.on_event do |event, payload|
         begin
           # not threadsafe; this is why we require `concurrency: single`
-          unless client_socket
-            client_socket = connect
-            peer_info = client_socket.peer
+          unless @client_socket
+            @client_socket = connect
+            peer_info = @client_socket.peer
           end
 
           writable_io = nil
           while writable_io.nil? || writable_io.any? == false
-            readable_io, writable_io, _ = IO.select([client_socket],[client_socket])
+            readable_io, writable_io, _ = IO.select([@client_socket],[@client_socket])
 
             # don't expect any reads, but a readable socket might
             # mean the remote end closed, so read it and throw it away.
@@ -192,15 +192,15 @@ def register
           # Now send the payload
           @logger.trace("transmitting #{payload.bytesize} bytes", socket: peer_info) if @logger.trace? && payload && !payload.empty?
           while payload && payload.bytesize > 0
-            written_bytes_size = client_socket.syswrite(payload)
+            written_bytes_size = @client_socket.syswrite(payload)
             payload = payload.byteslice(written_bytes_size..-1)
             @logger.trace(">transmitted #{written_bytes_size} bytes; #{payload.bytesize} bytes remain", socket: peer_info) if @logger.trace?
             sleep 0.1 unless payload.empty?
           end
         rescue => e
           log_warn "client socket failed:", e, host: @host, port: @port, socket: peer_info
-          client_socket.close rescue nil
-          client_socket = nil
+          @client_socket.close rescue nil
+          @client_socket = nil
           sleep @reconnect_interval
           retry
         end
@@ -210,13 +210,18 @@ def register
 
   # @overload Base#close
   def close
-    @closed.make_true
-    @server_socket.close rescue nil if @server_socket
+    if server?
+      # server-mode clean-up
+      @closed.make_true
+      @server_socket.shutdown rescue nil if @server_socket
 
-    return unless @client_threads
-    @client_threads.each do |thread|
-      client = thread[:client]
-      client.close rescue nil if client
+      @client_threads&.each do |thread|
+        client = thread[:client]
+        client.close rescue nil if client
+      end
+    else
+      # client-mode clean-up
+      @client_socket&.close
     end
   end
 

spec/outputs/tcp_spec.rb

@@ -3,7 +3,7 @@
 require "flores/pki"
 
 describe LogStash::Outputs::Tcp do
-  subject { described_class.new(config) }
+  subject(:instance) { described_class.new(config) }
   let(:config) { {
     "host" => "localhost",
     "port" => 2000 + rand(3000),
@@ -73,4 +73,121 @@
       end
     end
   end
+
+  ##
+  # Reads `in_io` until EOF and writes the bytes
+  # it receives to `out_io`, tolerating partial writes.
+  def siphon_until_eof(in_io, out_io)
+    while (buffer = in_io.read(16*1024)) do
+      while (buffer && !buffer.empty?) do
+        bytes_written = out_io.write(buffer)
+        buffer = buffer.byteslice(bytes_written..-1)
+      end
+    end
+  end
+
+  context 'client mode' do
+    context 'transmitting data' do
+      let!(:io) { StringIO.new } # somewhere for our server to stash the data it receives
+
+      let(:server_host) { 'localhost' }
+      let(:server_port) { server.addr[1] } # get actual since we bind to port 0
+
+      let!(:server) { TCPServer.new(server_host, 0) }
+
+      let(:config) do
+        { 'host' => server_host, 'port' => server_port, 'mode' => 'client' }
+      end
+
+      let(:event) { LogStash::Event.new({"hello" => "world"})}
+
+      subject(:instance) { described_class.new(config) }
+
+      before(:each) do
+        # accepts ONE connection
+        @server_socket_thread = Thread.start do
+          client = server.accept
+          siphon_until_eof(client, io)
+        end
+        instance.register
+      end
+
+      after(:each) do
+        @server_socket_thread&.join
+      end
+
+      it 'encodes and transmits data' do
+        instance.receive(event)
+        instance.close # release the connection
+        sleep 1
+        expect(io.string).to include('"hello"','"world"')
+      end
+
+      context 'when payload is very large' do
+        let(:one_hundred_megabyte_message) { "a" * 1024 * 1024 * 100 }
+        let(:event) { LogStash::Event.new("message" => one_hundred_megabyte_message) }
+
+
+        it 'encodes and transmits data' do
+          instance.receive(event)
+          instance.close # release the connection
+          sleep 1
+          expect(io.string).to include('"message"',%Q("#{one_hundred_megabyte_message}"))
+        end
+      end
+    end
+  end
+
+  context 'server mode' do
+    context 'transmitting data' do
+      let(:server_host) { 'localhost' }
+      let(:server_port) { Random.rand(1024...5000) }
+
+      let(:config) do
+        { 'host' => server_host, 'port' => server_port, 'mode' => 'server' }
+      end
+
+      subject(:instance) { described_class.new(config) }
+
+      before(:each) { instance.register } # start listener
+      after(:each) { instance.close }
+
+      let(:event) { LogStash::Event.new({"hello" => "world"})}
+
+      context 'when one client is connected' do
+        let(:io) { StringIO.new }
+        let(:client_socket) { TCPSocket.new(server_host, server_port) }
+
+        before(:each) do
+          @client_socket_thread = Thread.start { siphon_until_eof(client_socket, io) }
+          sleep 1 # wait for it to actually connect
+        end
+
+        it 'encodes and transmits data' do
+          instance.receive(event)
+
+          sleep 1 # wait for the event to get sent...
+          instance.close # release the connection
+
+          @client_socket_thread.join(30) || fail('client failed to join')
+          expect(io.string).to include('"hello"','"world"')
+        end
+
+        context 'when payload is very large' do
+          let(:one_hundred_megabyte_message) { "a" * 1024 * 1024 * 100 }
+          let(:event) { LogStash::Event.new("message" => one_hundred_megabyte_message) }
+
+          it 'encodes and transmits data' do
+          instance.receive(event)
+
+          sleep 1 # wait for the event to get sent...
+          instance.close # release the connection
+
+          @client_socket_thread.join(30) || fail('client failed to join')
+            expect(io.string).to include('"message"',%Q("#{one_hundred_megabyte_message}"))
+          end
+        end
+      end
+    end
+  end
 end