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 <simon.marchi@efficios.com>
Reviewed-on: https://review.lttng.org/c/babeltrace/+/1559
Reviewed-by: Philippe Proulx <eeppeliteloop@gmail.com>

diff --git a/src/bindings/python/bt2/bt2/message_iterator.py b/src/bindings/python/bt2/bt2/message_iterator.py
index bde3afa054a11a1ca4265fae050fd7f9df0ba8e6..b07e2450afc97b46f7b5e46c1198d6f07aeb2853 100644 (file)
--- 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 64f34d04c00cb51f7a75803e04b53d215d94690c..98ec64aba7c5796bff0052a793e1bca60c2be013 100644 (file)
--- 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):