fix graceful restart and add unit test

Author: clowwindy

shadowsocks/eventloop.py

@@ -192,14 +192,15 @@ def stop(self):
     def run(self):
         events = []
         while not self._stopping:
-            now = time.time()
+            asap = False
             try:
                 events = self.poll(TIMEOUT_PRECISION)
             except (OSError, IOError) as e:
                 if errno_from_exception(e) in (errno.EPIPE, errno.EINTR):
                     # EPIPE: Happens when the client closes the connection
                     # EINTR: Happens when received a signal
                     # handles them as soon as possible
+                    asap = True
                     logging.debug('poll:%s', e)
                 else:
                     logging.error('poll:%s', e)
@@ -214,7 +215,8 @@ def run(self):
                         handler.handle_event(sock, fd, event)
                     except (OSError, IOError) as e:
                         shell.print_exception(e)
-            if now - self._last_time >= TIMEOUT_PRECISION:
+            now = time.time()
+            if asap or now - self._last_time >= TIMEOUT_PRECISION:
                 for callback in self._periodic_callbacks:
                     callback()
                 self._last_time = now

shadowsocks/tcprelay.py

@@ -689,18 +689,19 @@ def handle_event(self, sock, fd, event):
                 logging.warn('poll removed fd')
 
     def handle_periodic(self):
-        self._sweep_timeout()
         if self._closed:
             if self._server_socket:
                 self._eventloop.remove(self._server_socket, self)
-                self._eventloop.remove_periodic(self.handle_periodic)
                 self._server_socket.close()
                 self._server_socket = None
-                logging.info('closed listen port %d', self._listen_port)
+                logging.info('closed TCP port %d', self._listen_port)
             if not self._fd_to_handlers:
+                logging.info('stopping')
                 self._eventloop.stop()
+        self._sweep_timeout()
 
     def close(self, next_tick=False):
+        logging.debug('TCP close')
         self._closed = True
         if not next_tick:
             if self._eventloop:

shadowsocks/udprelay.py

@@ -272,14 +272,18 @@ def handle_event(self, sock, fd, event):
             self._handle_client(sock)
 
     def handle_periodic(self):
+        if self._closed:
+            if self._server_socket:
+                logging.info('closed UDP port %d', self._listen_port)
+                self._server_socket.close()
+                self._server_socket = None
+                for sock in self._sockets:
+                    sock.close()
         self._cache.sweep()
         self._client_fd_to_server_addr.sweep()
-        if self._closed:
-            self._server_socket.close()
-            for sock in self._sockets:
-                sock.close()
 
     def close(self, next_tick=False):
+        logging.debug('UDP close')
         self._closed = True
         if not next_tick:
             if self._eventloop:

tests/graceful.json

@@ -0,0 +1,10 @@
+{
+    "server":"127.0.0.1",
+    "server_port":8387,
+    "local_port":1081,
+    "password":"aes_password",
+    "timeout":15,
+    "method":"aes-256-cfb",
+    "local_address":"127.0.0.1",
+    "fast_open":false
+}

tests/graceful_cli.py

@@ -0,0 +1,18 @@
+#!/usr/bin/python
+
+import socket
+import socks
+import time
+
+
+SERVER_IP = '127.0.0.1'
+SERVER_PORT = 8001
+
+
+if __name__ == '__main__':
+    s = socks.socksocket()
+    s.set_proxy(socks.SOCKS5, SERVER_IP, 1081)
+    s.connect((SERVER_IP, SERVER_PORT))
+    s.send(b'test')
+    time.sleep(30)
+    s.close()

tests/graceful_server.py

@@ -0,0 +1,13 @@
+#!/usr/bin/python
+
+import socket
+
+
+if __name__ == '__main__':
+    s = socket.socket()
+    s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+    s.bind(('127.0.0.1', 8001))
+    s.listen(1024)
+    c = None
+    while True:
+        c = s.accept()

tests/jenkins.sh

@@ -69,6 +69,7 @@ if [ -f /proc/sys/net/ipv4/tcp_fastopen ] ; then
 fi
 
 run_test tests/test_large_file.sh
+run_test tests/test_graceful_restart.sh
 run_test tests/test_udp_src.sh
 run_test tests/test_command.sh
 

tests/test_graceful_restart.sh

@@ -0,0 +1,63 @@
+#!/bin/bash
+
+PYTHON="coverage run -p -a"
+URL=http://127.0.0.1/file
+
+
+# setup processes
+$PYTHON shadowsocks/local.py -c tests/graceful.json &
+LOCAL=$!
+
+$PYTHON shadowsocks/server.py -c tests/graceful.json --forbidden-ip "" &
+SERVER=$!
+
+python tests/graceful_server.py &
+GSERVER=$!
+
+sleep 1
+
+python tests/graceful_cli.py &
+GCLI=$!
+
+sleep 1
+
+# graceful restart server: send SIGQUIT to old process and start a new one
+kill -s SIGQUIT $SERVER
+$PYTHON shadowsocks/server.py -c tests/graceful.json --forbidden-ip "" &
+NEWSERVER=$!
+
+sleep 1
+
+# check old server
+ps x | grep -v grep | grep $SERVER
+OLD_SERVER_RUNNING1=$?
+# old server should not quit at this moment
+echo old server running: $OLD_SERVER_RUNNING1
+
+sleep 1
+
+# close connections on old server
+kill -s SIGINT $GCLI
+kill -s SIGKILL $GSERVER
+kill -s SIGINT $LOCAL
+
+sleep 11
+
+# check old server
+ps x | grep -v grep | grep $SERVER
+OLD_SERVER_RUNNING2=$?
+# old server should quit at this moment
+echo old server running: $OLD_SERVER_RUNNING2
+
+# new server is expected running
+kill -s SIGINT $NEWSERVER || exit 1
+
+if [ $OLD_SERVER_RUNNING1 -ne 0 ]; then
+    exit 1
+fi
+
+if [ $OLD_SERVER_RUNNING2 -ne 1 ]; then
+    kill -s SIGINT $SERVER
+    sleep 1
+    exit 1
+fi