From d79a835343447b4e7f62af721cfb2a05fa9f450a Mon Sep 17 00:00:00 2001 From: Simon Marchi Date: Thu, 27 Jun 2019 15:56:28 -0400 Subject: [PATCH] Fix: bt2: fix reference counting of messages returned by Python components The refcount of messages returned by components implemented in Python (in _UserMessageIterator._next_from_native) is not right. When the C iterator method "next" puts bt_message pointers in the output array, it gives its reference to the bt_messages to the array. If it wishes to keep owning a copy of the message, it needs to acquire a new reference. For this reason, we currently make the Message Python object release its reference to the bt_message in _UserMessageIterator.__next__ (otherwise there would be a double free). Releasing the reference has the effect of setting the _ptr property to None, without changing the refcount. However, it's possible for the user to keep a reference to that Python object and re-use it. This is a problem, since the Message Python object now unexpectedly no longer wraps a valid bt_message. Trying to return that message in the following call to __next__, for example, will fail. To fix this, this patch makes _next_from_native acquire a new reference before returning the bt_message. If the user didn't keep a reference to the Python object, the Python object will die and put the bt_message. The message array will therefore be the sole owner of the message. If the user did keep a reference to the Python object, then both the Python object and the message array will now own a reference to the message. Change-Id: I6f7c472fddeea35509ee669e11d913bec7091f40 Signed-off-by: Simon Marchi Reviewed-on: https://review.lttng.org/c/babeltrace/+/1559 Reviewed-by: Philippe Proulx --- .../python/bt2/bt2/message_iterator.py | 8 ++-- .../python/bt2/test_message_iterator.py | 47 +++++++++++++++++++ 2 files changed, 52 insertions(+), 3 deletions(-) diff --git a/src/bindings/python/bt2/bt2/message_iterator.py b/src/bindings/python/bt2/bt2/message_iterator.py index bde3afa0..b07e2450 100644 --- a/src/bindings/python/bt2/bt2/message_iterator.py +++ b/src/bindings/python/bt2/bt2/message_iterator.py @@ -131,9 +131,11 @@ class _UserMessageIterator(_MessageIterator): utils._check_type(msg, bt2.message._Message) - # Release the reference to the native part. - ptr = msg._release() - return int(ptr) + # The reference we return will be given to the message array. + # However, the `msg` Python object may stay alive, if the user has kept + # a reference to it. Acquire a new reference to account for that. + msg._get_ref(msg._ptr) + return int(msg._ptr) def _create_event_message(self, event_class, packet, default_clock_snapshot=None): utils._check_type(event_class, bt2.event_class._EventClass) diff --git a/tests/bindings/python/bt2/test_message_iterator.py b/tests/bindings/python/bt2/test_message_iterator.py index 64f34d04..98ec64ab 100644 --- a/tests/bindings/python/bt2/test_message_iterator.py +++ b/tests/bindings/python/bt2/test_message_iterator.py @@ -119,6 +119,53 @@ class UserMessageIteratorTestCase(unittest.TestCase): self.assertIsNotNone(addr) self.assertNotEqual(addr, 0) + # Test that messages returned by _UserMessageIterator.__next__ remain valid + # and can be re-used. + def test_reuse_message(self): + class MyIter(bt2._UserMessageIterator): + def __init__(self, port): + tc, sc, ec = port.user_data + trace = tc() + stream = trace.create_stream(sc) + packet = stream.create_packet() + + # This message will be returned twice by __next__. + event_message = self._create_event_message(ec, packet) + + self._msgs = [ + self._create_stream_beginning_message(stream), + self._create_stream_activity_beginning_message(stream), + self._create_packet_beginning_message(packet), + event_message, + event_message, + ] + + def __next__(self): + return self._msgs.pop(0) + + class MySource(bt2._UserSourceComponent, message_iterator_class=MyIter): + def __init__(self, params): + tc = self._create_trace_class() + sc = tc.create_stream_class() + ec = sc.create_event_class() + self._add_output_port('out', (tc, sc, ec)) + + graph = bt2.Graph() + src = graph.add_component(MySource, 'src') + it = graph.create_output_port_message_iterator(src.output_ports['out']) + + # Skip beginning messages. + next(it) + next(it) + next(it) + + msg_ev1 = next(it) + msg_ev2 = next(it) + + self.assertIsInstance(msg_ev1, bt2.message._EventMessage) + self.assertIsInstance(msg_ev2, bt2.message._EventMessage) + self.assertEqual(msg_ev1.addr, msg_ev2.addr) + class OutputPortMessageIteratorTestCase(unittest.TestCase): def test_component(self): -- 2.34.1