Add an optional timeout to dmypy_server

This allows it ot shut down after a period of inactivity to avoid
hogging memory.

Also don't crash the server on an error communicating with the client.

Author: msullivan

mypy/dmypy.py

@@ -30,13 +30,17 @@
 start_parser = p = subparsers.add_parser('start', help="Start daemon")
 p.add_argument('--log-file', metavar='FILE', type=str,
                help="Direct daemon stdout/stderr to FILE")
+p.add_argument('--timeout', metavar='TIMEOUT', type=int,
+               help="Server shutdown timeout (in seconds)")
 p.add_argument('flags', metavar='FLAG', nargs='*', type=str,
                help="Regular mypy flags (precede with --)")
 
 restart_parser = p = subparsers.add_parser('restart',
     help="Restart daemon (stop or kill followed by start)")
 p.add_argument('--log-file', metavar='FILE', type=str,
                help="Direct daemon stdout/stderr to FILE")
+p.add_argument('--timeout', metavar='TIMEOUT', type=int,
+               help="Server shutdown timeout (in seconds)")
 p.add_argument('flags', metavar='FLAG', nargs='*', type=str,
                help="Regular mypy flags (precede with --)")
 
@@ -63,6 +67,8 @@
 hang_parser = p = subparsers.add_parser('hang', help="Hang for 100 seconds")
 
 daemon_parser = p = subparsers.add_parser('daemon', help="Run daemon in foreground")
+p.add_argument('--timeout', metavar='TIMEOUT', type=int,
+               help="Server shutdown timeout (in seconds)")
 p.add_argument('flags', metavar='FLAG', nargs='*', type=str,
                help="Regular mypy flags (precede with --)")
 
@@ -148,7 +154,8 @@ def start_server(args: argparse.Namespace) -> None:
     """Start the server from command arguments and wait for it."""
     # Lazy import so this import doesn't slow down other commands.
     from mypy.dmypy_server import daemonize, Server, process_start_options
-    if daemonize(Server(process_start_options(args.flags)).serve, args.log_file) != 0:
+    if daemonize(Server(process_start_options(args.flags), timeout=args.timeout).serve,
+                 args.log_file) != 0:
         sys.exit(1)
     wait_for_server()
 
@@ -284,7 +291,7 @@ def do_daemon(args: argparse.Namespace) -> None:
     """Serve requests in the foreground."""
     # Lazy import so this import doesn't slow down other commands.
     from mypy.dmypy_server import Server, process_start_options
-    Server(process_start_options(args.flags)).serve()
+    Server(process_start_options(args.flags), timeout=args.timeout).serve()
 
 
 @action(help_parser)

mypy/dmypy_server.py

@@ -109,11 +109,14 @@ class Server:
     # NOTE: the instance is constructed in the parent process but
     # serve() is called in the grandchild (by daemonize()).
 
-    def __init__(self, options: Options, alt_lib_path: Optional[str] = None) -> None:
+    def __init__(self, options: Options,
+                 timeout: Optional[int] = None,
+                 alt_lib_path: Optional[str] = None) -> None:
         """Initialize the server with the desired mypy flags."""
         self.saved_cache = {}  # type: mypy.build.SavedCache
         self.fine_grained = options.fine_grained_incremental
         self.options = options
+        self.timeout = timeout
         self.alt_lib_path = alt_lib_path
         self.fine_grained_manager = None  # type: Optional[FineGrainedBuildManager]
 
@@ -134,13 +137,19 @@ def serve(self) -> None:
         """Serve requests, synchronously (no thread or fork)."""
         try:
             sock = self.create_listening_socket()
+            if self.timeout is not None:
+                sock.settimeout(self.timeout)
             try:
                 with open(STATUS_FILE, 'w') as f:
                     json.dump({'pid': os.getpid(), 'sockname': sock.getsockname()}, f)
                     f.write('\n')  # I like my JSON with trailing newline
                 while True:
                     conn, addr = sock.accept()
-                    data = receive(conn)
+                    try:
+                        data = receive(conn)
+                    except OSError as err:
+                        conn.close()  # Maybe the client hung up
+                        continue
                     resp = {}  # type: Dict[str, Any]
                     if 'command' not in data:
                         resp = {'error': "No command found in request"}
@@ -159,12 +168,15 @@ def serve(self) -> None:
                     if command == 'stop':
                         sock.close()
                         sys.exit(0)
+            except socket.timeout:
+                print("Exiting due to inactivity.")
+                sys.exit(0)
             finally:
                 os.unlink(STATUS_FILE)
         finally:
             os.unlink(self.sockname)
             exc_info = sys.exc_info()
-            if exc_info[0]:
+            if exc_info[0] and exc_info[0] is not SystemExit:
                 traceback.print_exception(*exc_info)  # type: ignore
 
     def create_listening_socket(self) -> socket.socket: