PikaTransport: Enable Transactions (#92)

Make every subscription run on a separate channel so that
transactions can be used on channels. So if a message is
to be send as part of a transaction then it can be sent on the
relevant channel, and if a message is sent outside of a transaction
it can be sent via the default channel.

This disables confirm mode on all channels for now.
The confirm mode currently does not have any effect as we are not
observing the callback. However, confirm and transaction modes are
mutually exclusive, so transactions can't be used if we use confirm
mode on subscription channels. For the default channel we may
eventually want to have one in normal mode and one in confirm
mode, but this is another issue for another day.

Add optional subscription ID to transaction calls so that a transaction
can be associated with the correct channel where needed.

Author: Anthchirp

src/workflows/services/sample_transaction.py

@@ -17,7 +17,10 @@ class SampleTxn(CommonService):
     def initializing(self):
         """Subscribe to a channel. Received messages must be acknowledged."""
         self.subid = self._transport.subscribe(
-            "transient.transaction", self.receive_message, acknowledgement=True
+            "transient.transaction",
+            self.receive_message,
+            acknowledgement=True,
+            prefetch_count=1000,
         )
 
     @staticmethod

src/workflows/transport/common_transport.py

@@ -288,15 +288,24 @@ def nack(self, message, subscription_id: Optional[int] = None, **kwargs):
         )
         self._nack(message_id, subscription_id=subscription_id, **kwargs)
 
-    def transaction_begin(self, **kwargs) -> int:
+    def transaction_begin(self, subscription_id: Optional[int] = None, **kwargs) -> int:
         """Start a new transaction.
-        :param **kwargs: Further parameters for the transport layer. For example
+        :param **kwargs: Further parameters for the transport layer.
         :return: A transaction ID that can be passed to other functions.
         """
         self.__transaction_id += 1
         self.__transactions.add(self.__transaction_id)
-        self.log.debug("Starting transaction with ID %d", self.__subscription_id)
-        self._transaction_begin(self.__transaction_id, **kwargs)
+        if subscription_id:
+            self.log.debug(
+                "Starting transaction with ID %d on subscription %d",
+                self.__transaction_id,
+                subscription_id,
+            )
+        else:
+            self.log.debug("Starting transaction with ID %d", self.__transaction_id)
+        self._transaction_begin(
+            self.__transaction_id, subscription_id=subscription_id, **kwargs
+        )
         return self.__transaction_id
 
     def transaction_abort(self, transaction_id: int, **kwargs):
@@ -405,21 +414,23 @@ def _nack(self, message_id, subscription_id, **kwargs):
         """
         raise NotImplementedError("Transport interface not implemented")
 
-    def _transaction_begin(self, transaction_id, **kwargs):
+    def _transaction_begin(
+        self, transaction_id: int, *, subscription_id: Optional[int] = None, **kwargs
+    ) -> None:
         """Start a new transaction.
         :param transaction_id: ID for this transaction in the transport layer.
         :param **kwargs: Further parameters for the transport layer.
         """
         raise NotImplementedError("Transport interface not implemented")
 
-    def _transaction_abort(self, transaction_id, **kwargs):
+    def _transaction_abort(self, transaction_id: int, **kwargs) -> None:
         """Abort a transaction and roll back all operations.
         :param transaction_id: ID of transaction to be aborted.
         :param **kwargs: Further parameters for the transport layer.
         """
         raise NotImplementedError("Transport interface not implemented")
 
-    def _transaction_commit(self, transaction_id, **kwargs):
+    def _transaction_commit(self, transaction_id: int, **kwargs) -> None:
         """Commit a transaction.
         :param transaction_id: ID of transaction to be committed.
         :param **kwargs: Further parameters for the transport layer.

src/workflows/transport/pika_transport.py

@@ -513,26 +513,26 @@ def _broadcast(
             mandatory=False,
         ).result()
 
-    def _transaction_begin(self, **kwargs):
-        """Enter transaction mode.
-        :param **kwargs: Further parameters for the transport layer.
+    def _transaction_begin(
+        self, transaction_id: int, *, subscription_id: Optional[int] = None, **kwargs
+    ) -> None:
+        """Start a new transaction.
+        :param transaction_id: ID for this transaction in the transport layer.
+        :param subscription_id: Tie the transaction to a specific channel containing this subscription.
         """
-        raise NotImplementedError()
-        # self._channel.tx_select()
+        self._pika_thread.tx_select(transaction_id, subscription_id)
 
-    def _transaction_abort(self, **kwargs):
+    def _transaction_abort(self, transaction_id: int, **kwargs) -> None:
         """Abort a transaction and roll back all operations.
-        :param **kwargs: Further parameters for the transport layer.
+        :param transaction_id: ID of transaction to be aborted.
         """
-        raise NotImplementedError()
-        # self._channel.tx_rollback()
+        self._pika_thread.tx_rollback(transaction_id)
 
-    def _transaction_commit(self, **kwargs):
+    def _transaction_commit(self, transaction_id: int, **kwargs) -> None:
         """Commit a transaction.
-        :param **kwargs: Further parameters for the transport layer.
+        :param transaction_id: ID of transaction to be committed.
         """
-        raise NotImplementedError()
-        # self._channel.tx_commit()
+        self._pika_thread.tx_commit(transaction_id)
 
     def _ack(
         self, message_id, subscription_id: int, *, multiple: bool = False, **_kwargs
@@ -694,9 +694,12 @@ def __init__(
         self._subscriptions: Dict[int, _PikaSubscription] = {}
         # The pika connection object
         self._connection: Optional[pika.BlockingConnection] = None
-        # Per-subscription channels. May be pointing to the shared channel
+        # Index of per-subscription channels.
         self._pika_channels: Dict[int, BlockingChannel] = {}
-        # A common, shared channel, used for non-QoS subscriptions
+        # Bidirectional index of all ongoing transactions. May include the shared channel
+        self._transactions_by_id: Dict[int, BlockingChannel] = {}
+        self._transactions_by_channel: Dict[BlockingChannel, int] = {}
+        # A common, shared channel, used for sending messages outside of transactions.
         self._pika_shared_channel: Optional[BlockingChannel]
         # Are we allowed to reconnect. Can only be turned off, never on
         self._reconnection_allowed: bool = True
@@ -907,6 +910,11 @@ def _unsubscribe():
                         logger.debug("Closing channel that is now unused")
                         channel.close()
 
+                        # Forget about any ongoing transactions on the channel
+                        if channel in self._transactions_by_channel:
+                            transaction_id = self._transactions_by_channel.pop(channel)
+                            self._transactions_by_id.pop(transaction_id)
+
                     result.set_result(None)
             except BaseException as e:
                 result.set_exception(e)
@@ -974,6 +982,100 @@ def nack(
             lambda: channel.basic_nack(delivery_tag, multiple=multiple, requeue=requeue)
         )
 
+    def tx_select(
+        self, transaction_id: int, subscription_id: Optional[int]
+    ) -> Future[None]:
+        """Set a channel to transaction mode. Thread-safe.
+        :param transaction_id: ID for this transaction in the transport layer.
+        :param subscription_id: Tie the transaction to a specific channel containing this subscription.
+        """
+
+        if not self._connection:
+            raise RuntimeError("Cannot transact on unstarted connection")
+
+        future: Future[None] = Future()
+
+        def _tx_select():
+            if future.set_running_or_notify_cancel():
+                try:
+                    if subscription_id:
+                        if subscription_id not in self._pika_channels:
+                            raise KeyError(
+                                f"Could not find subscription {subscription_id} to begin transaction"
+                            )
+                        channel = self._pika_channels[subscription_id]
+                    else:
+                        channel = self._get_shared_channel()
+                    if channel in self._transactions_by_channel:
+                        raise KeyError(
+                            f"Channel {channel} is already running transaction {self._transactions_by_channel[channel]}, so can't start transaction {transaction_id}"
+                        )
+                    channel.tx_select()
+                    self._transactions_by_channel[channel] = transaction_id
+                    self._transactions_by_id[transaction_id] = channel
+
+                    future.set_result(None)
+                except BaseException as e:
+                    future.set_exception(e)
+                    raise
+
+        self._connection.add_callback_threadsafe(_tx_select)
+        return future
+
+    def tx_rollback(self, transaction_id: int) -> Future[None]:
+        """Abort a transaction and roll back all operations. Thread-safe.
+        :param transaction_id: ID of transaction to be aborted.
+        """
+        if not self._connection:
+            raise RuntimeError("Cannot transact on unstarted connection")
+
+        future: Future[None] = Future()
+
+        def _tx_rollback():
+            if future.set_running_or_notify_cancel():
+                try:
+                    channel = self._transactions_by_id.pop(transaction_id, None)
+                    if not channel:
+                        raise KeyError(
+                            f"Could not find transaction {transaction_id} to roll back"
+                        )
+                    self._transactions_by_channel.pop(channel)
+                    channel.tx_rollback()
+                    future.set_result(None)
+                except BaseException as e:
+                    future.set_exception(e)
+                    raise
+
+        self._connection.add_callback_threadsafe(_tx_rollback)
+        return future
+
+    def tx_commit(self, transaction_id: int) -> Future[None]:
+        """Commit a transaction.
+        :param transaction_id: ID of transaction to be committed. Thread-safe..
+        """
+        if not self._connection:
+            raise RuntimeError("Cannot transact on unstarted connection")
+
+        future: Future[None] = Future()
+
+        def _tx_commit():
+            if future.set_running_or_notify_cancel():
+                try:
+                    channel = self._transactions_by_id.pop(transaction_id, None)
+                    if not channel:
+                        raise KeyError(
+                            f"Could not find transaction {transaction_id} to commit"
+                        )
+                    self._transactions_by_channel.pop(channel)
+                    channel.tx_commit()
+                    future.set_result(None)
+                except BaseException as e:
+                    future.set_exception(e)
+                    raise
+
+        self._connection.add_callback_threadsafe(_tx_commit)
+        return future
+
     @property
     def connection_alive(self) -> bool:
         """
@@ -989,7 +1091,7 @@ def connection_alive(self) -> bool:
         )
 
     # NOTE: With reconnection lifecycle this probably doesn't make sense
-    #       on it's own. It might make sense to add this returning a
+    #       on its own. It might make sense to add this returning a
     #       connection-specific 'token' - presumably the user might want
     #       to ensure that a connection is still the same connection
     #       and thus adhering to various within-connection guarantees.
@@ -1017,7 +1119,7 @@ def _get_shared_channel(self) -> BlockingChannel:
 
         if not self._pika_shared_channel:
             self._pika_shared_channel = self._connection.channel()
-            self._pika_shared_channel.confirm_delivery()
+            ##### self._pika_shared_channel.confirm_delivery()
         return self._pika_shared_channel
 
     def _recreate_subscriptions(self):
@@ -1050,22 +1152,18 @@ def _add_subscription(self, subscription_id: int, subscription: _PikaSubscriptio
                 f"Subscription {subscription_id} to '{subscription.destination}' is not reconnectable. Turning reconnection off."
             )
 
-        # Either open a channel (if prefetch) or use the shared one
-        if subscription.prefetch_count == 0:
-            channel = self._get_shared_channel()
-        else:
-            channel = self._connection.channel()
-            channel.confirm_delivery()
-            channel.basic_qos(prefetch_count=subscription.prefetch_count)
+        # Open a dedicated channel for this subscription
+        channel = self._connection.channel()
+        channel.basic_qos(prefetch_count=subscription.prefetch_count)
 
-        if subscription.kind == _PikaSubscriptionKind.FANOUT:
+        if subscription.kind is _PikaSubscriptionKind.FANOUT:
             # If a FANOUT subscription, then we need to create and bind
             # a temporary queue to receive messages from the exchange
             queue = channel.queue_declare("", exclusive=True).method.queue
             assert queue is not None
             channel.queue_bind(queue, subscription.destination)
             subscription.queue = queue
-        elif subscription.kind == _PikaSubscriptionKind.DIRECT:
+        elif subscription.kind is _PikaSubscriptionKind.DIRECT:
             subscription.queue = subscription.destination
         else:
             raise NotImplementedError(f"Unknown subscription kind: {subscription.kind}")

src/workflows/transport/stomp_transport.py

@@ -410,21 +410,18 @@ def _broadcast(
     def _transaction_begin(self, transaction_id, **kwargs):
         """Start a new transaction.
         :param transaction_id: ID for this transaction in the transport layer.
-        :param **kwargs: Further parameters for the transport layer.
         """
         self._conn.begin(transaction=transaction_id)
 
     def _transaction_abort(self, transaction_id, **kwargs):
         """Abort a transaction and roll back all operations.
         :param transaction_id: ID of transaction to be aborted.
-        :param **kwargs: Further parameters for the transport layer.
         """
         self._conn.abort(transaction_id)
 
     def _transaction_commit(self, transaction_id, **kwargs):
         """Commit a transaction.
         :param transaction_id: ID of transaction to be committed.
-        :param **kwargs: Further parameters for the transport layer.
         """
         self._conn.commit(transaction_id)
 

tests/services/test_sample_transaction.py

@@ -52,7 +52,7 @@ def test_txnservice_subscribes_to_channel():
     p.initializing()
 
     mock_transport.subscribe.assert_called_once_with(
-        mock.ANY, p.receive_message, acknowledgement=True
+        mock.ANY, p.receive_message, acknowledgement=True, prefetch_count=1000
     )
 
 

tests/transport/test_common.py

@@ -248,7 +248,7 @@ def test_create_and_destroy_transactions():
     t = ct.transaction_begin()
 
     assert t
-    ct._transaction_begin.assert_called_once_with(t)
+    ct._transaction_begin.assert_called_once_with(t, subscription_id=None)
 
     ct.transaction_abort(t)
     with pytest.raises(workflows.Error):