Merge pull request #539 from dellis23/memcached-recv

memcached: call socket.recv multiple times to get all stats

Author: josegonzalez

src/collectors/memcached/memcached.py

@@ -87,13 +87,26 @@ def get_raw_stats(self, host, port):
             else:
                 sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                 sock.connect((host, int(port)))
+
+            # give up after a reasonable amount of time
+            sock.settimeout(3)
+
             # request stats
             sock.send('stats\n')
-            # something big enough to get whatever is sent back
-            data = sock.recv(4096)
+
+            # stats can be sent across multiple packets, so make sure we've
+            # read up until the END marker
+            while True:
+                received = sock.recv(4096)
+                if not received:
+                    break
+                data += received
+                if data.endswith('END\r\n'):
+                    break
         except socket.error:
             self.log.exception('Failed to get stats from %s:%s',
                                host, port)
+        sock.close()
         return data
 
     def get_stats(self, host, port):

src/collectors/memcached/test/testmemcached.py

@@ -5,6 +5,7 @@
 from test import CollectorTestCase
 from test import get_collector_config
 from test import unittest
+from mock import MagicMock
 from mock import Mock
 from mock import patch
 
@@ -27,6 +28,15 @@ def setUp(self):
     def test_import(self):
         self.assertTrue(MemcachedCollector)
 
+    @patch('socket.socket')
+    def test_get_raw_stats_works_across_packet_boundaries(self, socket_mock):
+        socket_instance = MagicMock()
+        socket_mock.return_value = socket_instance
+        stats_packets = ['stat foo 1\r\n', 'END\r\n']
+        socket_instance.recv.side_effect = stats_packets
+        stats = self.collector.get_raw_stats('', None)
+        self.assertEqual(stats, ''.join(stats_packets))
+
     @patch.object(Collector, 'publish')
     def test_should_work_with_real_data(self, publish_mock):
         patch_raw_stats = patch.object(