Bound camera RPC waits and feed shutdown

anki_vector/camera.py

@@ -29,15 +29,12 @@
            "CameraComponent", "CameraConfig", "CameraImage"]
 
 import asyncio
-from concurrent.futures import CancelledError
+from concurrent import futures
 import io
-import time
 import sys
+import time
 
-from . import annotate, connection, util
-from .events import Events
-from .exceptions import VectorCameraFeedException, VectorCameraImageCaptureException
-from .messaging import protocol
+import grpc
 
 try:
     import numpy as np
@@ -49,6 +46,16 @@
 except ImportError:
     sys.exit("Cannot import from PIL: Do `pip3 install --user Pillow` to install")
 
+from . import annotate, connection, util
+from .events import Events
+from .exceptions import (VectorCameraFeedException,
+                         VectorCameraImageCaptureException,
+                         VectorTimeoutException)
+from .messaging import protocol
+
+
+_CAMERA_RPC_TIMEOUT = 10.0
+
 
 def _convert_to_pillow_image(image_data: bytes) -> Image.Image:
     """Convert raw image bytes to a Pillow Image."""
@@ -278,6 +285,7 @@ def __init__(self, robot):
         self._latest_image: CameraImage = None
         self._latest_image_id: int = None
         self._camera_feed_task: asyncio.Task = None
+        self._camera_feed_stream = None
         self._enabled = False
         self._config = None  # type CameraConfig
         self._gain = 0.0
@@ -414,42 +422,50 @@ def init_camera_feed(self) -> None:
             self._enabled = True
             self._camera_feed_task = self.conn.loop.create_task(self._request_and_handle_images())
 
-    def close_camera_feed(self) -> None:
-        """Cancel camera feed task."""
+    def close_camera_feed(self, timeout: float = _CAMERA_RPC_TIMEOUT) -> None:
+        """Cancel the camera feed task within a bounded time.
+
+        :param timeout: Number of seconds to wait for the stream to close.
+        """
+        if timeout <= 0:
+            raise ValueError('timeout must be greater than zero')
         if self._camera_feed_task:
+            deadline = time.monotonic() + timeout
             self._enabled = False
-            self._camera_feed_task.cancel()
+            if self._camera_feed_stream is not None:
+                self._camera_feed_stream.cancel()
+            else:
+                self._camera_feed_task.cancel()
             future = self.conn.run_coroutine(self._camera_feed_task)
             try:
-                future.result()
-            except CancelledError:
-                self.logger.debug('Camera feed task was cancelled. This is expected during disconnection.')
-            # wait for streaming to end, up to 10 seconds
-            iterations = 0
-            max_iterations = 100
-            while self.image_streaming_enabled():
-                time.sleep(0.1)
-                iterations += 1
-                if iterations > max_iterations:
-                    # leave loop, even if streaming is still enabled
-                    # because other SDK functions will still work and
-                    # the RPC should have had enough time to finish
-                    # which means we _should_ be in a good state.
-                    self.logger.info('Camera Feed closed, but streaming on'
-                                     ' robot remained enabled.  This is unexpected.')
+                future.result(timeout=timeout)
+            except futures.CancelledError:
+                self.logger.debug(
+                    'Camera feed task was cancelled. This is expected during '
+                    'disconnection.')
+            except futures.TimeoutError as exc:
+                future.cancel()
+                raise VectorTimeoutException(None) from exc
+            while True:
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise VectorTimeoutException(None)
+                if not self.image_streaming_enabled(timeout=min(1.0, remaining)):
                     break
+                time.sleep(min(0.1, remaining))
             self._camera_feed_task = None
 
-    async def _image_streaming_enabled(self) -> bool:
+    async def _image_streaming_enabled(self, timeout: float) -> bool:
         """request streaming enabled status from the robot"""
         request = protocol.IsImageStreamingEnabledRequest()
-        response = await self.conn.grpc_interface.IsImageStreamingEnabled(request)
+        response = await self.conn.grpc_interface.IsImageStreamingEnabled(
+            request, timeout=timeout)
         enabled = False
         if response:
             enabled = response.is_image_streaming_enabled
         return enabled
 
-    def image_streaming_enabled(self) -> bool:
+    def image_streaming_enabled(self, timeout: float = _CAMERA_RPC_TIMEOUT) -> bool:
         """True if image streaming is enabled on the robot
 
         .. testcode::
@@ -461,9 +477,17 @@ def image_streaming_enabled(self) -> bool:
                     print("Robot is streaming video")
                 else:
                     print("Robot is not streaming video")
+
+        :param timeout: Number of seconds to wait for the status RPC.
         """
-        future = self.conn.run_coroutine(self._image_streaming_enabled())
-        return future.result()
+        if timeout <= 0:
+            raise ValueError('timeout must be greater than zero')
+        future = self.conn.run_coroutine(self._image_streaming_enabled(timeout))
+        try:
+            return future.result(timeout=timeout)
+        except futures.TimeoutError as exc:
+            future.cancel()
+            raise VectorTimeoutException(None) from exc
 
     def _unpack_image(self, msg: protocol.CameraFeedResponse) -> None:
         """Processes raw data from the robot into a more useful image structure."""
@@ -484,20 +508,36 @@ def _unpack_image(self, msg: protocol.CameraFeedResponse) -> None:
     async def _request_and_handle_images(self) -> None:
         """Queries and listens for camera feed events from the robot.
         Received events are parsed by a helper function."""
+        stream = None
         try:
             req = protocol.CameraFeedRequest()
-            async for evt in self.grpc_interface.CameraFeed(req):
+            stream = self.grpc_interface.CameraFeed(req)
+            self._camera_feed_stream = stream
+            async for evt in stream:
                 # If the camera feed is disabled after stream is setup, exit the stream
                 # (the camera feed on the robot is disabled internally on stream exit)
                 if not self._enabled:
                     self.logger.warning('Camera feed has been disabled. Enable the feed to start/continue receiving camera feed data')
                     return
                 self._unpack_image(evt)
-        except CancelledError:
-            self.logger.debug('Camera feed task was cancelled. This is expected during disconnection.')
+        except asyncio.CancelledError:
+            self.logger.debug(
+                'Camera feed task was cancelled. This is expected during '
+                'disconnection.')
+            raise
+        except grpc.RpcError as exc:
+            if exc.code() != grpc.StatusCode.CANCELLED:
+                raise
+        finally:
+            self._camera_feed_stream = None
+            if stream is not None:
+                stream.cancel()
 
     @connection.on_connection_thread()
-    async def capture_single_image(self, enable_high_resolution: bool = False) -> CameraImage:
+    async def capture_single_image(
+            self,
+            enable_high_resolution: bool = False,
+            timeout: float = _CAMERA_RPC_TIMEOUT) -> CameraImage:
         """Request to capture a single image from the robot's camera.
 
         This call requests the robot to capture an image and returns the
@@ -516,14 +556,17 @@ async def capture_single_image(self, enable_high_resolution: bool = False) -> Ca
 
         :param enable_high_resolution: Enable/disable request for high resolution images. The default resolution
                                        is 640x360, while the high resolution is 1280x720.
+        :param timeout: Number of seconds to wait for the gRPC response.
         """
+        if timeout <= 0:
+            raise ValueError('timeout must be greater than zero')
         if self._enabled:
             self.logger.warning('Camera feed is enabled. Receiving image from the feed at default resolution.')
             return self._latest_image
         if enable_high_resolution:
             self.logger.warning('Capturing a high resolution (1280*720) image. Image events for this frame need to be scaled.')
         req = protocol.CaptureSingleImageRequest(enable_high_resolution=enable_high_resolution)
-        res = await self.grpc_interface.CaptureSingleImage(req)
+        res = await self.grpc_interface.CaptureSingleImage(req, timeout=timeout)
         if res and res.data:
             image = _convert_to_pillow_image(res.data)
             return CameraImage(image, self._image_annotator, res.image_id)