Implement caching of whether RepeatableRead works

msullivan · msullivan · commit 95cb8a62244f · 2025-04-24T16:33:11.000-07:00
diff --git a/gel/base_client.py b/gel/base_client.py
@@ -415,6 +415,7 @@ class BasePoolImpl(abc.ABC):
         "_connect_args",
         "_codecs_registry",
         "_query_cache",
+        "_tx_needs_serializable_cache",
         "_connection_factory",
         "_queue",
         "_user_max_concurrency",
@@ -444,6 +445,11 @@ def __init__(
         self._connect_args = connect_args
         self._codecs_registry = protocol.CodecsRegistry()
         self._query_cache = protocol.LRUMapping(maxsize=QUERY_CACHE_SIZE)
+        # Whether a transaction() call from a particular source location
+        # needs to use Serializable. See transaction.BaseRetry for details.
+        self._tx_needs_serializable_cache = (
+            protocol.LRUMapping(maxsize=QUERY_CACHE_SIZE)
+        )
 
         if max_concurrency is not None and max_concurrency <= 0:
             raise ValueError(
@@ -639,6 +645,13 @@ class BaseClient(abstract.BaseReadOnlyExecutor, _options._OptionsMixin):
     __slots__ = ("_impl", "_options")
     _impl_class = NotImplemented
 
+    # Number of stack frames that the Retry objects used by
+    # transaction() need to look up to find their caller from for
+    # caching purposes.  We define this in a variable on BaseClient
+    # basically to make it possible for a user to hackily override it
+    # if they always have some wrapper they need to skip past...
+    _TRANSACTION_FRAME_OFFSET = 2
+
     def __init__(
         self,
         *,
diff --git a/gel/protocol/lru.pxd b/gel/protocol/lru.pxd
@@ -25,4 +25,4 @@ cdef class LRUMapping:
         object _dict_move_to_end
         object _dict_get
 
-    cdef get(self, key, default)
+    cpdef get(self, key, default)
diff --git a/gel/protocol/lru.pyx b/gel/protocol/lru.pyx
@@ -53,7 +53,7 @@ cdef class LRUMapping:
         self._dict_get = self._dict.get
         self._maxsize = maxsize
 
-    cdef get(self, key, default):
+    cpdef get(self, key, default):
         o = self._dict_get(key, _LRU_MARKER)
         if o is _LRU_MARKER:
             return default
diff --git a/gel/transaction.py b/gel/transaction.py
@@ -20,6 +20,7 @@
 import typing
 
 import enum
+import sys
 
 from . import abstract
 from . import errors
@@ -242,13 +243,25 @@ def __init__(self, owner):
         self._done = False
         self._next_backoff = 0
         self._options = owner._options
+        self._key = None
 
         prefer_rr = (
             self._options.transaction_options._isolation
             == options.IsolationLevel.PreferRepeatableRead
         )
+        if prefer_rr:
+            owner = self._owner
+            frame = sys._getframe(owner._TRANSACTION_FRAME_OFFSET)
+            self._key = key = (frame.f_code.co_filename, frame.f_lineno)
+
+            # If we have seen this cache key before and it needed
+            # Serializable, then do serializable.
+            if owner._impl._tx_needs_serializable_cache.get(key, False):
+                prefer_rr = False
+
         self._optimistic_rr = prefer_rr
 
+
     def _retry(self, exc):
         self._last_exception = exc
         rule = self._options.retry_options.get_rule_for_exception(exc)
@@ -262,6 +275,9 @@ def _retry_rr_failure(self, exc):
         # Retry a failure due to REPEATABLE READ not working
         if not self._optimistic_rr:
             return False
+        if self._key:
+            self._owner._impl._tx_needs_serializable_cache[self._key] = True
+
         self._optimistic_rr = False
         # Decrement _iteration count, since this one doesn't really count.
         self._iteration -= 1
diff --git a/tests/test_async_tx.py b/tests/test_async_tx.py
@@ -32,10 +32,18 @@ class TestAsyncTx(tb.AsyncQueryTestCase):
         CREATE TYPE test::TransactionTest EXTENDING std::Object {
             CREATE PROPERTY name -> std::str;
         };
-    '''
 
-    TEARDOWN = '''
-        DROP TYPE test::TransactionTest;
+
+        CREATE TYPE test::Tmp {
+            CREATE REQUIRED PROPERTY tmp -> std::str;
+        };
+        CREATE TYPE test::TmpConflict {
+            CREATE REQUIRED PROPERTY tmp -> std::str {
+                CREATE CONSTRAINT exclusive;
+            }
+        };
+
+        CREATE TYPE test::TmpConflictChild extending test::TmpConflict;
     '''
 
     async def test_async_transaction_regular_01(self):
@@ -104,3 +112,61 @@ async def test_async_transaction_exclusive(self):
                 ):
                     await asyncio.wait_for(f1, timeout=5)
                     await asyncio.wait_for(f2, timeout=5)
+
+
+    async def _try_bogus_rr_tx(self, con, first_try):
+        # A transaction that needs to be serializable
+        async for tx in con.transaction():
+            async with tx:
+                res1 = await tx.query_single('''
+                    select {
+                        ins := (insert test::Tmp { tmp := "test1" }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+                # If this is the second time we've tried to run this
+                # transaction, then the cache should ensure we *only*
+                # try Serializable.
+                if not first_try:
+                    self.assertEqual(res1.level, 'Serializable')
+
+                res2 = await tx.query_single('''
+                    select {
+                        ins := (insert test::TmpConflict {
+                            tmp := <str>random()
+                        }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+
+                # N.B: res1 will be RepeatableRead on the first
+                # iteration, maybe, but contingent on the second query
+                # succeeding it will be Serializable!
+                self.assertEqual(res1.level, 'Serializable')
+                self.assertEqual(res2.level, 'Serializable')
+
+    async def test_async_transaction_prefer_rr(self):
+        if (
+            str(self.server_version.stage) != 'dev'
+            and (self.server_version.major, self.server_version.minor) < (6, 5)
+        ):
+            self.skipTest("DML in RepeatableRead not supported yet")
+        con = self.client.with_transaction_options(
+            edgedb.TransactionOptions(
+                isolation=edgedb.IsolationLevel.PreferRepeatableRead
+            )
+        )
+        # A transaction that needs to be serializable
+        await self._try_bogus_rr_tx(con, first_try=True)
+        await self._try_bogus_rr_tx(con, first_try=False)
+
+        # And one that doesn't
+        async for tx in con.transaction():
+            async with tx:
+                res = await tx.query_single('''
+                    select {
+                        ins := (insert test::Tmp { tmp := "test" }),
+                        level := sys::get_transaction_isolation(),
+                    }
+                ''')
+            self.assertEqual(str(res.level), 'RepeatableRead')
diff --git a/tests/test_sync_tx.py b/tests/test_sync_tx.py
@@ -108,39 +108,51 @@ async def test_sync_transaction_kinds(self):
                 with tx:
                     pass
 
-    def test_sync_transaction_prefer_rr(self):
-        if (
-            str(self.server_version.stage) != 'dev'
-            and (self.server_version.major, self.server_version.minor) < (6, 5)
-        ):
-            self.skipTest("DML in RepeatableRead not supported yet")
-        con = self.client.with_transaction_options(
-            edgedb.TransactionOptions(
-                isolation=edgedb.IsolationLevel.PreferRepeatableRead
-            )
-        )
+    def _try_bogus_rr_tx(self, con, first_try):
         # A transaction that needs to be serializable
         for tx in con.transaction():
             with tx:
                 res1 = tx.query_single('''
                     select {
-                        ins := (insert test::TmpConflict { tmp := "test1" }),
+                        ins := (insert test::Tmp { tmp := "test1" }),
                         level := sys::get_transaction_isolation(),
                     }
                 ''')
+                # If this is the second time we've tried to run this
+                # transaction, then the cache should ensure we *only*
+                # try Serializable.
+                if not first_try:
+                    self.assertEqual(res1.level, 'Serializable')
 
                 res2 = tx.query_single('''
                     select {
-                        ins := (insert test::TmpConflict { tmp := "test2" }),
+                        ins := (insert test::TmpConflict {
+                            tmp := <str>random()
+                        }),
                         level := sys::get_transaction_isolation(),
                     }
                 ''')
 
                 # N.B: res1 will be RepeatableRead on the first
                 # iteration, maybe, but contingent on the second query
                 # succeeding it will be Serializable!
-                self.assertEqual(str(res1.level), 'Serializable')
-                self.assertEqual(str(res2.level), 'Serializable')
+                self.assertEqual(res1.level, 'Serializable')
+                self.assertEqual(res2.level, 'Serializable')
+
+    def test_sync_transaction_prefer_rr(self):
+        if (
+            str(self.server_version.stage) != 'dev'
+            and (self.server_version.major, self.server_version.minor) < (6, 5)
+        ):
+            self.skipTest("DML in RepeatableRead not supported yet")
+        con = self.client.with_transaction_options(
+            edgedb.TransactionOptions(
+                isolation=edgedb.IsolationLevel.PreferRepeatableRead
+            )
+        )
+        # A transaction that needs to be serializable
+        self._try_bogus_rr_tx(con, first_try=True)
+        self._try_bogus_rr_tx(con, first_try=False)
 
         # And one that doesn't
         for tx in con.transaction():
