feat: improve reliability of certificate refresh (#599)

Author: jackwotherspoon

google/cloud/sql/connector/instance.py

@@ -407,10 +407,8 @@ async def _refresh_task(self: Instance, delay: int) -> InstanceMetadata:
                 raise
             # if valid refresh, replace current with valid metadata and schedule next refresh
             self._current = refresh_task
-            # Ephemeral certificate expires in 1 hour, so we schedule a refresh to happen in 55 minutes.
-            delay = _seconds_until_refresh(
-                refresh_data.expiration, self._enable_iam_auth
-            )
+            # calculate refresh delay based on certificate expiration
+            delay = _seconds_until_refresh(refresh_data.expiration)
             self._next = self._schedule_refresh(delay)
 
             return refresh_data

google/cloud/sql/connector/refresh_utils.py

@@ -29,15 +29,9 @@
 
 _sql_api_version: str = "v1beta4"
 
-# default_refresh_buffer is the amount of time before a refresh's result expires
+# _refresh_buffer is the amount of time before a refresh's result expires
 # that a new refresh operation begins.
-_default_refresh_buffer: int = 5 * 60  # 5 minutes
-
-# _iam_auth_refresh_buffer is the amount of time before a refresh's result expires
-# that a new refresh operation begins when IAM DB AuthN is enabled. Because token
-# sources may be cached until ~60 seconds before expiration, this value must be smaller
-# than default_refresh_buffer.
-_iam_auth_refresh_buffer: int = 55  # seconds
+_refresh_buffer: int = 4 * 60  # 4 minutes
 
 
 async def _get_metadata(
@@ -194,29 +188,28 @@ async def _get_ephemeral(
 
 def _seconds_until_refresh(
     expiration: datetime.datetime,
-    enable_iam_auth: bool,
 ) -> int:
     """
-    Helper function to get time in seconds before performing next refresh.
+    Calculates the duration to wait before starting the next refresh.
+
+    Usually the duration will be half of the time until certificate
+    expiration.
 
     :rtype: int
     :returns: Time in seconds to wait before performing next refresh.
     """
-    if enable_iam_auth:
-        refresh_buffer = _iam_auth_refresh_buffer
-    else:
-        refresh_buffer = _default_refresh_buffer
 
-    delay = (expiration - datetime.datetime.now()) - datetime.timedelta(
-        seconds=refresh_buffer
-    )
+    duration = int((expiration - datetime.datetime.now()).total_seconds())
 
-    if delay.total_seconds() < 0:
-        # If the time until the certificate expires is less than the buffer,
-        # schedule the refresh closer to the expiration time
-        delay = (expiration - datetime.datetime.now()) - datetime.timedelta(seconds=5)
+    # if certificate duration is less than 1 hour
+    if duration < 3600:
+        # something is wrong with certificate, refresh now
+        if duration < _refresh_buffer:
+            return 0
+        # otherwise wait until 4 minutes before expiration for next refresh
+        return duration - _refresh_buffer
 
-    return int(delay.total_seconds())
+    return duration // 2
 
 
 async def _is_valid(task: asyncio.Task) -> bool:

tests/unit/test_refresh_utils.py

@@ -1,4 +1,4 @@
-""""
+"""
 Copyright 2021 Google LLC
 
 Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,7 +14,7 @@
 limitations under the License.
 """
 from typing import Any, no_type_check
-
+from datetime import datetime, timedelta
 import aiohttp
 from google.auth.credentials import Credentials
 import google.oauth2.credentials
@@ -28,6 +28,7 @@
     _get_metadata,
     _is_valid,
     _downscope_credentials,
+    _seconds_until_refresh,
 )
 from google.cloud.sql.connector.utils import generate_keys
 
@@ -37,7 +38,7 @@
     instance_metadata_expired,
     FakeCSQLInstance,
 )
-from tests.conftest import SCOPES  # type: ignore
+from conftest import SCOPES  # type: ignore
 
 
 @pytest.fixture
@@ -265,3 +266,39 @@ def test_downscope_credentials_user() -> None:
     # verify downscoped credentials have new scope
     assert credentials.scopes == ["https://www.googleapis.com/auth/sqlservice.login"]
     assert credentials != creds
+
+
+def test_seconds_until_refresh_over_1_hour() -> None:
+    """
+    Test _seconds_until_refresh returns proper time in seconds.
+
+    If expiration is over 1 hour, should return duration/2.
+    """
+    # using pytest.approx since sometimes can be off by a second
+    assert (
+        pytest.approx(_seconds_until_refresh(datetime.now() + timedelta(minutes=62)), 1)
+        == 31 * 60
+    )
+
+
+def test_seconds_until_refresh_under_1_hour_over_4_mins() -> None:
+    """
+    Test _seconds_until_refresh returns proper time in seconds.
+
+    If expiration is under 1 hour and over 4 minutes,
+    should return duration-refresh_buffer (refresh_buffer = 4 minutes).
+    """
+    # using pytest.approx since sometimes can be off by a second
+    assert (
+        pytest.approx(_seconds_until_refresh(datetime.now() + timedelta(minutes=5)), 1)
+        == 60
+    )
+
+
+def test_seconds_until_refresh_under_4_mins() -> None:
+    """
+    Test _seconds_until_refresh returns proper time in seconds.
+
+    If expiration is under 4 minutes, should return 0.
+    """
+    assert _seconds_until_refresh(datetime.now() + timedelta(minutes=3)) == 0