Fix deadlock in Measurements.close() and add test coverage

concordia/utils/async_measurements_test.py

@@ -0,0 +1,90 @@
+# Copyright 2025 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for ReactiveMeasurements."""
+
+from absl.testing import absltest
+from concordia.utils import async_measurements
+
+
+class ReactiveMeasurementsTest(absltest.TestCase):
+
+  def test_publish_datum_still_stores_in_channel(self):
+    m = async_measurements.ReactiveMeasurements()
+    m.publish_datum('a', 1)
+    self.assertEqual(m.get_channel('a'), [1])
+
+  def test_subscribe_receives_published_data(self):
+    m = async_measurements.ReactiveMeasurements()
+    received = []
+    subscription = m.subscribe(lambda pair: received.append(pair))
+    m.publish_datum('a', 1)
+    m.publish_datum('b', 2)
+    subscription.dispose()
+    self.assertEqual(received, [('a', 1), ('b', 2)])
+
+  def test_capture_collects_only_data_with_matching_key(self):
+    m = async_measurements.ReactiveMeasurements()
+    with m.capture('entity_1') as captured:
+      m.publish_datum('a', 'mine', capture_key='entity_1')
+      m.publish_datum('b', 'not_mine', capture_key='entity_2')
+      m.publish_datum('c', 'uncaptured')
+    self.assertEqual(captured, {'a': 'mine'})
+
+  def test_capture_only_keeps_last_datum_per_channel(self):
+    m = async_measurements.ReactiveMeasurements()
+    with m.capture('entity_1') as captured:
+      m.publish_datum('a', 'first', capture_key='entity_1')
+      m.publish_datum('a', 'second', capture_key='entity_1')
+    self.assertEqual(captured, {'a': 'second'})
+
+  def test_capture_stops_collecting_after_context_exits(self):
+    m = async_measurements.ReactiveMeasurements()
+    with m.capture('entity_1') as captured:
+      m.publish_datum('a', 'inside', capture_key='entity_1')
+    m.publish_datum('a', 'outside', capture_key='entity_1')
+    self.assertEqual(captured, {'a': 'inside'})
+
+  def test_nested_captures_with_different_keys_do_not_cross_contaminate(self):
+    m = async_measurements.ReactiveMeasurements()
+    with m.capture('entity_1') as captured_1:
+      with m.capture('entity_2') as captured_2:
+        m.publish_datum('a', 'for_1', capture_key='entity_1')
+        m.publish_datum('a', 'for_2', capture_key='entity_2')
+    self.assertEqual(captured_1, {'a': 'for_1'})
+    self.assertEqual(captured_2, {'a': 'for_2'})
+
+  def test_close_still_works_on_reactive_subclass(self):
+    # Regression coverage: close()/close_channel() must not deadlock on the
+    # ReactiveMeasurements subclass either.
+    m = async_measurements.ReactiveMeasurements()
+    m.publish_datum('a', 1)
+    m.publish_datum('b', 2)
+    m.close()
+    self.assertEqual(m.available_channels(), set())
+
+  def test_dispose_completes_subject(self):
+    m = async_measurements.ReactiveMeasurements()
+    received = []
+    completed = []
+    m.subscribe(lambda pair: received.append(pair))
+    m._subject.subscribe(on_completed=lambda: completed.append(True))  # pylint: disable=protected-access
+    m.publish_datum('a', 1)
+    m.dispose()
+    self.assertEqual(received, [('a', 1)])
+    self.assertTrue(completed)
+
+
+if __name__ == '__main__':
+  absltest.main()

concordia/utils/measurements.py

@@ -97,15 +97,13 @@ def close_channel(self, channel: str) -> None:
     """Closes the channel for the given name.
 
     Args:
-      channel: The channel to close. If the channel doesn't exist yet, it will
-        be created.
+      channel: The channel to close. If the channel doesn't exist, this is a
+        no-op.
     """
     with self._channels_lock:
-      del self._channels[channel]
+      self._channels.pop(channel, None)
 
   def close(self) -> None:
     """Closes all channels."""
     with self._channels_lock:
-      for channel in self._channels:
-        self.close_channel(channel)
       self._channels.clear()

concordia/utils/measurements_test.py

@@ -0,0 +1,117 @@
+# Copyright 2024 DeepMind Technologies Limited.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Tests for the Measurements registry."""
+
+import threading
+
+from absl.testing import absltest
+from concordia.utils import measurements as measurements_lib
+
+
+class MeasurementsTest(absltest.TestCase):
+
+  def test_publish_datum_creates_channel(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('channel_a', 1)
+    self.assertEqual(m.available_channels(), {'channel_a'})
+
+  def test_publish_datum_appends_in_order(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('channel_a', 1)
+    m.publish_datum('channel_a', 2)
+    m.publish_datum('channel_a', 3)
+    self.assertEqual(m.get_channel('channel_a'), [1, 2, 3])
+
+  def test_get_channel_creates_empty_channel_if_missing(self):
+    m = measurements_lib.Measurements()
+    self.assertEqual(m.get_channel('missing'), [])
+    self.assertEqual(m.available_channels(), {'missing'})
+
+  def test_get_last_datum_returns_most_recent(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('channel_a', 'first')
+    m.publish_datum('channel_a', 'second')
+    self.assertEqual(m.get_last_datum('channel_a'), 'second')
+
+  def test_get_last_datum_returns_none_for_empty_channel(self):
+    m = measurements_lib.Measurements()
+    self.assertIsNone(m.get_last_datum('missing'))
+
+  def test_get_all_channels_returns_all_data(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('a', 1)
+    m.publish_datum('b', 2)
+    self.assertEqual(m.get_all_channels(), {'a': [1], 'b': [2]})
+
+  def test_get_all_channels_returns_a_copy(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('a', 1)
+    snapshot = m.get_all_channels()
+    m.publish_datum('a', 2)
+    self.assertEqual(snapshot, {'a': [1]})
+    self.assertEqual(m.get_channel('a'), [1, 2])
+
+  def test_close_channel_removes_channel(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('a', 1)
+    m.close_channel('a')
+    self.assertEqual(m.available_channels(), set())
+
+  def test_close_channel_on_missing_channel_is_a_no_op(self):
+    m = measurements_lib.Measurements()
+    # Should not raise, per the method's documented behavior.
+    m.close_channel('never_published')
+    self.assertEqual(m.available_channels(), set())
+
+  def test_close_removes_all_channels(self):
+    m = measurements_lib.Measurements()
+    m.publish_datum('a', 1)
+    m.publish_datum('b', 2)
+    m.publish_datum('c', 3)
+    m.close()
+    self.assertEqual(m.available_channels(), set())
+
+  def test_close_does_not_deadlock_with_multiple_channels(self):
+    # Regression test: close() used to hold _channels_lock while calling
+    # close_channel(), which itself acquired the same non-reentrant lock,
+    # deadlocking whenever there was at least one channel.
+    m = measurements_lib.Measurements()
+    m.publish_datum('a', 1)
+    m.publish_datum('b', 2)
+
+    completed = threading.Event()
+
+    def run_close():
+      m.close()
+      completed.set()
+
+    thread = threading.Thread(target=run_close)
+    thread.start()
+    thread.join(timeout=5)
+    self.assertTrue(completed.is_set(), 'Measurements.close() deadlocked.')
+
+  def test_close_on_empty_measurements_does_not_raise(self):
+    m = measurements_lib.Measurements()
+    m.close()
+    self.assertEqual(m.available_channels(), set())
+
+  def test_get_channel_or_create_requires_lock(self):
+    m = measurements_lib.Measurements()
+    with self.assertRaises(RuntimeError):
+      m._get_channel_or_create('a')  # pylint: disable=protected-access
+
+
+if __name__ == '__main__':
+  absltest.main()