From 0ccc4697a40f37653b694a6e08032dada9e29c9a Mon Sep 17 00:00:00 2001
From: Pavel Minaev <int19h@gmail.com>
Date: Wed, 7 Nov 2018 19:53:17 -0800
Subject: [PATCH] Add disconnect handler support to JsonMessageChannel.

Only wrap RequestFailure exceptions as fail responses when handling a request, to avoid inadvertently wrapping and swallowing logic errors.
---
 ptvsd/messaging.py | 28 ++++++++++++++++------------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/ptvsd/messaging.py b/ptvsd/messaging.py
index 0a7915f8..4d0dd582 100644
--- a/ptvsd/messaging.py
+++ b/ptvsd/messaging.py
@@ -262,8 +262,6 @@ class JsonMessageChannel(object):
 
     def __init__(self, stream, handlers=None, name='vsc_messaging'):
         self.stream = stream
-        self.send_callback = lambda channel, message: None
-        self.receive_callback = lambda channel, message: None
         self._lock = threading.Lock()
         self._stop = threading.Event()
         self._seq_iter = itertools.count(1)
@@ -293,7 +291,6 @@ class JsonMessageChannel(object):
         with self._lock:
             yield seq
             self.stream.write_json(message)
-        self.send_callback(self, message)
 
     def send_request(self, command, arguments=None):
         d = {'command': command}
@@ -327,7 +324,6 @@ class JsonMessageChannel(object):
             pass
 
     def on_message(self, message):
-        self.receive_callback(self, message)
         seq = message['seq']
         typ = message['type']
         if typ == 'request':
@@ -359,7 +355,7 @@ class JsonMessageChannel(object):
         request = Request(self, seq, command, arguments)
         try:
             response_body = handler(request)
-        except Exception as ex:
+        except RequestFailure as ex:
             self._send_response(seq, False, command, str(ex), None)
         else:
             if isinstance(response_body, Exception):
@@ -387,6 +383,14 @@ class JsonMessageChannel(object):
             body = RequestFailure(error_message)
         return request._handle_response(seq, command, body)
 
+    def on_disconnect(self):
+        # There's no more incoming messages, so any requests that are still pending
+        # must be marked as failed to unblock anyone waiting on them.
+        with self._lock:
+            for request in self._requests.values():
+                request._handle_response(None, request.command, EOFError('No response'))
+        getattr(self._handlers, 'disconnect', lambda: None)()
+
     def _process_incoming_messages(self):
         try:
             while True:
@@ -401,17 +405,17 @@ class JsonMessageChannel(object):
                     traceback.print_exc(file=sys.__stderr__)
                     raise
         finally:
-            # There's no more incoming messages, so any requests that are still pending
-            # must be marked as failed to unblock anyone waiting on them.
-            with self._lock:
-                for request in self._requests.values():
-                    request._handle_response(None, request.command, EOFError('No response'))
-
+            try:
+                self.on_disconnect()
+            except Exception:
+                print('Error while processing disconnect', file=sys.__stderr__)
+                traceback.print_exc(file=sys.__stderr__)
+                raise
 
 class MessageHandlers(object):
     """A simple delegating message handlers object for use with JsonMessageChannel.
     For every argument provided, the object has an attribute with the corresponding
-    name and value. Example:
+    name and value.
     """
 
     def __init__(self, **kwargs):