Give correct functional signature to sendfile and add live-throughput

test
epio · Nov 3, 2011 · 1c699f1 · 1c699f1
1 parent a3b9e64
commit 1c699f1
Show file tree

Hide file tree

Showing 4 changed files with 74 additions and 16 deletions.
diff --git a/mantrid/actions.py b/mantrid/actions.py
@@ -55,7 +55,7 @@ def __init__(self, balancer, host, matched_host, type=None):
     try:
         import ctypes
         _sendfile = ctypes.CDLL("libc.so.6").sendfile
-        _sendfile.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_off_t, ctypes.c_size_t]
+        _sendfile.argtypes = [ctypes.c_int, ctypes.c_int, ctypes.c_long, ctypes.c_size_t]
         _sendfile.restype = ctypes.c_ssize_t
     except Exception:
         _sendfile = None
@@ -74,8 +74,9 @@ def handle(self, sock, read_data, path, headers):
                 self._sendfile(sock.fileno(), fh.fileno(), 0, os.fstat(fh.fileno()).st_size)
             except (TypeError, AttributeError):
                 sock.sendall(fh.read())
-            # Close the file
+            # Close the file and socket
             fh.close()
+            sock.close()
         except socket.error, e:
             if e.errno != errno.EPIPE:
                 raise

diff --git a/mantrid/loadbalancer.py b/mantrid/loadbalancer.py
@@ -218,12 +218,15 @@ def management_loop(self, address, family):
         # Actually serve management
         logging.info("Listening for management on %s" % (address, ))
         management_app = ManagementApp(self)
-        with open("/dev/null", "w") as log_dest:
-            wsgi.server(
-                sock,
-                management_app.handle,
-                log = log_dest,
-            )
+        try:
+            with open("/dev/null", "w") as log_dest:
+                wsgi.server(
+                    sock,
+                    management_app.handle,
+                    log = log_dest,
+                )
+        finally:
+            sock.close()
 
     ### Client handling ###
 
@@ -246,11 +249,14 @@ def listen_loop(self, address, family, internal=False):
         eventlet.sleep(0.5)
         # Start serving
         logging.info("Listening for requests on %s" % (address, ))
-        eventlet.serve(
-            sock,
-            lambda sock, addr: self.handle(sock, addr, internal),
-            concurrency = 10000,
-        )
+        try:
+            eventlet.serve(
+                sock,
+                lambda sock, addr: self.handle(sock, addr, internal),
+                concurrency = 10000,
+            )
+        finally:
+            sock.close()
 
     def resolve_host(self, host, protocol="http"):
         # Special case for empty hosts dict

diff --git a/mantrid/tests/__init__.py b/mantrid/tests/__init__.py
@@ -1,3 +1,3 @@
-from .actions import ActionTests
+from .actions import ActionTests, LiveActionTests
 from .loadbalancer import BalancerTests
 from .client import ClientTests
diff --git a/mantrid/tests/actions.py b/mantrid/tests/actions.py
@@ -3,8 +3,10 @@
 import socket
 import time
 import eventlet
+import unittest
+httplib2 = eventlet.import_patched("httplib2")
 from eventlet.timeout import Timeout
-from unittest import TestCase
+from ..loadbalancer import Balancer
 from ..actions import Empty, Static, Unknown, NoHosts, Redirect, Proxy, Spin
 
 
@@ -32,6 +34,9 @@ def send(self, data):
     def sendall(self, data):
         self.data += data
 
+    def close(self):
+        pass
+
 
 class MockErrorSocket(object):
     "Fake Socket class that raises a specific error message on use."
@@ -44,7 +49,7 @@ def _error(self, *args, **kwargs):
     sendall = _error
 
 
-class ActionTests(TestCase):
+class ActionTests(unittest.TestCase):
     "Tests the various actions"
 
     def test_empty(self):
@@ -188,3 +193,49 @@ def test_socket_errors(self):
             with self.assertRaises(socket.error) as cm:
                 action.handle(sock, "", "/", {})
             self.assertEqual(cm.exception.errno, errno.EBADF)
+
+
+class LiveActionTests(unittest.TestCase):
+    """
+    Tests that the client/API work correctly.
+    """
+
+    next_port = 30300
+
+    def setUp(self):
+        self.__class__.next_port += 3
+        self.balancer = Balancer(
+            [(("0.0.0.0", self.next_port), socket.AF_INET)],
+            [(("0.0.0.0", self.next_port + 1), socket.AF_INET)],
+            [(("0.0.0.0", self.next_port + 2), socket.AF_INET)],
+            "/tmp/mantrid-test-state-2",
+        )
+        self.balancer_thread = eventlet.spawn(self.balancer.run)
+        eventlet.sleep(0.1)
+        self.balancer.hosts = {
+            "test-host.com": ["static", {"type": "test"}, True],
+        }
+
+    def tearDown(self):
+        self.balancer.running = False
+        self.balancer_thread.kill()
+        eventlet.sleep(0.1)
+
+    def test_unknown(self):
+        # Send a HTTP request to the balancer, ensure the response
+        # is the same as the "unknown" template
+        h = httplib2.Http()
+        resp, content = h.request(
+            "http://127.0.0.1:%i" % self.next_port,
+            "GET",
+        )
+        self.assertEqual(
+            '503',
+            resp['status'],
+        )
+        expected_content = open(os.path.join(os.path.dirname(__file__), "..", "static", "unknown.http")).read()
+        expected_content = expected_content[expected_content.index("\r\n\r\n") + 4:]
+        self.assertEqual(
+            expected_content,
+            content,
+        )