Merge pull request #310 from ccatterina/patch-socket-recv

Use select to check if there is data available in the socket

Author: dhoomakethu

pymodbus/client/sync.py

@@ -1,4 +1,5 @@
 import socket
+import select
 import serial
 import time
 import sys
@@ -230,34 +231,43 @@ def _recv(self, size):
         """
         if not self.socket:
             raise ConnectionException(self.__str__())
-            # socket.recv(size) waits until it gets some data from the host but
-            # not necessarily the entire response that can be fragmented in
-            # many packets.
-            # To avoid the splitted responses to be recognized as invalid
-            # messages and to be discarded, loops socket.recv until full data
-            # is received or timeout is expired.
-            # If timeout expires returns the read data, also if its length is
-            # less than the expected size.
+
+        # socket.recv(size) waits until it gets some data from the host but
+        # not necessarily the entire response that can be fragmented in
+        # many packets.
+        # To avoid the splitted responses to be recognized as invalid
+        # messages and to be discarded, loops socket.recv until full data
+        # is received or timeout is expired.
+        # If timeout expires returns the read data, also if its length is
+        # less than the expected size.
         self.socket.setblocking(0)
-        begin = time.time()
 
-        data = b''
-        if size is not None:
-            while len(data) < size:
-                try:
-                    data += self.socket.recv(size - len(data))
-                except socket.error:
-                    pass
-                if not self.timeout or (time.time() - begin > self.timeout):
-                    break
+        timeout = self.timeout
+
+        # If size isn't specified read 1 byte at a time.
+        if size is None:
+            recv_size = 1
         else:
-            while True:
-                try:
-                    data += self.socket.recv(1)
-                except socket.error:
-                    pass
-                if not self.timeout or (time.time() - begin > self.timeout):
-                    break
+            recv_size = size
+
+        data = b''
+        begin = time.time()
+        while recv_size > 0:
+            ready = select.select([self.socket], [], [], timeout)
+            if ready[0]:
+                data += self.socket.recv(recv_size)
+
+            # If size isn't specified continue to read until timeout expires.
+            if size:
+                recv_size = size - len(data)
+
+            # Timeout is reduced also if some data has been received in order
+            # to avoid infinite loops when there isn't an expected response size
+            # and the slave sends noisy data continuosly.
+            timeout -= time.time() - begin
+            if timeout <= 0:
+                break
+
         return data
 
     def is_socket_open(self):

test/test_client_sync.py

@@ -166,10 +166,12 @@ def testSyncTcpClientInstantiation(self):
         client = ModbusTcpClient()
         self.assertNotEqual(client, None)
 
-    def testBasicSyncTcpClient(self):
+    @patch('pymodbus.client.sync.select')
+    def testBasicSyncTcpClient(self, mock_select):
         ''' Test the basic methods for the tcp sync client'''
 
         # receive/send
+        mock_select.select.return_value = [True]
         client = ModbusTcpClient()
         client.socket = mockSocket()
         self.assertEqual(0, client._send(None))
@@ -207,8 +209,11 @@ def testTcpClientSend(self):
         self.assertEqual(0, client._send(None))
         self.assertEqual(4, client._send('1234'))
 
-    def testTcpClientRecv(self):
+    @patch('pymodbus.client.sync.select')
+    def testTcpClientRecv(self, mock_select):
         ''' Test the tcp client receive method'''
+
+        mock_select.select.return_value = [True]
         client = ModbusTcpClient()
         self.assertRaises(ConnectionException, lambda: client._recv(1024))
 
@@ -223,10 +228,10 @@ def testTcpClientRecv(self):
         self.assertEqual(b'\x00\x01\x02', client._recv(3))
         mock_socket.recv.side_effect = iter([b'\x00', b'\x01', b'\x02'])
         self.assertEqual(b'\x00\x01', client._recv(2))
-        mock_socket.recv.side_effect = socket.error('No data')
+        mock_select.select.return_value = [False]
         self.assertEqual(b'', client._recv(2))
         client.socket = mockSocket()
-        client.socket.timeout = 0.1
+        mock_select.select.return_value = [True]
         self.assertIn(b'\x00', client._recv(None))
 
     def testSerialClientRpr(self):
@@ -351,4 +356,4 @@ def testSerialClientRepr(self):
 # Main
 # ---------------------------------------------------------------------------#
 if __name__ == "__main__":
-    unittest.main()
+    unittest.main()