Merge pull request igrigorik#151 from greensnark/tcpsocket-read_nonbl…

…ock-fix

Improve TCPSocket read_nonblock & run TCP specs.

lib/em-synchrony/tcpsocket.rb

@@ -85,7 +85,13 @@ def write(msg)
       def read(num_bytes = nil, dest = nil)
         handle_read(:read, num_bytes, dest)
       end
-      alias_method :read_nonblock, :read
+
+      def read_nonblock(maxlen, dest = nil)
+        raise ArgumentError, "maxlen must be > 0" if !maxlen || maxlen <= 0
+        read_bytes = handle_read(:read_nonblock, maxlen, dest)
+        raise EOFError if read_bytes.nil?
+        read_bytes
+      end
 
       def recv(num_bytes, flags = 0)
         raise "Unknown flags in recv(): #{flags}" if flags.nonzero?
@@ -159,11 +165,14 @@ def handle_read(type, num_bytes, dest=nil)
         end
 
         def try_read_data
-          if @read_type == :read
+          if @read_type == :read || @read_type == :read_nonblock
+            nonblocking = @read_type == :read_nonblock
             unless @remote_closed
               if @read_bytes
                 # read(n) on an open socket, with >= than n buffered data, returns n data
-                if @in_buff.size >= @read_bytes then @in_buff.slice!(0, @read_bytes)
+                if (@in_buff.size >= @read_bytes ||
+                    (nonblocking && @in_buff.size > 0)) then
+                  @in_buff.slice!(0, @read_bytes)
                 # read(n) on an open socket, with < than n buffered data, blocks
                 else :block end
               else

spec/helper/all.rb

@@ -1,9 +1,4 @@
-require 'rubygems'
-require 'rspec'
-require 'pp'
-
-require 'lib/em-synchrony'
-require 'lib/em-synchrony/em-http'
+require 'spec/helper/core'
 require 'lib/em-synchrony/mysql2'
 require 'lib/em-synchrony/em-remcached'
 require 'lib/em-synchrony/em-memcache'

spec/helper/core.rb

@@ -0,0 +1,6 @@
+require 'rubygems'
+require 'rspec'
+require 'pp'
+
+require 'lib/em-synchrony'
+require 'lib/em-synchrony/em-http'

spec/tcpsocket_spec.rb

@@ -1,4 +1,4 @@
-require "spec/helper/all"
+require "spec/helper/core"
 
 module SendAndClose
   def post_init
@@ -21,15 +21,13 @@ def post_init
 end
 
 def tcp_test(server_type, ops={}, &block)
-  Proc.new do
-    EventMachine.synchrony do
-      ops = {:stop => true}.merge ops
-      EM::start_server('localhost', 12345, server_type)
-      @socket = EventMachine::Synchrony::TCPSocket.new 'localhost', 12345
-      @socket.close if ops[:close]
-      block.call
-      EM.stop if ops[:stop]
-    end
+  EventMachine.synchrony do
+    ops = {:stop => true}.merge ops
+    EM::start_server('localhost', 12345, server_type)
+    @socket = EventMachine::Synchrony::TCPSocket.new 'localhost', 12345
+    @socket.close if ops[:close]
+    block.call
+    EM.stop if ops[:stop]
   end
 end
 
@@ -231,6 +229,82 @@ def tcp_test(server_type, ops={}, &block)
     end
   end
 
+  context '#read_nonblock' do
+    context 'with a positive length argument' do
+      context 'when the connection is open' do
+        context 'with greater or equal than the requested data buffered' do
+          it 'returns the requested data and no more' do
+            tcp_test(SendAndKeepOpen) do
+              @socket.read_nonblock(2).size.should eq 2
+              @socket.read_nonblock(1).size.should eq 1
+            end
+          end
+        end
+        context 'with less than the requested data buffered' do
+          it 'returns the available data' do
+            tcp_test(SendAndKeepOpen) do
+              @socket.read_nonblock(10).size.should eq 4
+            end
+          end
+        end
+      end
+      context 'when the peer has closed the connection' do
+        context 'with no data buffered' do
+          it 'raises EOFError' do
+            tcp_test(SendAndClose) do
+              @socket.read_nonblock(4).size.should eq 4
+              lambda {
+                @socket.read_nonblock(1)
+              }.should raise_error(EOFError)
+            end
+          end
+        end
+        context 'with less than the requested data buffered' do
+          it 'returns the buffered data' do
+            tcp_test(SendAndClose) do
+              @socket.read_nonblock(50).size.should eq 4
+            end
+          end
+        end
+        context 'with greater or equal than the requested data buffered' do
+          it 'returns the requested data and no more' do
+            tcp_test(SendAndClose) do
+              @socket = EventMachine::Synchrony::TCPSocket.new 'localhost', 12345
+              @socket.read_nonblock(2).size.should eq 2
+            end
+          end
+        end
+      end
+      context 'when we closed the connection' do
+        it 'raises IOError' do
+          tcp_test(SendAndKeepOpen, :close => true) do
+            proc {
+              @socket.read_nonblock(4)
+            }.should raise_error(IOError)
+          end
+        end
+      end
+    end
+    context 'with a negative length argument' do
+      it 'raises ArgumentError' do
+        tcp_test(SendAndKeepOpen) do
+          proc {
+            @socket.read_nonblock(-10)
+          }.should raise_error(ArgumentError)
+        end
+      end
+    end
+    context 'with a zero length argument' do
+      it 'raises ArgumentError' do
+        tcp_test(SendAndKeepOpen) do
+          proc {
+            @socket.read_nonblock(0)
+          }.should raise_error(ArgumentError)
+        end
+      end
+    end
+  end
+
   context '#recv' do
     context 'with a length argument' do
       context 'with a possitive length argument' do