Fix rendering context cleanup logic.

saran-t · copybara-github · commit 4e6426a5a739 · 2022-09-05T10:04:35.000-07:00
Fixes #310. PiperOrigin-RevId: 472296006 Change-Id: I558c1fdf3753ed660f4463ae2dc054f496164c8a
diff --git a/dm_control/_render/base.py b/dm_control/_render/base.py
@@ -88,10 +88,13 @@ def thread(self):
     return self._render_executor.thread
 
   def _free_on_executor_thread(self):  # pylint: disable=missing-function-docstring
-    if _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread] == id(self):
-      del _CURRENT_THREAD_FOR_CONTEXT[id(self)]
-      del _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]
+    current_ctx = _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]
+    if current_ctx is not None:
+      del _CURRENT_THREAD_FOR_CONTEXT[current_ctx]
+    del _CURRENT_CONTEXT_FOR_THREAD[self._render_executor.thread]
+
     self._platform_make_current()
+
     try:
       dummy = []
       while self._patients:
diff --git a/dm_control/_render/base_test.py b/dm_control/_render/base_test.py
@@ -110,6 +110,27 @@ def test_free(self):
     self.assertNotIn(id(self.context), base._CURRENT_THREAD_FOR_CONTEXT)
     self.assertNotIn(thread, base._CURRENT_CONTEXT_FOR_THREAD)
 
+  def test_free_with_multiple_contexts(self):
+    context1 = ContextBaseTests.ContextMock(WIDTH, HEIGHT,
+                                            executor.PassthroughRenderExecutor)
+    with context1.make_current():
+      pass
+
+    context2 = ContextBaseTests.ContextMock(WIDTH, HEIGHT,
+                                            executor.PassthroughRenderExecutor)
+    with context2.make_current():
+      pass
+
+    self.assertEqual(base._CURRENT_CONTEXT_FOR_THREAD[threading.main_thread()],
+                     id(context2))
+    self.assertIs(base._CURRENT_THREAD_FOR_CONTEXT[id(context2)],
+                  threading.main_thread())
+
+    context1.free()
+    self.assertIsNone(
+        base._CURRENT_CONTEXT_FOR_THREAD[self.context.free_thread])
+    self.assertIsNone(base._CURRENT_THREAD_FOR_CONTEXT[id(context2)])
+
   def test_refcounting(self):
     thread = self.context.thread
 
