PYTHON-3345 CSOT use connection handshake RTT for load balanced mode (#998)

Author: ShaneHarvey

pymongo/mongo_client.py

@@ -1336,6 +1336,11 @@ def is_retrying():
                 bulk.started_retryable_write = True
 
         while True:
+            if is_retrying():
+                remaining = _csot.remaining()
+                if remaining is not None and remaining <= 0:
+                    assert last_error is not None
+                    raise last_error
             try:
                 server = self._select_server(writable_server_selector, session)
                 supports_session = (
@@ -1394,6 +1399,11 @@ def _retryable_read(self, func, read_pref, session, address=None, retryable=True
         multiple_retries = _csot.get_timeout() is not None
 
         while True:
+            if retrying:
+                remaining = _csot.remaining()
+                if remaining is not None and remaining <= 0:
+                    assert last_error is not None
+                    raise last_error
             try:
                 server = self._select_server(read_pref, session, address=address)
                 with self._socket_from_server(read_pref, server, session) as (sock_info, read_pref):

pymongo/pool.py

@@ -559,6 +559,7 @@ def __init__(self, sock, pool, address, id):
         self.pinned_cursor = False
         self.active = False
         self.last_timeout = self.opts.socket_timeout
+        self.connect_rtt = 0.0
 
     def set_socket_timeout(self, timeout):
         """Cache last timeout to avoid duplicate calls to sock.settimeout."""
@@ -580,6 +581,8 @@ def apply_timeout(self, client, cmd, write_concern=None):
             return None
         # RTT validation.
         rtt = _csot.get_rtt()
+        if rtt is None:
+            rtt = self.connect_rtt
         max_time_ms = timeout - rtt
         if max_time_ms < 0:
             # CSOT: raise an error without running the command since we know it will time out.
@@ -655,7 +658,11 @@ def _hello(self, cluster_time, topology_version, heartbeat_frequency):
         else:
             auth_ctx = None
 
+        if performing_handshake:
+            start = time.monotonic()
         doc = self.command("admin", cmd, publish_events=False, exhaust_allowed=awaitable)
+        if performing_handshake:
+            self.connect_rtt = time.monotonic() - start
         hello = Hello(doc, awaitable=awaitable)
         self.is_writable = hello.is_writable
         self.max_wire_version = hello.max_wire_version

test/test_csot.py

@@ -33,6 +33,9 @@
 
 
 class TestCSOT(IntegrationTest):
+    RUN_ON_SERVERLESS = True
+    RUN_ON_LOAD_BALANCER = True
+
     def test_timeout_nested(self):
         coll = self.db.coll
         self.assertEqual(_csot.get_timeout(), None)