grpc added client disconnect detection for streams

daemon/core/grpc/server.py

@@ -2,7 +2,7 @@ import logging
 import os
 import tempfile
 import time
-from Queue import Queue
+from Queue import Queue, Empty
 from itertools import repeat
 
 import grpc
@@ -377,8 +377,9 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
         queue = Queue()
         session.node_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            node = queue.get()
+            try:
+                node = queue.get(timeout=1)
                 node_event = core_pb2.NodeEvent()
                 update_proto(
                     node_event.node,
@@ -394,14 +395,17 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                 services = node.services or ""
                 node_event.node.services.extend(services.split("|"))
                 yield node_event
+            except Empty:
+                continue
 
     def LinkEvents(self, request, context):
         session = self.get_session(request.id, context)
         queue = Queue()
         session.link_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            event = queue.get()
+            try:
+                event = queue.get(timeout=1)
                 link_event = core_pb2.LinkEvent()
                 if event.interface1_id is not None:
                     interface_one = link_event.link.interface_one
@@ -451,14 +455,17 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                     unidirectional=event.unidirectional
                 )
                 yield link_event
+            except Empty:
+                continue
 
     def SessionEvents(self, request, context):
         session = self.get_session(request.id, context)
         queue = Queue()
         session.event_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            event = queue.get()
+            try:
+                event = queue.get(timeout=1)
                 session_event = core_pb2.SessionEvent()
                 event_time = event.time
                 if event_time is not None:
@@ -473,14 +480,17 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                     session=session.session_id
                 )
                 yield session_event
+            except Empty:
+                continue
 
     def ConfigEvents(self, request, context):
         session = self.get_session(request.id, context)
         queue = Queue()
         session.config_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            event = queue.get()
+            try:
+                event = queue.get(timeout=1)
                 config_event = core_pb2.ConfigEvent()
                 update_proto(
                     config_event,
@@ -500,14 +510,17 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                 )
                 config_event.data_types.extend(event.data_types)
                 yield config_event
+            except Empty:
+                continue
 
     def ExceptionEvents(self, request, context):
         session = self.get_session(request.id, context)
         queue = Queue()
         session.exception_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            event = queue.get()
+            try:
+                event = queue.get(timeout=1)
                 exception_event = core_pb2.ExceptionEvent()
                 event_time = event.date
                 if event_time is not None:
@@ -523,14 +536,17 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                     opaque=event.opaque
                 )
                 yield exception_event
+            except Empty:
+                continue
 
     def FileEvents(self, request, context):
         session = self.get_session(request.id, context)
         queue = Queue()
         session.file_handlers.append(lambda x: queue.put(x))
 
-        while True:
+        while context.is_active():
-            event = queue.get()
+            try:
+                event = queue.get(timeout=1)
                 file_event = core_pb2.FileEvent()
                 update_proto(
                     file_event,
@@ -546,6 +562,8 @@ class CoreApiServer(core_pb2_grpc.CoreApiServicer):
                     compressed_data=event.compressed_data
                 )
                 yield file_event
+            except Empty:
+                continue
 
     def CreateNode(self, request, context):
         session = self.get_session(request.session, context)