Rev 6194: Expose infrastructure so that we can test the Inet portion of the server. in http://bazaar.launchpad.net/~jameinel/bzr/2.5-soft-hangup-795025

[John Arbash Meinel]

At http://bazaar.launchpad.net/~jameinel/bzr/2.5-soft-hangup-795025

------------------------------------------------------------
revno: 6194
revision-id: john at arbash-meinel.com-20110923200950-eps56g6unln7s7qw
parent: john at arbash-meinel.com-20110923181128-d0qrj4u1celg9gkn
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 2.5-soft-hangup-795025
timestamp: Fri 2011-09-23 22:09:50 +0200
message:
  Expose infrastructure so that we can test the Inet portion of the server.
  
  It seems to be working, you should be able to SIGHUP while transferring, and it will
  'gently' shutdown.
-------------- next part --------------
=== modified file 'bzrlib/smart/medium.py'
--- a/bzrlib/smart/medium.py	2011-09-23 17:45:15 +0000
+++ b/bzrlib/smart/medium.py	2011-09-23 20:09:50 +0000
@@ -45,7 +45,7 @@
     ui,
     urlutils,
     )
-from bzrlib.smart import client, protocol, request, vfs
+from bzrlib.smart import client, protocol, request, signals, vfs
 from bzrlib.transport import ssh
 """)
 from bzrlib import osutils
@@ -246,6 +246,7 @@
         except Exception, e:
             stderr.write("%s terminating on exception %s\n" % (self, e))
             raise
+        self._disconnect_client()
 
     def _stop_gracefully(self):
         """When we finish this message, stop looking for more."""
@@ -430,6 +431,17 @@
         self._in = in_file
         self._out = out_file
 
+    def serve(self):
+        """See SmartServerStreamMedium.serve"""
+        # This is the regular serve, except it adds signal trapping for soft
+        # shutdown.
+        stop_gracefully = self._stop_gracefully
+        signals.register_on_hangup(id(self), stop_gracefully)
+        try:
+            return super(SmartServerPipeStreamMedium, self).serve()
+        finally:
+            signals.unregister_on_hangup(id(self))
+
     def _serve_one_request_unguarded(self, protocol):
         while True:
             # We need to be careful not to read past the end of the current

=== modified file 'bzrlib/smart/server.py'
--- a/bzrlib/smart/server.py	2011-09-23 18:11:28 +0000
+++ b/bzrlib/smart/server.py	2011-09-23 20:09:50 +0000
@@ -431,13 +431,17 @@
             transport = _mod_transport.get_transport_from_url(expand_userdirs.get_url())
         self.transport = transport
 
+    def _get_stdin_stdout(self):
+        return sys.stdin, sys.stdout
+
     def _make_smart_server(self, host, port, inet, timeout):
         if timeout is None:
             c = config.GlobalStack()
             timeout = c.get('serve.client_timeout')
         if inet:
+            stdin, stdout = self._get_stdin_stdout()
             smart_server = medium.SmartServerPipeStreamMedium(
-                sys.stdin, sys.stdout, self.transport, timeout=timeout)
+                stdin, stdout, self.transport, timeout=timeout)
         else:
             if host is None:
                 host = medium.BZR_DEFAULT_INTERFACE

=== modified file 'bzrlib/tests/test_smart_signals.py'
--- a/bzrlib/tests/test_smart_signals.py	2011-09-23 18:11:28 +0000
+++ b/bzrlib/tests/test_smart_signals.py	2011-09-23 20:09:50 +0000
@@ -15,12 +15,13 @@
 # Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
 
 
-import sys
+import os
 import signal
+import threading
 import weakref
 
-from bzrlib import tests
-from bzrlib.smart import signals
+from bzrlib import tests, transport
+from bzrlib.smart import client, medium, server, signals
 
 # Windows doesn't define SIGHUP. And while we could just skip a lot of these
 # tests, we often don't actually care about interaction with 'signal', so we
@@ -141,3 +142,50 @@
         self.assertIsNot(None, signals._on_sighup)
         signals.restore_sighup_handler(orig)
         self.assertIs(None, signals._on_sighup)
+
+
+class TestInetServer(tests.TestCase):
+
+    def create_file_pipes(self):
+        r, w = os.pipe()
+        rf = os.fdopen(r, 'rb')
+        wf = os.fdopen(w, 'wb')
+        return rf, wf
+
+    def test_inet_server_responds_to_sighup(self):
+        t = transport.get_transport('memory:///')
+        content = 'a'*1024*1024
+        t.put_bytes('bigfile', content)
+        factory = server.BzrServerFactory()
+        # Override stdin/stdout so that we can inject our own handles
+        client_read, server_write = self.create_file_pipes()
+        server_read, client_write = self.create_file_pipes()
+        factory._get_stdin_stdout = lambda: (server_read, server_write)
+        factory.set_up(t, None, None, inet=True, timeout=4.0)
+        self.addCleanup(factory.tear_down)
+        started = threading.Event()
+        stopped = threading.Event()
+        def serving():
+            started.set()
+            factory.smart_server.serve()
+            stopped.set()
+        server_thread = threading.Thread(target=serving)
+        server_thread.start()
+        started.wait()
+        client_medium = medium.SmartSimplePipesClientMedium(client_read,
+                            client_write, 'base')
+        client_client = client._SmartClient(client_medium)
+        resp, response_handler = client_client.call_expecting_body('get',
+            'bigfile')
+        signals._sighup_handler(SIGHUP, None)
+        self.assertTrue(factory.smart_server.finished)
+        # We can still finish reading the file content, but more than that, and
+        # the file is closed.
+        v = response_handler.read_body_bytes()
+        if v != content:
+            self.fail('Got the wrong content back, expected 1M "a"')
+        # writer is closed, so read returns nothing.
+        self.assertEqual('', client_read.read())
+        stopped.wait()
+        server_thread.join()
+