RSDK-5109 - add close function (#458)

purplenicole730 · web-flow · commit a372745b963a · 2023-10-12T14:52:40.000-04:00
diff --git a/docs/examples/example.ipynb b/docs/examples/example.ipynb
@@ -269,6 +269,9 @@
     "from typing import Any, Dict, Mapping, Optional\n",
     "\n",
     "from viam.components.sensor import Sensor\n",
+    "from viam.logging import getLogger\n",
+    "\n",
+    "LOGGER = getLogger(__name__)\n",
     "\n",
     "\n",
     "class MySensor(Sensor):\n",
@@ -282,6 +285,11 @@
     "    async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeout: Optional[float] = None) -> List[Geometry]:\n",
     "        raise NotImplementedError\n",
     "\n",
+    "    def close(self):\n",
+    "        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module\n",
+    "        # is shutting down.\n",
+    "        LOGGER.debug(f\"{self.name} is closed.\")\n",
+    "\n",
     "# Anything below this line is optional and will be replaced later, but may come in handy for debugging and testing.\n",
     "# To use, call `python wifi_sensor_module.py` in the command line while in the `src` directory.\n",
     "async def main():\n",
diff --git a/docs/examples/module_step2.py b/docs/examples/module_step2.py
@@ -4,12 +4,15 @@
 from typing_extensions import Self
 
 from viam.components.sensor import Sensor
+from viam.logging import getLogger
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import ResourceName
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
 
+LOGGER = getLogger(__name__)
+
 
 class MySensor(Sensor):
     # Subclass the Viam Sensor component and implement the required functions
@@ -26,6 +29,11 @@ async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -
         wifi_signal = [x for x in content[2].split(" ") if x != ""]
         return {"link": wifi_signal[2], "level": wifi_signal[3], "noise": wifi_signal[4]}
 
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
+
 
 async def main():
     Registry.register_resource_creator(Sensor.SUBTYPE, MySensor.MODEL, ResourceCreatorRegistration(MySensor.new))
diff --git a/docs/examples/module_step2_optional.py b/docs/examples/module_step2_optional.py
@@ -4,12 +4,15 @@
 from typing_extensions import Self
 
 from viam.components.sensor import Sensor
+from viam.logging import getLogger
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import ResourceName
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
 
+LOGGER = getLogger(__name__)
+
 
 class MySensor(Sensor):
     # Subclass the Viam Sensor component and implement the required functions
@@ -48,6 +51,11 @@ def reconfigure(self, config: ComponentConfig, dependencies: Mapping[ResourceNam
             multiplier = 1.0
         self.multiplier = multiplier
 
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
+
 
 async def main():
     Registry.register_resource_creator(Sensor.SUBTYPE, MySensor.MODEL, ResourceCreatorRegistration(MySensor.new, MySensor.validate_config))
diff --git a/docs/examples/module_step3.py b/docs/examples/module_step3.py
@@ -4,13 +4,16 @@
 from typing_extensions import Self
 
 from viam.components.sensor import Sensor
+from viam.logging import getLogger
 from viam.module.module import Module
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import ResourceName
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry, ResourceCreatorRegistration
 from viam.resource.types import Model, ModelFamily
 
+LOGGER = getLogger(__name__)
+
 
 class MySensor(Sensor):
     # Subclass the Viam Sensor component and implement the required functions
@@ -27,6 +30,11 @@ async def get_readings(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -
         wifi_signal = [x for x in content[2].split(" ") if x != ""]
         return {"link": wifi_signal[2], "level": wifi_signal[3], "noise": wifi_signal[4]}
 
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
+
 
 async def main():
     """
diff --git a/docs/examples/my_cool_arm.py b/docs/examples/my_cool_arm.py
@@ -5,9 +5,12 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from viam.components.arm import Arm, JointPositions, KinematicsFileFormat, Pose
+from viam.logging import getLogger
 from viam.operations import run_with_operation
 from viam.proto.common import Capsule, Geometry, Sphere
 
+LOGGER = getLogger(__name__)
+
 
 class MyCoolArm(Arm):
     # Subclass the Viam Arm component and implement the required functions
@@ -102,3 +105,8 @@ async def get_geometries(self, *, extra: Optional[Dict[str, Any]] = None, timeou
 
     async def get_kinematics(self, extra: Optional[Dict[str, Any]] = None, **kwargs) -> Tuple[KinematicsFileFormat.ValueType, bytes]:
         return KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, self.kinematics
+
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
diff --git a/examples/complex_module/src/arm/my_arm.py b/examples/complex_module/src/arm/my_arm.py
@@ -5,12 +5,15 @@
 from typing_extensions import Self
 
 from viam.components.arm import Arm, JointPositions, KinematicsFileFormat, Pose
+from viam.logging import getLogger
 from viam.operations import run_with_operation
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import Capsule, Geometry, ResourceName, Sphere
 from viam.resource.base import ResourceBase
 from viam.resource.types import Model, ModelFamily
 
+LOGGER = getLogger(__name__)
+
 
 class MyArm(Arm):
     # Subclass the Viam Arm component and implement the required functions
@@ -104,3 +107,8 @@ async def get_kinematics(self, extra: Optional[Dict[str, Any]] = None, **kwargs)
         with open(filepath, mode="rb") as f:
             file_data = f.read()
         return (KinematicsFileFormat.KINEMATICS_FILE_FORMAT_SVA, file_data)
+
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
diff --git a/examples/complex_module/src/gizmo/my_gizmo.py b/examples/complex_module/src/gizmo/my_gizmo.py
@@ -2,6 +2,7 @@
 
 from typing_extensions import Self
 
+from viam.logging import getLogger
 from viam.module.types import Reconfigurable
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import ResourceName
@@ -10,6 +11,8 @@
 
 from ..gizmo.api import Gizmo
 
+LOGGER = getLogger(__name__)
+
 
 class MyGizmo(Gizmo, Reconfigurable):
     """This is the specific implementation of a ``Gizmo`` (defined in api.py).
@@ -68,3 +71,8 @@ async def do_two(self, arg1: bool, **kwargs) -> str:
 
     def reconfigure(self, config: ComponentConfig, dependencies: Mapping[ResourceName, ResourceBase]):
         self.my_arg = config.attributes.fields["arg1"].string_value
+
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
diff --git a/examples/simple_module/src/main.py b/examples/simple_module/src/main.py
@@ -4,6 +4,7 @@
 from typing_extensions import Self
 
 from viam.components.sensor import Sensor
+from viam.logging import getLogger
 from viam.module.module import Module
 from viam.proto.app.robot import ComponentConfig
 from viam.proto.common import ResourceName
@@ -12,6 +13,8 @@
 from viam.resource.types import Model, ModelFamily
 from viam.utils import ValueTypes
 
+LOGGER = getLogger(__name__)
+
 
 class MySensor(Sensor):
     # Subclass the Viam Sensor component and implement the required functions
@@ -49,6 +52,11 @@ def reconfigure(self, config: ComponentConfig, dependencies: Mapping[ResourceNam
             multiplier = 1.0
         self.multiplier = multiplier
 
+    def close(self):
+        # This is a completely optional function to include. This will be called when the resource is removed from the config or the module
+        # is shutting down.
+        LOGGER.debug(f"{self.name} is closed.")
+
 
 async def main():
     """This function creates and starts a new module, after adding all desired resource models.
diff --git a/src/viam/module/module.py b/src/viam/module/module.py
@@ -150,7 +150,7 @@ async def reconfigure_resource(self, request: ReconfigureResourceRequest):
                 else:
                     resource.stop()
             add_request = AddResourceRequest(config=request.config, dependencies=request.dependencies)
-            self.server.remove_resource(rn)
+            await self.server.remove_resource(rn)
             await self.add_resource(add_request)
 
     async def remove_resource(self, request: RemoveResourceRequest):
@@ -161,7 +161,7 @@ async def remove_resource(self, request: RemoveResourceRequest):
                 await resource.stop()
             else:
                 resource.stop()
-        self.server.remove_resource(rn)
+        await self.server.remove_resource(rn)
 
     async def ready(self, request: ReadyRequest) -> ReadyResponse:
         self._parent_address = request.parent_address
diff --git a/src/viam/resource/base.py b/src/viam/resource/base.py
@@ -84,3 +84,12 @@ def get_operation(self, kwargs: Mapping[str, Any]) -> Operation:
             viam.operations.Operation: The operation associated with this function
         """
         return kwargs.get(Operation.ARG_NAME, Operation._noop())
+
+    async def close(self):
+        """Safely shut down the resource and prevent further use.
+
+        Close must be idempotent. Later configuration may allow a resource to be "open" again.
+        If a resource does not want or need a close function, it is assumed that the resource does not need to retun errors when future
+        non-Close methods are called.
+        """
+        return
diff --git a/src/viam/resource/manager.py b/src/viam/resource/manager.py
@@ -1,6 +1,7 @@
 from threading import RLock
 from typing import Dict, List, Type, TypeVar
 
+from viam.logging import getLogger
 from viam.proto.common import ResourceName
 from viam.resource.base import ResourceBase
 from viam.resource.registry import Registry
@@ -9,6 +10,7 @@
 from ..errors import DuplicateResourceError, ResourceNotFoundError
 from ..services.service_base import ServiceBase
 
+LOGGER = getLogger(__name__)
 ResourceType = TypeVar("ResourceType", bound=ResourceBase)
 
 
@@ -91,14 +93,31 @@ def get_resource(self, of_type: Type[ResourceType], name: ResourceName) -> Resou
                 return self.get_resource(of_type, self._short_to_long_name[name.name][0])
             raise ResourceNotFoundError(name.subtype, name.name)
 
-    def remove_resource(self, name: ResourceName):
+    async def remove_resource(self, name: ResourceName):
         """Remove the resource with the specified ```ResourceName```.
 
         Args:
             name (viam.proto.common.ResourceName): The ResourceName of the resource
         """
         with self._lock:
-            del self.resources[name]
+            try:
+                resource = self.resources[name]
+                await resource.close()
+            except Exception as e:
+                raise e
+            finally:
+                del self.resources[name]
+
+    async def close(self):
+        """Close the resourcce manager by removing all resources.
+        Please note that any errors will not raise an exception. Errors will still be logged."""
+        rns = [key for key in self.resources.keys()]
+        with self._lock:
+            for rn in rns:
+                try:
+                    await self.remove_resource(rn)
+                except Exception as e:
+                    LOGGER.error(f"Error while closing {rn.name}:", e)
 
     def _resource_by_name_only(self, name: str) -> ResourceBase:
         for rname, resource in self.resources.items():
diff --git a/src/viam/robot/client.py b/src/viam/robot/client.py
@@ -213,15 +213,15 @@ async def refresh(self):
                 if rname.subtype == Sensor.SUBTYPE.resource_subtype and MovementSensor.get_resource_name(rname.name) in resource_names:
                     continue
 
-                self._create_or_reset_client(rname)
+                await self._create_or_reset_client(rname)
 
             for rname in self.resource_names:
                 if rname not in resource_names:
-                    self._manager.remove_resource(rname)
+                    await self._manager.remove_resource(rname)
 
             self._resource_names = resource_names
 
-    def _create_or_reset_client(self, resourceName: ResourceName):
+    async def _create_or_reset_client(self, resourceName: ResourceName):
         if resourceName in self._manager.resources:
             res = self._manager.get_resource(ResourceBase, resourceName)
 
@@ -233,7 +233,7 @@ def _create_or_reset_client(self, resourceName: ResourceName):
             if isinstance(res, ReconfigurableResourceRPCClientBase):
                 res.reset_channel(self._channel)
             else:
-                self._manager.remove_resource(resourceName)
+                await self._manager.remove_resource(resourceName)
                 self._manager.register(
                     Registry.lookup_subtype(Subtype.from_resource_name(resourceName)).create_rpc_client(resourceName.name, self._channel)
                 )
diff --git a/src/viam/rpc/server.py b/src/viam/rpc/server.py
@@ -113,6 +113,7 @@ async def serve(
                 await self._server.start(host, port)
                 LOGGER.info(f"Serving on {host}:{port}")
             await self._server.wait_closed()
+            await self.close()
             LOGGER.debug("gRPC server closed")
 
     @classmethod
diff --git a/tests/mocks/module/gizmo/my_gizmo.py b/tests/mocks/module/gizmo/my_gizmo.py
@@ -15,6 +15,7 @@
 class MyGizmo(Gizmo, Reconfigurable):
     MODEL: ClassVar[Model] = Model(ModelFamily("acme", "demo"), "mygizmo")
     my_arg: str
+    closed: bool = False
 
     @classmethod
     def new(cls, config: ComponentConfig, dependencies: Mapping[ResourceName, ResourceBase]) -> Self:
@@ -59,3 +60,7 @@ async def do_two(self, arg1: bool, **kwargs) -> str:
 
     def reconfigure(self, config: ComponentConfig, dependencies: Mapping[ResourceName, ComponentBase]):
         self.my_arg = config.attributes.fields["arg1"].string_value
+
+    async def close(self):
+        self.closed = True
+        return
diff --git a/tests/test_module.py b/tests/test_module.py
@@ -146,15 +146,23 @@ async def test_add_resource_with_deps(self, robot_service: RobotService):  # noq
 
     @pytest.mark.asyncio
     async def test_remove_resource(self):
+        gizmo = self.module.server.get_resource(MyGizmo, Gizmo.get_resource_name("gizmo1"))
+        assert gizmo.closed is False
         assert Gizmo.get_resource_name("gizmo1") in self.module.server.resources
         req = RemoveResourceRequest(name="acme:component:gizmo/gizmo1")
         await self.module.remove_resource(req)
         assert Gizmo.get_resource_name("gizmo1") not in self.module.server.resources
+        assert gizmo.closed is True
 
-        assert SummationService.get_resource_name("mysum1") in self.module.server.resources
-        req = RemoveResourceRequest(name="acme:service:summation/mysum1")
-        await self.module.remove_resource(req)
-        assert SummationService.get_resource_name("mysum1") not in self.module.server.resources
+        with mock.patch("tests.mocks.module.summation.MySummationService.close") as mocked:
+            assert SummationService.get_resource_name("mysum1") in self.module.server.resources
+            req = RemoveResourceRequest(name="acme:service:summation/mysum1")
+
+            mocked.assert_not_called()
+            await self.module.remove_resource(req)
+            assert SummationService.get_resource_name("mysum1") not in self.module.server.resources
+            # test default close
+            mocked.assert_called_once()
 
     @pytest.mark.asyncio
     async def test_ready(self):
diff --git a/tests/test_resource_manager.py b/tests/test_resource_manager.py
diff --git a/tests/test_robot.py b/tests/test_robot.py
