Handle exit explicitly. (#499)

(fixes gh-459)

Our exit handling was all over the place. I've refactored the code to properly separate concerns. One consequence is that this fixes wait-on-exit.

ptvsd/_util.py

@@ -51,10 +51,44 @@ def call_all(callables, *args, **kwargs):
 # threading stuff
 
 try:
-    TimeoutError = __builtins__.TimeoutError
+    ThreadError = threading.ThreadError
 except AttributeError:
-    class TimeoutError(OSError):
-        """Timeout expired."""
+    ThreadError = RuntimeError
+
+
+try:
+    base = __builtins__.TimeoutError
+except AttributeError:
+    base = OSError
+class TimeoutError(base):  # noqa
+    """Timeout expired."""
+    timeout = None
+    reason = None
+
+    @classmethod
+    def from_timeout(cls, timeout, reason=None):
+        """Return a TimeoutError with the given timeout."""
+        msg = 'timed out (after {} seconds)'.format(timeout)
+        if reason is not None:
+            msg += ' ' + reason
+        self = cls(msg)
+        self.timeout = timeout
+        self.reason = reason
+        return self
+del base  # noqa
+
+
+def wait(check, timeout=None, reason=None):
+    """Wait for the given func to return True.
+
+    If a timeout is given and reached then raise TimeoutError.
+    """
+    if timeout is None or timeout <= 0:
+        while not check():
+            time.sleep(0.01)
+    else:
+        if not _wait(check, timeout):
+            raise TimeoutError.from_timeout(timeout, reason)
 
 
 def is_locked(lock):
@@ -73,18 +107,18 @@ def lock_release(lock):
         return
     try:
         lock.release()
-    except RuntimeError:  # already unlocked
+    except ThreadError:  # already unlocked
         pass
 
 
-def lock_wait(lock, timeout=None):
+def lock_wait(lock, timeout=None, reason='waiting for lock'):
     """Wait until the lock is not locked."""
     if not _lock_acquire(lock, timeout):
-        raise TimeoutError
+        raise TimeoutError.from_timeout(timeout, reason)
     lock_release(lock)
 
 
-if sys.version_info > (2,):
+if sys.version_info >= (3,):
     def _lock_acquire(lock, timeout):
         if timeout is None:
             timeout = -1
@@ -93,14 +127,21 @@ def _lock_acquire(lock, timeout):
     def _lock_acquire(lock, timeout):
         if timeout is None or timeout <= 0:
             return lock.acquire()
-        if lock.acquire(False):
+
+        def check():
+            return lock.acquire(False)
+        return _wait(check, timeout)
+
+
+def _wait(check, timeout):
+    if check():
+        return True
+    for _ in range(int(timeout * 100)):
+        time.sleep(0.01)
+        if check():
             return True
-        for _ in range(int(timeout * 100)):
-            time.sleep(0.01)
-            if lock.acquire(False):
-                return True
-        else:
-            return False
+    else:
+        return False
 
 
 ########################

ptvsd/daemon.py

@@ -9,7 +9,8 @@
     ExitHandlers, UnsupportedSignalError,
     kill_current_proc)
 from .session import PyDevdDebugSession
-from ._util import ClosedError, NotRunningError, ignore_errors, debug
+from ._util import (
+    ClosedError, NotRunningError, ignore_errors, debug, lock_wait)
 
 
 def _wait_for_user():
@@ -283,18 +284,26 @@ def _stop_quietly(self):
         with ignore_errors():
             self._stop()
 
-    def _handle_session_closing(self, session, can_disconnect=None):
-        debug('handling closing session')
+    def _handle_session_disconnecting(self, session):
+        debug('handling disconnecting session')
+        if self._singlesession:
+            if self._killonclose:
+                with self._lock:
+                    if not self._exiting_via_atexit_handler:
+                        # Ensure the proc is exiting before closing
+                        # socket.  Note that we kill the proc instead
+                        # of calling sys.exit(0).
+                        # Note that this will trigger either the atexit
+                        # handler or the signal handler.
+                        kill_current_proc()
+            else:
+                try:
+                    self.close()
+                except DaemonClosedError:
+                    pass
 
-        if self._exiting_via_atexit_handler:
-            # This must be done before we send a disconnect response
-            # (which implies before we close the client socket).
-            # TODO: Call session.wait_on_exit() directly?
-            wait_on_exit = session.get_wait_on_exit()
-            if wait_on_exit(self.exitcode or 0):
-                self._wait_for_user()
-        if can_disconnect is not None:
-            can_disconnect()
+    def _handle_session_closing(self, session):
+        debug('handling closing session')
 
         if self._singlesession:
             if self._killonclose:
@@ -332,6 +341,7 @@ def _bind_session(self, session):
         session = self.SESSION.from_raw(
             session,
             notify_closing=self._handle_session_closing,
+            notify_disconnecting=self._handle_session_disconnecting,
             ownsock=True,
         )
         self._session = session
@@ -366,13 +376,8 @@ def _release_session(self):
             # TODO: This shouldn't happen if we are exiting?
             self._session = None
 
-        # Possibly trigger VSC "exited" and "terminated" events.
-        exitcode = self.exitcode
-        if exitcode is None:
-            if self._exiting_via_atexit_handler or self._singlesession:
-                exitcode = 0
         try:
-            session.stop(exitcode)
+            session.stop()
         except NotRunningError:
             pass
         try:
@@ -408,6 +413,22 @@ def _handle_atexit(self):
         with self._lock:
             self._exiting_via_atexit_handler = True
         session = self.session
+
+        if session is not None:
+            lock = threading.Lock()
+            lock.acquire()
+
+            def wait_debugger(timeout=None):
+                lock_wait(lock, timeout)
+
+            def wait_exiting(cfg):
+                if cfg:
+                    self._wait_for_user()
+                lock.release()
+            # TODO: Rely on self._stop_debugger().
+            session.handle_debugger_stopped(wait_debugger)
+            session.handle_exiting(self.exitcode, wait_exiting)
+
         try:
             self.close()
         except DaemonClosedError:

ptvsd/session.py

@@ -30,15 +30,20 @@ def from_server_socket(cls, server, **kwargs):
         client, _ = server.accept()
         return cls(client, ownsock=True, **kwargs)
 
-    def __init__(self, sock, notify_closing=None, ownsock=False):
+    def __init__(self, sock, notify_closing=None, notify_disconnecting=None,
+                 ownsock=False):
         super(DebugSession, self).__init__()
 
         if notify_closing is not None:
             def handle_closing(before):
                 if before:
-                    notify_closing(self, can_disconnect=self._can_disconnect)
+                    notify_closing(self)
             self.add_close_handler(handle_closing)
 
+        if notify_disconnecting is None:
+            notify_disconnecting = (lambda _: None)
+        self._notify_disconnecting = notify_disconnecting
+
         self._sock = sock
         if ownsock:
             # Close the socket *after* calling sys.exit() (via notify_closing).
@@ -55,7 +60,6 @@ def handle_closing(before):
             self.add_close_handler(handle_closing)
 
         self._msgprocessor = None
-        self._can_disconnect = None
 
     @property
     def socket(self):
@@ -65,12 +69,19 @@ def socket(self):
     def msgprocessor(self):
         return self._msgprocessor
 
-    def wait_options(self):
-        """Return (normal, abnormal) based on the session's launch config."""
+    def handle_debugger_stopped(self, wait=None):
+        """Deal with the debugger exiting."""
         proc = self._msgprocessor
         if proc is None:
-            return (False, False)
-        return proc._wait_options()
+            return
+        proc.handle_debugger_stopped(wait)
+
+    def handle_exiting(self, exitcode=None, wait=None):
+        """Deal with the debuggee exiting."""
+        proc = self._msgprocessor
+        if proc is None:
+            return
+        proc.handle_exiting(exitcode, wait)
 
     def wait_until_stopped(self):
         """Block until all resources (e.g. message processor) have stopped."""
@@ -80,18 +91,6 @@ def wait_until_stopped(self):
         # TODO: Do this in VSCodeMessageProcessor.close()?
         proc._wait_for_server_thread()
 
-    def get_wait_on_exit(self):
-        """Return a wait_on_exit(exitcode) func.
-
-        The func returns True if process should wait for the user
-        before exiting.
-        """
-        normal, abnormal = self.wait_options()
-
-        def wait_on_exit(exitcode):
-            return (normal and not exitcode) or (abnormal and exitcode)
-        return wait_on_exit
-
     # internal methods
 
     def _new_msg_processor(self, **kwargs):
@@ -109,13 +108,16 @@ def _start(self, threadname, **kwargs):
         self._msgprocessor.start(threadname)
         return self._msgprocessor_running
 
-    def _stop(self, exitcode=None):
-        if self._msgprocessor is None:
+    def _stop(self):
+        proc = self._msgprocessor
+        if proc is None:
             return
 
         debug('proc stopping')
-        self._msgprocessor.handle_session_stopped(exitcode)
-        self._msgprocessor.close()
+        # TODO: We should not need to wait if not exiting.
+        # The editor will send a "disconnect" request at this point.
+        proc._wait_for_disconnect()
+        proc.close()
         self._msgprocessor = None
 
     def _close(self):
@@ -130,18 +132,17 @@ def _msgprocessor_running(self):
 
     # internal methods for VSCodeMessageProcessor
 
-    def _handle_vsc_disconnect(self, can_disconnect=None):
+    def _handle_vsc_disconnect(self):
         debug('disconnecting')
-        self._can_disconnect = can_disconnect
+        self._notify_disconnecting(self)
+
+    def _handle_vsc_close(self):
+        debug('processor closing')
         try:
             self.close()
         except ClosedError:
             pass
 
-    def _handle_vsc_close(self):
-        debug('processor closing')
-        self.close()
-
 
 class PyDevdDebugSession(DebugSession):
     """A single DAP session for a network client socket."""