PYTHON-2560 Retry KMS requests on transient errors

Author: ShaneHarvey

pymongo/asynchronous/encryption.py

@@ -72,7 +72,6 @@
     EncryptedCollectionError,
     EncryptionError,
     InvalidOperation,
-    PyMongoError,
     ServerSelectionTimeoutError,
 )
 from pymongo.network_layer import BLOCKING_IO_ERRORS, async_sendall
@@ -176,6 +175,9 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             ssl_context=ctx,
         )
         host, port = parse_host(endpoint, _HTTPS_PORT)
+        sleep_usec = kms_context.usleep
+        if sleep_usec:
+            await asyncio.sleep(float(sleep_usec) / 1e6)
         try:
             conn = await _configured_socket((host, port), opts)
             try:
@@ -201,13 +203,15 @@ async def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 raise socket.timeout("timed out") from None
             finally:
                 conn.close()
-        except (PyMongoError, MongoCryptError):
-            raise  # Propagate pymongo errors directly.
-        except asyncio.CancelledError:
-            raise
-        except Exception as error:
-            # Wrap I/O errors in PyMongo exceptions.
-            _raise_connection_failure((host, port), error)
+        except MongoCryptError:
+            raise  # Propagate MongoCryptError errors directly.
+        except Exception as exc:
+            remaining = _csot.remaining()
+            if remaining is not None and remaining <= 0:
+                # Wrap I/O errors in PyMongo exceptions.
+                _raise_connection_failure((host, port), exc)
+            # Mark this attempt as failed and defer to libmongocrypt to retry.
+            kms_context.fail()
 
     async def collection_info(self, database: str, filter: bytes) -> Optional[bytes]:
         """Get the collection info for a namespace.

pymongo/synchronous/encryption.py

@@ -67,7 +67,6 @@
     EncryptedCollectionError,
     EncryptionError,
     InvalidOperation,
-    PyMongoError,
     ServerSelectionTimeoutError,
 )
 from pymongo.network_layer import BLOCKING_IO_ERRORS, sendall
@@ -176,6 +175,9 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
             ssl_context=ctx,
         )
         host, port = parse_host(endpoint, _HTTPS_PORT)
+        sleep_usec = kms_context.usleep
+        if sleep_usec:
+            asyncio.sleep(float(sleep_usec) / 1e6)
         try:
             conn = _configured_socket((host, port), opts)
             try:
@@ -201,13 +203,15 @@ def kms_request(self, kms_context: MongoCryptKmsContext) -> None:
                 raise socket.timeout("timed out") from None
             finally:
                 conn.close()
-        except (PyMongoError, MongoCryptError):
-            raise  # Propagate pymongo errors directly.
-        except asyncio.CancelledError:
-            raise
-        except Exception as error:
-            # Wrap I/O errors in PyMongo exceptions.
-            _raise_connection_failure((host, port), error)
+        except MongoCryptError:
+            raise  # Propagate MongoCryptError errors directly.
+        except Exception as exc:
+            remaining = _csot.remaining()
+            if remaining is not None and remaining <= 0:
+                # Wrap I/O errors in PyMongo exceptions.
+                _raise_connection_failure((host, port), exc)
+            # Mark this attempt as failed and defer to libmongocrypt to retry.
+            kms_context.fail()
 
     def collection_info(self, database: str, filter: bytes) -> Optional[bytes]:
         """Get the collection info for a namespace.