Handle SystemExit gracefully in two phase commit mode

CHANGES.rst

@@ -6,7 +6,8 @@ CHANGES
 2.0.2 (unreleased)
 ------------------
 
-- Nothing changed yet.
+- Handle SystemExit gracefully if raised while PJDataManager is
+  joining the transaction in two-phase-commit mode.
 
 
 2.0.1 (2020-10-13)

src/pjpersist/datamanager.py

@@ -853,6 +853,20 @@ def _begin(self, transaction):
         except psycopg2.Error as e:
             check_for_disconnect(e, 'PJDataManager._begin')
             raise
+        except BaseException as exc:
+            # There's a corner case when SystemExit is raised while
+            # tpc_begin is still being processed.  Because we don't
+            # know if that operation completed, we need to check
+            # transaction status so we can clean up correctly.
+            if self._conn.status == psycopg2.extensions.STATUS_BEGIN:
+                self._tpc_activated = True
+            else:
+                LOG.warning(
+                    "PJDataManager._begin: tried to call connection.tpc_begin,"
+                    " but it did not change connection status to STATUS_BEGIN,"
+                    " failing with a generic exception %s."
+                    "  Not setting TPC mode on PJDataManager.", exc)
+            raise
         self._tpc_activated = True
 
     def _log_rw_stats(self):

src/pjpersist/tests/test_datamanager.py

@@ -1609,6 +1609,74 @@ def cursor(self, abort=False):
             transaction.commit()
 
 
+class ConnectionProxy:
+    __conn__ = None
+    def __init__(self, conn, **kw):
+        self.__dict__['__conn__'] = conn
+        self.__dict__.update(kw)
+
+    def __getattr__(self, name):
+        return getattr(self.__conn__, name)
+
+    def __setattr__(self, name, value):
+        return setattr(self.__conn__, name, value)
+
+class TwoPhaseCommitTestCase(testing.PJTestCase):
+
+    def setUp(self):
+        super(TwoPhaseCommitTestCase, self).setUp()
+
+        # get rid of the previous transaction
+        transaction.abort()
+
+        tpc_patch = mock.patch(
+            "pjpersist.datamanager.PJ_TWO_PHASE_COMMIT_ENABLED", True)
+        no_prep_patch = mock.patch(
+            "pjpersist.datamanager."
+            "CALL_TPC_PREPARE_ON_NO_WRITE_TRANSACTION", False)
+        log_patch = mock.patch(
+            "pjpersist.datamanager.LOG_READ_WRITE_TRANSACTION", True)
+        self.patches = [tpc_patch, no_prep_patch, log_patch]
+        for p in self.patches:
+            p.start()
+
+        self.conn = testing.getConnection(testing.DBNAME)
+
+    def tearDown(self):
+        for p in self.patches:
+            p.stop()
+
+        super(TwoPhaseCommitTestCase, self).tearDown()
+
+    def test_system_exit_while_joining_transaction_on_tpc_begin(self):
+        def tpc_begin(tx):
+            self.conn.tpc_begin(tx)
+            raise SystemExit(-241)
+        proxy = ConnectionProxy(self.conn, tpc_begin=tpc_begin)
+
+        with self.assertRaises(SystemExit):
+            datamanager.PJDataManager(proxy)
+
+    def test_system_exit_while_joining_transaction_before_tpc_begin(self):
+        def tpc_begin(tx):
+            raise SystemExit(-241)
+        proxy = ConnectionProxy(self.conn, tpc_begin=tpc_begin)
+
+        with self.assertRaises(SystemExit):
+            datamanager.PJDataManager(proxy)
+
+    def test_system_exit_while_joining_transaction_fail_rollback(self):
+        def tpc_begin(tx):
+            self.conn.tpc_begin(tx)
+            raise SystemExit(-241)
+        def tpc_rollback(tx):
+            raise psycopg2.OperationalError('boom')
+        proxy = ConnectionProxy(self.conn, tpc_begin=tpc_begin)
+
+        with self.assertRaises(SystemExit):
+            datamanager.PJDataManager(proxy)
+
+
 def test_suite():
     dtsuite = doctest.DocTestSuite(
         setUp=testing.setUp, tearDown=testing.tearDown,
@@ -1622,4 +1690,5 @@ def test_suite():
         unittest.makeSuite(QueryLoggingTestCase),
         unittest.makeSuite(TransactionOptionsTestCase),
         unittest.makeSuite(DirtyTestCase),
+        unittest.makeSuite(TwoPhaseCommitTestCase),
         ))