From 4686b0e5dd56d32241de672dad9c51b16a038a50 Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Tue, 20 Apr 2021 12:07:28 -0400
Subject: [PATCH 1/6] fix(api): do not cache removed modules in thread_manager

Closes #5359
---
 api/src/opentrons/drivers/temp_deck/driver.py | 26 ++++++++++++-------
 .../opentrons/drivers/thermocycler/driver.py  | 10 ++++---
 api/src/opentrons/hardware_control/api.py     |  2 +-
 .../hardware_control/modules/magdeck.py       |  2 +-
 .../hardware_control/modules/mod_abc.py       |  8 ++++++
 .../hardware_control/modules/tempdeck.py      |  5 ++--
 .../hardware_control/modules/thermocycler.py  |  3 +++
 .../hardware_control/thread_manager.py        | 18 +++++++++----
 8 files changed, 53 insertions(+), 21 deletions(-)

diff --git a/api/src/opentrons/drivers/temp_deck/driver.py b/api/src/opentrons/drivers/temp_deck/driver.py
index cf37d52882a..0df44e54ce9 100644
--- a/api/src/opentrons/drivers/temp_deck/driver.py
+++ b/api/src/opentrons/drivers/temp_deck/driver.py
@@ -210,17 +210,25 @@ def update_temperature(self, default=None) -> str:
             updated_temperature = default or self._temperature.copy()
             self._temperature.update(updated_temperature)
         else:
-            # comment
-            try:
-                self._update_thread = Thread(
-                    target=self._recursive_update_temperature,
-                    args=[DEFAULT_COMMAND_RETRIES],
-                    name='Tempdeck recursive update temperature')
-                self._update_thread.start()
-            except (TempDeckError, SerialException, SerialNoResponse) as e:
-                return str(e)
+            def _update():
+                try:
+                    self._recursive_update_temperature(retries=DEFAULT_COMMAND_RETRIES)
+                except (OSError, TempDeckError, SerialException, SerialNoResponse):
+                    log.exception("UPDATE FAILED")
+
+            self._update_thread = Thread(
+                target=_update(),
+                name='Tempdeck recursive update temperature')
+            self._update_thread.start()
         return ''
 
+    def _update_temperature(self):
+        self._update_thread = Thread(
+            target=self._recursive_update_temperature,
+            args=[DEFAULT_COMMAND_RETRIES],
+            name='Tempdeck recursive update temperature')
+        self._update_thread.start()
+
     @property
     def target(self) -> Optional[int]:
         return self._temperature.get('target')
diff --git a/api/src/opentrons/drivers/thermocycler/driver.py b/api/src/opentrons/drivers/thermocycler/driver.py
index 2e99a5d6f3f..d8decbbdec9 100644
--- a/api/src/opentrons/drivers/thermocycler/driver.py
+++ b/api/src/opentrons/drivers/thermocycler/driver.py
@@ -331,10 +331,12 @@ def send(self, command, callback):
             self._send_write_fd.write(b'c')
 
     def close(self):
+        log.debug("Halting TCPoller")
         self._halt_write_fd.write(b'q')
 
     def __del__(self):
         """ Clean up thread fifos"""
+        log.debug("Cleaning up thread fifos in TCPoller.")
         try:
             os.unlink(self._send_path)
         except NameError:
@@ -376,9 +378,10 @@ async def connect(self, port: str) -> 'Thermocycler':
         return self
 
     def disconnect(self) -> 'Thermocycler':
-        if self.is_connected():
+        if self.is_connected() or self._poller:
             self._poller.close()  # type: ignore
             self._poller.join()  # type: ignore
+            log.debug("TC poller stopped.")
         self._poller = None
         return self
 
@@ -627,6 +630,7 @@ async def enter_programming_mode(self):
 
     def __del__(self):
         try:
-            self._poller.close()  # type: ignore
+            if self._poller:
+                self._poller.close()  # type: ignore
         except Exception:
-            log.exception('Exception while cleaning up Thermocycler:')
+            log.exception('Exception while cleaning up Thermocycler')
diff --git a/api/src/opentrons/hardware_control/api.py b/api/src/opentrons/hardware_control/api.py
index c1c49bf2024..51d31cf9c94 100644
--- a/api/src/opentrons/hardware_control/api.py
+++ b/api/src/opentrons/hardware_control/api.py
@@ -1733,7 +1733,7 @@ def _unregister_modules(self,
         for removed_mod in removed_modules:
             self._log.info(f"Module {removed_mod.name()} detached"
                            f" from port {removed_mod.port}")
-            del removed_mod
+            removed_mod.cleanup()
 
     async def register_modules(
             self,
diff --git a/api/src/opentrons/hardware_control/modules/magdeck.py b/api/src/opentrons/hardware_control/modules/magdeck.py
index b5e3b840e4f..f6d60ad37d5 100644
--- a/api/src/opentrons/hardware_control/modules/magdeck.py
+++ b/api/src/opentrons/hardware_control/modules/magdeck.py
@@ -209,7 +209,7 @@ def _disconnect(self):
         if self._driver:
             self._driver.disconnect(port=self._port)
 
-    def __del__(self):
+    def cleanup(self) -> None:
         self._disconnect()
 
     async def prep_for_update(self) -> str:
diff --git a/api/src/opentrons/hardware_control/modules/mod_abc.py b/api/src/opentrons/hardware_control/modules/mod_abc.py
index a2edb1f9054..a6496e05a19 100644
--- a/api/src/opentrons/hardware_control/modules/mod_abc.py
+++ b/api/src/opentrons/hardware_control/modules/mod_abc.py
@@ -162,3 +162,11 @@ def name(cls) -> str:
     def bootloader(cls) -> UploadFunction:
         """ Method used to upload file to this module's bootloader. """
         pass
+
+    def cleanup(self) -> None:
+        """ Clean up the module instance.
+
+        Clean up, i.e. stop pollers, disconnect serial, etc in preparation for
+        object destruction.
+        """
+        pass
diff --git a/api/src/opentrons/hardware_control/modules/tempdeck.py b/api/src/opentrons/hardware_control/modules/tempdeck.py
index 10b32b88a5b..fb36fd97536 100644
--- a/api/src/opentrons/hardware_control/modules/tempdeck.py
+++ b/api/src/opentrons/hardware_control/modules/tempdeck.py
@@ -242,8 +242,9 @@ async def _connect(self):
         self._poller = Poller(self._driver)
         self._poller.start()
 
-    def __del__(self):
-        if hasattr(self, '_poller') and self._poller:
+    def cleanup(self) -> None:
+        if self._poller:
+            log.debug("Stopping tempdeck poller.")
             self._poller.stop()
 
     async def prep_for_update(self) -> str:
diff --git a/api/src/opentrons/hardware_control/modules/thermocycler.py b/api/src/opentrons/hardware_control/modules/thermocycler.py
index 172bac17271..37b30057e8a 100644
--- a/api/src/opentrons/hardware_control/modules/thermocycler.py
+++ b/api/src/opentrons/hardware_control/modules/thermocycler.py
@@ -347,3 +347,6 @@ async def prep_for_update(self):
         new_port = await update.find_bootloader_port()
 
         return new_port or self.port
+
+    def cleanup(self) -> None:
+        self._driver.disconnect()
diff --git a/api/src/opentrons/hardware_control/thread_manager.py b/api/src/opentrons/hardware_control/thread_manager.py
index a3ad835e83c..2b498c53b4b 100644
--- a/api/src/opentrons/hardware_control/thread_manager.py
+++ b/api/src/opentrons/hardware_control/thread_manager.py
@@ -4,7 +4,7 @@
 import logging
 import asyncio
 import functools
-from typing import Generic, TypeVar, Any, Optional
+from typing import Generic, TypeVar, Any, Optional, Dict
 from .adapters import SynchronousAdapter
 from .modules.mod_abc import AbstractModule
 
@@ -108,7 +108,7 @@ def __init__(self, builder, *args, **kwargs):
         self._sync_managed_obj: Optional[SynchronousAdapter] = None
         is_running = threading.Event()
         self._is_running = is_running
-
+        self._cached_modules: Dict[AbstractModule, CallBridger[AbstractModule]] = {}
         # TODO: remove this if we switch to python 3.8
         # https://docs.python.org/3/library/asyncio-subprocess.html#subprocess-and-threads  # noqa
         # On windows, the event loop and system interface is different and
@@ -175,10 +175,9 @@ def clean_up(self):
             loop.call_soon_threadsafe(loop.stop)
         except Exception:
             pass
-        object.__getattribute__(self, 'wrap_module').cache_clear()
+        self._cached_modules = {}
         object.__getattribute__(self, '_thread').join()
 
-    @functools.lru_cache(8)
     def wrap_module(
             self, module: AbstractModule) -> CallBridger[AbstractModule]:
         return CallBridger(module, object.__getattribute__(self, '_loop'))
@@ -193,7 +192,16 @@ def __getattribute__(self, attr_name):
             wrap = object.__getattribute__(self, 'wrap_module')
             managed = object.__getattribute__(self, 'managed_obj')
             attr = getattr(managed, attr_name)
-            return [wrap(mod) for mod in attr]
+            cached_mods = object.__getattribute__(self, '_cached_modules')
+
+            # Update self._cached_modules to delete all removed modules' entries and add
+            # newly attached modules. Removing references to stale instances
+            # is necessary to allow the garbage collector to delete those objects from
+            # memory and cleanly stop all the threads associated with them.
+            cached_mods = {
+                    module: cached_mods.get(module, wrap(module)) for module in attr}
+            object.__setattr__(self, '_cached_modules', cached_mods)
+            return cached_mods.values()
         elif attr_name == 'clean_up':
             # the wrapped object probably has this attr as well as us, and we
             # want to call both, with the wrapped one first

From bc94f86b6226aa71c6253d02ff56ff1e1bc41f12 Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Tue, 20 Apr 2021 12:26:09 -0400
Subject: [PATCH 2/6] remove unused function

---
 api/src/opentrons/drivers/temp_deck/driver.py | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/api/src/opentrons/drivers/temp_deck/driver.py b/api/src/opentrons/drivers/temp_deck/driver.py
index 0df44e54ce9..3a29f40918e 100644
--- a/api/src/opentrons/drivers/temp_deck/driver.py
+++ b/api/src/opentrons/drivers/temp_deck/driver.py
@@ -222,13 +222,6 @@ def _update():
             self._update_thread.start()
         return ''
 
-    def _update_temperature(self):
-        self._update_thread = Thread(
-            target=self._recursive_update_temperature,
-            args=[DEFAULT_COMMAND_RETRIES],
-            name='Tempdeck recursive update temperature')
-        self._update_thread.start()
-
     @property
     def target(self) -> Optional[int]:
         return self._temperature.get('target')

From 1c2e28f6a920ba78feb61744b2ac47eade419546 Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Tue, 20 Apr 2021 12:29:16 -0400
Subject: [PATCH 3/6] improve exception msg

---
 api/src/opentrons/drivers/temp_deck/driver.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/api/src/opentrons/drivers/temp_deck/driver.py b/api/src/opentrons/drivers/temp_deck/driver.py
index 3a29f40918e..f57330e9d97 100644
--- a/api/src/opentrons/drivers/temp_deck/driver.py
+++ b/api/src/opentrons/drivers/temp_deck/driver.py
@@ -214,7 +214,7 @@ def _update():
                 try:
                     self._recursive_update_temperature(retries=DEFAULT_COMMAND_RETRIES)
                 except (OSError, TempDeckError, SerialException, SerialNoResponse):
-                    log.exception("UPDATE FAILED")
+                    log.exception("Failed to execute _recursive_update_temperature.")
 
             self._update_thread = Thread(
                 target=_update(),

From b5cdd31705b8b17dec2eff816dc3f59d93f25a51 Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Tue, 20 Apr 2021 12:49:18 -0400
Subject: [PATCH 4/6] remove TC driver destructor, cleanup

---
 api/src/opentrons/drivers/thermocycler/driver.py           | 7 -------
 api/src/opentrons/hardware_control/modules/thermocycler.py | 5 ++++-
 api/src/opentrons/hardware_control/thread_manager.py       | 4 ++--
 3 files changed, 6 insertions(+), 10 deletions(-)

diff --git a/api/src/opentrons/drivers/thermocycler/driver.py b/api/src/opentrons/drivers/thermocycler/driver.py
index d8decbbdec9..889662878fc 100644
--- a/api/src/opentrons/drivers/thermocycler/driver.py
+++ b/api/src/opentrons/drivers/thermocycler/driver.py
@@ -627,10 +627,3 @@ async def enter_programming_mode(self):
         await asyncio.sleep(0.05)
         trigger_connection.close()
         self.disconnect()
-
-    def __del__(self):
-        try:
-            if self._poller:
-                self._poller.close()  # type: ignore
-        except Exception:
-            log.exception('Exception while cleaning up Thermocycler')
diff --git a/api/src/opentrons/hardware_control/modules/thermocycler.py b/api/src/opentrons/hardware_control/modules/thermocycler.py
index 37b30057e8a..aa622a0da03 100644
--- a/api/src/opentrons/hardware_control/modules/thermocycler.py
+++ b/api/src/opentrons/hardware_control/modules/thermocycler.py
@@ -349,4 +349,7 @@ async def prep_for_update(self):
         return new_port or self.port
 
     def cleanup(self) -> None:
-        self._driver.disconnect()
+        try:
+            self._driver.disconnect()
+        except Exception:
+            MODULE_LOG.exception('Exception while cleaning up Thermocycler')
diff --git a/api/src/opentrons/hardware_control/thread_manager.py b/api/src/opentrons/hardware_control/thread_manager.py
index 2b498c53b4b..d85129fa011 100644
--- a/api/src/opentrons/hardware_control/thread_manager.py
+++ b/api/src/opentrons/hardware_control/thread_manager.py
@@ -196,8 +196,8 @@ def __getattribute__(self, attr_name):
 
             # Update self._cached_modules to delete all removed modules' entries and add
             # newly attached modules. Removing references to stale instances
-            # is necessary to allow the garbage collector to delete those objects from
-            # memory and cleanly stop all the threads associated with them.
+            # is necessary in order to allow the garbage collector to delete
+            # those objects from memory and cleanly stop all associated threads.
             cached_mods = {
                     module: cached_mods.get(module, wrap(module)) for module in attr}
             object.__setattr__(self, '_cached_modules', cached_mods)

From a0916f6eb9a6ef357f646a95caaad2e6aec169cd Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Wed, 21 Apr 2021 11:25:38 -0400
Subject: [PATCH 5/6] address PR comments, don't use self. in thread_manager.
 Some debug lines to remove

---
 api/src/opentrons/drivers/thermocycler/driver.py       | 6 +++---
 api/src/opentrons/hardware_control/modules/tempdeck.py | 4 ++--
 api/src/opentrons/hardware_control/thread_manager.py   | 4 ++--
 api/tests/opentrons/hardware_control/test_modules.py   | 3 +++
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/api/src/opentrons/drivers/thermocycler/driver.py b/api/src/opentrons/drivers/thermocycler/driver.py
index 889662878fc..0803a6df04f 100644
--- a/api/src/opentrons/drivers/thermocycler/driver.py
+++ b/api/src/opentrons/drivers/thermocycler/driver.py
@@ -331,12 +331,12 @@ def send(self, command, callback):
             self._send_write_fd.write(b'c')
 
     def close(self):
-        log.debug("Halting TCPoller")
+        log.info("===== Halting TCPoller")
         self._halt_write_fd.write(b'q')
 
     def __del__(self):
         """ Clean up thread fifos"""
-        log.debug("Cleaning up thread fifos in TCPoller.")
+        log.info("===== Cleaning up thread fifos in TCPoller.")
         try:
             os.unlink(self._send_path)
         except NameError:
@@ -381,7 +381,7 @@ def disconnect(self) -> 'Thermocycler':
         if self.is_connected() or self._poller:
             self._poller.close()  # type: ignore
             self._poller.join()  # type: ignore
-            log.debug("TC poller stopped.")
+            log.info("===== TC poller stopped.")
         self._poller = None
         return self
 
diff --git a/api/src/opentrons/hardware_control/modules/tempdeck.py b/api/src/opentrons/hardware_control/modules/tempdeck.py
index fb36fd97536..5d9ee818e24 100644
--- a/api/src/opentrons/hardware_control/modules/tempdeck.py
+++ b/api/src/opentrons/hardware_control/modules/tempdeck.py
@@ -243,8 +243,8 @@ async def _connect(self):
         self._poller.start()
 
     def cleanup(self) -> None:
-        if self._poller:
-            log.debug("Stopping tempdeck poller.")
+        if hasattr(self, '_poller') and self._poller:
+            log.info("===== Stopping tempdeck poller.")
             self._poller.stop()
 
     async def prep_for_update(self) -> str:
diff --git a/api/src/opentrons/hardware_control/thread_manager.py b/api/src/opentrons/hardware_control/thread_manager.py
index d85129fa011..a337cf7c0df 100644
--- a/api/src/opentrons/hardware_control/thread_manager.py
+++ b/api/src/opentrons/hardware_control/thread_manager.py
@@ -175,7 +175,7 @@ def clean_up(self):
             loop.call_soon_threadsafe(loop.stop)
         except Exception:
             pass
-        self._cached_modules = {}
+        object.__setattr__(self, '_cached_modules', {})
         object.__getattribute__(self, '_thread').join()
 
     def wrap_module(
@@ -199,7 +199,7 @@ def __getattribute__(self, attr_name):
             # is necessary in order to allow the garbage collector to delete
             # those objects from memory and cleanly stop all associated threads.
             cached_mods = {
-                    module: cached_mods.get(module, wrap(module)) for module in attr}
+                    module: cached_mods.get(module) or wrap(module) for module in attr}
             object.__setattr__(self, '_cached_modules', cached_mods)
             return cached_mods.values()
         elif attr_name == 'clean_up':
diff --git a/api/tests/opentrons/hardware_control/test_modules.py b/api/tests/opentrons/hardware_control/test_modules.py
index 831c122f162..68c2addbc2a 100644
--- a/api/tests/opentrons/hardware_control/test_modules.py
+++ b/api/tests/opentrons/hardware_control/test_modules.py
@@ -14,6 +14,7 @@
 from opentrons.drivers.rpi_drivers.types import USBPort
 
 
+# Freezes
 async def test_get_modules_simulating():
     import opentrons.hardware_control as hardware_control
     mods = ['tempdeck', 'magdeck', 'thermocycler']
@@ -61,6 +62,7 @@ async def test_module_caching():
     assert two_magdecks[1] is not two_magdecks[0]
 
 
+# freezes
 async def test_filtering_modules():
     import opentrons.hardware_control as hardware_control
     mods = [
@@ -178,6 +180,7 @@ async def mock_find_bossa_bootloader_port():
     )
 
 
+# freezes
 async def test_get_bundled_fw(monkeypatch, tmpdir):
     from opentrons.hardware_control import modules
 

From 35b89b7dc0d9bad0303880e43781207193923fa1 Mon Sep 17 00:00:00 2001
From: Sanniti <sanniti@opentrons.com>
Date: Thu, 22 Apr 2021 17:53:02 -0400
Subject: [PATCH 6/6] use WeakKeyDictionary, add tests, cleanup

---
 .../opentrons/drivers/thermocycler/driver.py  |  6 +--
 .../hardware_control/modules/magdeck.py       |  3 ++
 .../hardware_control/modules/tempdeck.py      |  5 ++-
 .../hardware_control/modules/thermocycler.py  |  3 ++
 .../hardware_control/thread_manager.py        | 40 +++++++++++-------
 .../hardware_control/test_modules.py          |  3 --
 .../hardware_control/test_thread_manager.py   | 42 +++++++++++++++++++
 7 files changed, 79 insertions(+), 23 deletions(-)

diff --git a/api/src/opentrons/drivers/thermocycler/driver.py b/api/src/opentrons/drivers/thermocycler/driver.py
index 0803a6df04f..889662878fc 100644
--- a/api/src/opentrons/drivers/thermocycler/driver.py
+++ b/api/src/opentrons/drivers/thermocycler/driver.py
@@ -331,12 +331,12 @@ def send(self, command, callback):
             self._send_write_fd.write(b'c')
 
     def close(self):
-        log.info("===== Halting TCPoller")
+        log.debug("Halting TCPoller")
         self._halt_write_fd.write(b'q')
 
     def __del__(self):
         """ Clean up thread fifos"""
-        log.info("===== Cleaning up thread fifos in TCPoller.")
+        log.debug("Cleaning up thread fifos in TCPoller.")
         try:
             os.unlink(self._send_path)
         except NameError:
@@ -381,7 +381,7 @@ def disconnect(self) -> 'Thermocycler':
         if self.is_connected() or self._poller:
             self._poller.close()  # type: ignore
             self._poller.join()  # type: ignore
-            log.info("===== TC poller stopped.")
+            log.debug("TC poller stopped.")
         self._poller = None
         return self
 
diff --git a/api/src/opentrons/hardware_control/modules/magdeck.py b/api/src/opentrons/hardware_control/modules/magdeck.py
index f6d60ad37d5..3318b82338d 100644
--- a/api/src/opentrons/hardware_control/modules/magdeck.py
+++ b/api/src/opentrons/hardware_control/modules/magdeck.py
@@ -212,6 +212,9 @@ def _disconnect(self):
     def cleanup(self) -> None:
         self._disconnect()
 
+    def __del__(self):
+        self.cleanup()
+
     async def prep_for_update(self) -> str:
         self._driver.enter_programming_mode()
         new_port = await update.find_bootloader_port()
diff --git a/api/src/opentrons/hardware_control/modules/tempdeck.py b/api/src/opentrons/hardware_control/modules/tempdeck.py
index 5d9ee818e24..7df5fc7d5c6 100644
--- a/api/src/opentrons/hardware_control/modules/tempdeck.py
+++ b/api/src/opentrons/hardware_control/modules/tempdeck.py
@@ -244,9 +244,12 @@ async def _connect(self):
 
     def cleanup(self) -> None:
         if hasattr(self, '_poller') and self._poller:
-            log.info("===== Stopping tempdeck poller.")
+            log.debug("Stopping tempdeck poller.")
             self._poller.stop()
 
+    def __del__(self):
+        self.cleanup()
+
     async def prep_for_update(self) -> str:
         model = self._device_info and self._device_info.get('model')
         if model in ('temp_deck_v1', 'temp_deck_v1.1', 'temp_deck_v2'):
diff --git a/api/src/opentrons/hardware_control/modules/thermocycler.py b/api/src/opentrons/hardware_control/modules/thermocycler.py
index aa622a0da03..e890982ea65 100644
--- a/api/src/opentrons/hardware_control/modules/thermocycler.py
+++ b/api/src/opentrons/hardware_control/modules/thermocycler.py
@@ -353,3 +353,6 @@ def cleanup(self) -> None:
             self._driver.disconnect()
         except Exception:
             MODULE_LOG.exception('Exception while cleaning up Thermocycler')
+
+    def __del__(self):
+        self.cleanup()
diff --git a/api/src/opentrons/hardware_control/thread_manager.py b/api/src/opentrons/hardware_control/thread_manager.py
index a337cf7c0df..380d8eb1cc9 100644
--- a/api/src/opentrons/hardware_control/thread_manager.py
+++ b/api/src/opentrons/hardware_control/thread_manager.py
@@ -4,7 +4,8 @@
 import logging
 import asyncio
 import functools
-from typing import Generic, TypeVar, Any, Optional, Dict
+import weakref
+from typing import Generic, TypeVar, Any, Optional
 from .adapters import SynchronousAdapter
 from .modules.mod_abc import AbstractModule
 
@@ -108,7 +109,8 @@ def __init__(self, builder, *args, **kwargs):
         self._sync_managed_obj: Optional[SynchronousAdapter] = None
         is_running = threading.Event()
         self._is_running = is_running
-        self._cached_modules: Dict[AbstractModule, CallBridger[AbstractModule]] = {}
+        self._cached_modules: weakref.WeakKeyDictionary[
+            AbstractModule, CallBridger[AbstractModule]] = weakref.WeakKeyDictionary()
         # TODO: remove this if we switch to python 3.8
         # https://docs.python.org/3/library/asyncio-subprocess.html#subprocess-and-threads  # noqa
         # On windows, the event loop and system interface is different and
@@ -175,12 +177,27 @@ def clean_up(self):
             loop.call_soon_threadsafe(loop.stop)
         except Exception:
             pass
-        object.__setattr__(self, '_cached_modules', {})
+        object.__setattr__(self, '_cached_modules', weakref.WeakKeyDictionary({}))
         object.__getattribute__(self, '_thread').join()
 
-    def wrap_module(
-            self, module: AbstractModule) -> CallBridger[AbstractModule]:
-        return CallBridger(module, object.__getattribute__(self, '_loop'))
+    def wrap_module(self, module: AbstractModule) -> CallBridger[AbstractModule]:
+        """ Return the module object wrapped in a CallBridger and cache it.
+
+        The wrapped module objects are cached in `self._cached_modules` so they can be
+        re-used throughout the module object's life, as creating a wrapper is expensive.
+        We use a WeakKeyDictionary for caching so that module objects can be
+        garbage collected when modules are detached (since entries in WeakKeyDictionary
+        get discarded when there is no longer a strong reference to the key).
+        """
+        wrapper_cache = object.__getattribute__(self, '_cached_modules')
+        this_module_wrapper = wrapper_cache.get(module)
+
+        if this_module_wrapper is None:
+            this_module_wrapper = CallBridger(module,
+                                              object.__getattribute__(self, '_loop'))
+            wrapper_cache.update({module: this_module_wrapper})
+
+        return this_module_wrapper
 
     def __getattribute__(self, attr_name):
         # hardware_control.api.API.attached_modules is the only hardware
@@ -192,16 +209,7 @@ def __getattribute__(self, attr_name):
             wrap = object.__getattribute__(self, 'wrap_module')
             managed = object.__getattribute__(self, 'managed_obj')
             attr = getattr(managed, attr_name)
-            cached_mods = object.__getattribute__(self, '_cached_modules')
-
-            # Update self._cached_modules to delete all removed modules' entries and add
-            # newly attached modules. Removing references to stale instances
-            # is necessary in order to allow the garbage collector to delete
-            # those objects from memory and cleanly stop all associated threads.
-            cached_mods = {
-                    module: cached_mods.get(module) or wrap(module) for module in attr}
-            object.__setattr__(self, '_cached_modules', cached_mods)
-            return cached_mods.values()
+            return [wrap(mod) for mod in attr]
         elif attr_name == 'clean_up':
             # the wrapped object probably has this attr as well as us, and we
             # want to call both, with the wrapped one first
diff --git a/api/tests/opentrons/hardware_control/test_modules.py b/api/tests/opentrons/hardware_control/test_modules.py
index 68c2addbc2a..831c122f162 100644
--- a/api/tests/opentrons/hardware_control/test_modules.py
+++ b/api/tests/opentrons/hardware_control/test_modules.py
@@ -14,7 +14,6 @@
 from opentrons.drivers.rpi_drivers.types import USBPort
 
 
-# Freezes
 async def test_get_modules_simulating():
     import opentrons.hardware_control as hardware_control
     mods = ['tempdeck', 'magdeck', 'thermocycler']
@@ -62,7 +61,6 @@ async def test_module_caching():
     assert two_magdecks[1] is not two_magdecks[0]
 
 
-# freezes
 async def test_filtering_modules():
     import opentrons.hardware_control as hardware_control
     mods = [
@@ -180,7 +178,6 @@ async def mock_find_bossa_bootloader_port():
     )
 
 
-# freezes
 async def test_get_bundled_fw(monkeypatch, tmpdir):
     from opentrons.hardware_control import modules
 
diff --git a/api/tests/opentrons/hardware_control/test_thread_manager.py b/api/tests/opentrons/hardware_control/test_thread_manager.py
index 2e55651c055..91530365826 100644
--- a/api/tests/opentrons/hardware_control/test_thread_manager.py
+++ b/api/tests/opentrons/hardware_control/test_thread_manager.py
@@ -1,6 +1,9 @@
 import pytest
+import weakref
+from opentrons.hardware_control.modules import ModuleAtPort
 from opentrons.hardware_control.thread_manager import ThreadManagerException,\
     ThreadManager
+from opentrons.hardware_control.api import API
 
 
 def test_build_fail_raises_exception():
@@ -12,3 +15,42 @@ def f():
         raise Exception()
     with pytest.raises(ThreadManagerException):
         ThreadManager(f)
+
+
+def test_module_cache_add_entry():
+    """ Test that _cached_modules updates correctly."""
+
+    mod_names = ['tempdeck']
+    thread_manager = ThreadManager(API.build_hardware_simulator,
+                                   attached_modules=mod_names)
+
+    # Test that module gets added to the cache
+    mods = thread_manager.attached_modules
+    wrapper_cache = thread_manager._cached_modules.copy()
+    assert isinstance(wrapper_cache, weakref.WeakKeyDictionary)
+    assert len(wrapper_cache) == 1
+    assert mods[0] in wrapper_cache.values()
+
+    # Test that calling attached modules doesn't add duplicate entries to cache
+    mods2 = thread_manager.attached_modules
+    wrapper_cache2 = thread_manager._cached_modules.copy()
+    assert len(wrapper_cache2) == 1
+    assert mods == mods2
+
+
+async def test_module_cache_remove_entry():
+    """Test that module entry gets removed from cache when module detaches. """
+    mod_names = ['tempdeck', 'magdeck']
+    thread_manager = ThreadManager(API.build_hardware_simulator,
+                                   attached_modules=mod_names)
+
+    mods_before = thread_manager.attached_modules
+    assert len(mods_before) == 2
+
+    await thread_manager.register_modules(
+        removed_mods_at_ports=[
+            ModuleAtPort(port='/dev/ot_module_sim_tempdeck0',
+                         name='tempdeck')
+        ])
+    mods_after = thread_manager.attached_modules
+    assert len(mods_after) == 1