Graceful shutdown; close #179

setup.py

@@ -6,7 +6,7 @@ with open('README.rst') as f:
 
 setup(
     name="shadowsocks",
-    version="2.2.0",
+    version="2.2.1",
     license='MIT',
     description="A fast tunnel proxy that help you get through firewalls",
     author='clowwindy',

shadowsocks/asyncdns.py

@@ -324,7 +324,7 @@ class DNSResolver(object):
                                    socket.SOL_UDP)
         self._sock.setblocking(False)
         loop.add(self._sock, eventloop.POLL_IN)
-        loop.add_handler(self.handle_events)
+        loop.add_handler(self.handle_events, ref=False)
 
     def _call_callback(self, hostname, ip, error=None):
         callbacks = self._hostname_to_cb.get(hostname, [])

shadowsocks/eventloop.py

@@ -168,7 +168,7 @@ class EventLoop(object):
                             'package')
         self._fd_to_f = {}
         self._handlers = []
-        self.stopping = False
+        self._ref_handlers = []
         logging.debug('using event model: %s', model)
 
     def poll(self, timeout=None):
@@ -189,17 +189,24 @@ class EventLoop(object):
         fd = f.fileno()
         self._impl.modify_fd(fd, mode)
 
-    def add_handler(self, handler):
+    def add_handler(self, handler, ref=True):
         self._handlers.append(handler)
+        if ref:
+            # when all ref handlers are removed, loop stops
+            self._ref_handlers.append(handler)
+
+    def remove_handler(self, handler):
+        self._handlers.remove(handler)
+        if handler in self._ref_handlers:
+            self._ref_handlers.remove(handler)
 
     def run(self):
-        while not self.stopping:
+        while self._ref_handlers:
             try:
                 events = self.poll(1)
             except (OSError, IOError) as e:
-                if errno_from_exception(e) == errno.EPIPE:
+                if errno_from_exception(e) in (errno.EPIPE, errno.EINTR):
                     # Happens when the client closes the connection
-                    logging.error('poll:%s', e)
                     continue
                 else:
                     logging.error('poll:%s', e)

shadowsocks/server.py

@@ -30,6 +30,7 @@ import eventloop
 import tcprelay
 import udprelay
 import asyncdns
+import signal
 
 
 def main():
@@ -67,13 +68,14 @@ def main():
         udp_servers.append(udprelay.UDPRelay(a_config, dns_resolver, False))
 
     def run_server():
+        def child_handler(signum, _):
+            logging.warn('received SIGQUIT, doing graceful shutting down..')
+            map(lambda s: s.close(next_tick=True), tcp_servers + udp_servers)
+        signal.signal(signal.SIGQUIT, child_handler)
         try:
             loop = eventloop.EventLoop()
             dns_resolver.add_to_loop(loop)
-            for tcp_server in tcp_servers:
-                tcp_server.add_to_loop(loop)
-            for udp_server in udp_servers:
-                udp_server.add_to_loop(loop)
+            map(lambda s: s.add_to_loop(loop), tcp_servers + udp_servers)
             loop.run()
         except (KeyboardInterrupt, IOError, OSError) as e:
             logging.error(e)
@@ -97,11 +99,13 @@ def main():
             if not is_child:
                 def handler(signum, _):
                     for pid in children:
+                        try:
                             os.kill(pid, signum)
-                        os.waitpid(pid, 0)
+                        except OSError:  # child may already exited
+                            pass
                     sys.exit()
-                import signal
                 signal.signal(signal.SIGTERM, handler)
+                signal.signal(signal.SIGQUIT, handler)
 
                 # master
                 for a_tcp_server in tcp_servers:

shadowsocks/tcprelay.py

@@ -631,7 +631,15 @@ class TCPRelay(object):
         if now - self._last_time > TIMEOUT_PRECISION:
             self._sweep_timeout()
             self._last_time = now
+            if self._closed:
+                if self._server_socket:
+                    self._eventloop.remove(self._server_socket)
+                    self._server_socket.close()
+                    self._server_socket = None
+                if not self._fd_to_handlers:
+                    self._eventloop.remove_handler(self._handle_events)
 
-    def close(self):
+    def close(self, next_tick=False):
         self._closed = True
+        if not next_tick:
             self._server_socket.close()

shadowsocks/udprelay.py

@@ -260,12 +260,17 @@ class UDPRelay(object):
                     logging.error('UDP client_socket err')
                 self._handle_client(sock)
         now = time.time()
-        if now - self._last_time > 3.5:
+        if now - self._last_time > 3:
             self._cache.sweep()
-        if now - self._last_time > 7:
             self._client_fd_to_server_addr.sweep()
             self._last_time = now
+            if self._closed:
+                self._server_socket.close()
+                for sock in self._sockets:
+                    sock.close()
+                self._eventloop.remove_handler(self._handle_events)
 
-    def close(self):
+    def close(self, next_tick=False):
         self._closed = True
+        if not next_tick:
             self._server_socket.close()

shadowsocks/utils.py

@@ -90,7 +90,7 @@ def get_config(is_local):
         longopts = ['fast-open']
     else:
         shortopts = 'hs:p:k:m:c:t:vq'
-        longopts = ['fast-open', 'workers:']
+        longopts = ['fast-open', 'workers=']
     try:
         config_path = find_config()
         optlist, args = getopt.getopt(sys.argv[1:], shortopts, longopts)
@@ -134,7 +134,7 @@ def get_config(is_local):
             elif key == '--fast-open':
                 config['fast_open'] = True
             elif key == '--workers':
-                config['workers'] = value
+                config['workers'] = int(value)
             elif key == '-h':
                 if is_local:
                     print_local_help()