From 52316b1722742c92f7667ea566987358d0f9fbed Mon Sep 17 00:00:00 2001 From: Praveen Mittal Date: Thu, 15 Oct 2026 09:20:16 +0200 Subject: [PATCH] feat(core): add KeyAuth.close() to stop the token refresh thread Each KeyAuth started a daemon thread that refreshed tokens forever, with no way to stop it, so a discarded instance kept running and requesting tokens. close() stops the loop, shuts down the refresh executor and removes the atexit hook. It is safe to call more than once. --- CHANGELOG.md | 2 + core/CHANGELOG.md | 3 + core/pyproject.toml | 2 +- .../src/stackit/core/auth_methods/key_auth.py | 23 ++++++- core/tests/core/test_auth.py | 61 +++++++++++++++++++ core/uv.lock | 2 +- 6 files changed, 89 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9becc1388..5bd78dae2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,8 @@ ## Release (2026-MM-DD) - `core`: + - [v0.4.0](core/CHANGELOG.md#v040) + - **Feature:** Add `KeyAuth.close()` to stop the background token refresh thread and its worker, so a discarded instance no longer keeps running and requesting tokens - [v0.3.1](core/CHANGELOG.md#v031) - **Fix:** Refresh access tokens five minutes before expiration instead of after, and wait for in-progress refreshes before authorizing requests - `alb`: diff --git a/core/CHANGELOG.md b/core/CHANGELOG.md index b41b2470c..e671d6cd4 100644 --- a/core/CHANGELOG.md +++ b/core/CHANGELOG.md @@ -1,3 +1,6 @@ +## v0.4.0 +- **Feature:** Add `KeyAuth.close()` to stop the background token refresh thread and its worker, so a discarded instance (for example after key rotation) no longer keeps running and requesting tokens + ## v0.3.1 - **Fix:** Refresh access tokens five minutes before expiration instead of after, and wait for in-progress refreshes before authorizing requests diff --git a/core/pyproject.toml b/core/pyproject.toml index 0bac65e7b..18350e95b 100644 --- a/core/pyproject.toml +++ b/core/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "stackit-core" -version = "v0.3.1" +version = "v0.4.0" description = "Core functionality for the STACKIT SDK for Python" authors = [{ name = "STACKIT Developer Tools", email = "developer-tools@stackit.cloud" }] requires-python = ">=3.10,<4.0" diff --git a/core/src/stackit/core/auth_methods/key_auth.py b/core/src/stackit/core/auth_methods/key_auth.py index 6f24eaff2..bb6da2c60 100644 --- a/core/src/stackit/core/auth_methods/key_auth.py +++ b/core/src/stackit/core/auth_methods/key_auth.py @@ -54,6 +54,8 @@ class KeyAuth(AuthBase): executor: ThreadPoolExecutor refresh_future: Optional[futures.Future] service_account_key: ServiceAccountKey + _stop: threading.Event + _refresh_thread: Optional[threading.Thread] def __init__( self, @@ -66,6 +68,8 @@ def __init__( self.executor = ThreadPoolExecutor(max_workers=1) self.access_token = None self.refresh_future = None + self._stop = threading.Event() + self._refresh_thread = None self.__create_initial_token() self.__fetch_token_from_endpoint() self.__start_token_refresh_task() @@ -92,6 +96,21 @@ def __call__(self, r: Request) -> Request: r.headers["Authorization"] = f"Bearer {self.access_token}" return r + def close(self) -> None: + """Stop the background token refresh and release its worker. + + Call this when the instance is no longer needed, for example after + replacing it during key rotation in a long-running process. Without it + the refresh thread keeps running and keeps requesting tokens. Safe to + call more than once. + """ + self._stop.set() + self.executor.shutdown(wait=False, cancel_futures=True) + atexit.unregister(self.__shutdown) + thread = self._refresh_thread + if thread is not None and thread is not threading.current_thread(): + thread.join() + def __create_initial_token(self) -> None: payload = { "iss": self.service_account_key.credentials.issuer, @@ -126,8 +145,7 @@ def __fetch_token_from_endpoint(self) -> None: def __start_token_refresh_task(self): def token_refresh_task(): - while True: - time.sleep(self.TOKEN_EXPIRY_CHECK_INTERVAL.total_seconds()) + while not self._stop.wait(self.TOKEN_EXPIRY_CHECK_INTERVAL.total_seconds()): with self.lock: if self.refresh_future is not None and self.refresh_future.done(): self.refresh_future.result() @@ -139,6 +157,7 @@ def token_refresh_task(): thread = threading.Thread(target=token_refresh_task) thread.daemon = True thread.start() + self._refresh_thread = thread def __refresh_token(self): with self.lock: diff --git a/core/tests/core/test_auth.py b/core/tests/core/test_auth.py index d73009f10..903e5755c 100644 --- a/core/tests/core/test_auth.py +++ b/core/tests/core/test_auth.py @@ -2,6 +2,9 @@ from pathlib import Path, PurePath from threading import Event, Lock, Thread +import time +from datetime import timedelta + import pytest import json import jwt @@ -382,3 +385,61 @@ def result(self): assert not call_thread.is_alive() assert request.headers["Authorization"] == f"Bearer {fresh_token}" + + +class TestKeyAuthClose: + """A discarded KeyAuth must be able to stop its background token refresh.""" + + @staticmethod + def _make_auth(service_account_key_file_json, mock_post, access_token): + service_account_key = ServiceAccountKey.model_validate_json(service_account_key_file_json) + service_account_key.credentials.private_key = "test-private-key" + response = Mock() + response.raise_for_status.return_value = None + response.json.return_value = { + "access_token": access_token, + "refresh_token": jwt.encode({"exp": 4102444800}, "secret", algorithm="HS256"), + } + mock_post.return_value = response + return KeyAuth(service_account_key) + + def test_close_stops_the_refresh_thread_and_further_token_requests(self, service_account_key_file_json): + def set_initial_token(auth): + auth.initial_token = "test-initial-token" + + expired = jwt.encode({"exp": 1}, "secret", algorithm="HS256") + with ( + patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token), + patch.object(KeyAuth, "TOKEN_EXPIRY_CHECK_INTERVAL", timedelta(milliseconds=10)), + patch("requests.post") as mock_post, + ): + # An access token that is always expired keeps the refresh loop busy. + auth = self._make_auth(service_account_key_file_json, mock_post, expired) + thread = auth._refresh_thread + assert thread is not None and thread.is_alive() + + auth.close() + + assert not thread.is_alive() + requests_after_close = mock_post.call_count + time.sleep(0.1) # ten refresh intervals + assert mock_post.call_count == requests_after_close + + def test_close_is_idempotent(self, service_account_key_file_json): + def set_initial_token(auth): + auth.initial_token = "test-initial-token" + + valid = jwt.encode({"exp": 4102444800}, "secret", algorithm="HS256") + with ( + patch.object(KeyAuth, "_KeyAuth__create_initial_token", new=set_initial_token), + patch("requests.post") as mock_post, + patch("stackit.core.auth_methods.key_auth.atexit") as mock_atexit, + ): + auth = self._make_auth(service_account_key_file_json, mock_post, valid) + + auth.close() + auth.close() + + assert not auth._refresh_thread.is_alive() + assert auth.executor._shutdown + mock_atexit.unregister.assert_called_with(auth._KeyAuth__shutdown) diff --git a/core/uv.lock b/core/uv.lock index 59c9b39f5..d89822b6a 100644 --- a/core/uv.lock +++ b/core/uv.lock @@ -998,7 +998,7 @@ wheels = [ [[package]] name = "stackit-core" -version = "0.3.1" +version = "0.4.0" source = { editable = "." } dependencies = [ { name = "cryptography" },