diff --git a/mock_tests/test_auth.py b/mock_tests/test_auth.py index 192f0eb6d..cc4480c93 100644 --- a/mock_tests/test_auth.py +++ b/mock_tests/test_auth.py @@ -1,8 +1,9 @@ import asyncio import json +import threading import time import warnings -from typing import Union +from typing import List, Union import grpc import pytest @@ -84,6 +85,246 @@ def test_client_credentials(weaviate_auth_mock: HTTPServer, start_grpc_server: g weaviate_auth_mock.check_assertions() +@pytest.mark.asyncio +async def test_client_credentials_refresh_async( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """Test the refresh_session branch of the async token refresher. + + Client-credentials tokens carry no refresh token, so the refresher must get a whole + new token from the saved credentials. + """ + token_requests = 0 + + def handler(request: Request) -> Response: + nonlocal token_requests + token_requests += 1 + return Response( + json.dumps({"access_token": ACCESS_TOKEN, "expires_in": 1}), + content_type="application/json", + ) + + weaviate_auth_mock.expect_request("/auth").respond_with_handler(handler) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthClientCredentials( + client_secret=CLIENT_SECRET, scope=SCOPE + ), + ) as client: + await client.collections.list_all() + first = token_requests + await asyncio.sleep(3) # refresh interval is max(expires_in - 30, 1) -> 1s + assert token_requests > first # a fresh token was fetched with the credentials + + +def _reject_refreshes(weaviate_auth_mock: HTTPServer) -> List[float]: + """Make the IdP reject every refresh (400 invalid_grant); returns the hit timestamps.""" + hits: List[float] = [] + + def handler(request: Request) -> Response: + hits.append(time.monotonic()) + return Response( + json.dumps({"error": "invalid_grant", "error_description": "refresh token expired"}), + status=400, + content_type="application/json", + ) + + weaviate_auth_mock.expect_request("/auth").respond_with_handler(handler) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + return hits + + +@pytest.mark.asyncio +async def test_token_refresh_survives_failures_async( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, recwarn +) -> None: + """A failing refresh must not kill the refresher: warn, retry, stay alive. + + A 400 invalid_grant makes authlib raise OAuthError, which is not an httpx.HTTPError. + """ + hits = _reject_refreshes(weaviate_auth_mock) + + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, + refresh_token=REFRESH_TOKEN, + expires_in=1, # force an immediate (and failing) refresh + ), + ) as client: + task = getattr(client._connection, "_ConnectionBase__token_refresh_task") # noqa: B009 + assert task is not None + await asyncio.sleep(3) + assert not task.done() # the refresher survived the failures + await client.collections.list_all() # ... and the client still works + + assert len(hits) >= 2 # it kept retrying + # recwarn's "default" filter shows an identical warning once per location + assert len([w for w in recwarn if str(w.message).startswith("Con001")]) >= 1 + assert task.done() # close() cancelled and awaited it + + +def test_token_refresh_survives_failures( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, recwarn +) -> None: + """Sync twin of the test above. + + The daemon thread used to die silently on anything but an httpx.HTTPError. + """ + hits = _reject_refreshes(weaviate_auth_mock) + + threads_before = set(threading.enumerate()) + with weaviate.connect_to_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=1 + ), + ) as client: + refreshers = [ + t for t in set(threading.enumerate()) - threads_before if t.name == "TokenRefresh" + ] + assert len(refreshers) == 1 + time.sleep(3) + assert refreshers[0].is_alive() # survived the failures + client.collections.list_all() + + assert len(hits) >= 2 + # recwarn's "default" filter shows an identical warning once per location + assert len([w for w in recwarn if str(w.message).startswith("Con001")]) >= 1 + refreshers[0].join(timeout=2) + assert not refreshers[0].is_alive() # close() stops the daemon thread promptly + + +@pytest.mark.asyncio +async def test_async_auth_starts_no_threads( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """The async client must refresh tokens with an asyncio task, not threads. + + Under WASM/Pyodide threads cannot start at all, so the TokenRefresh daemon thread + and the event-loop sidecar thread would make every async OIDC flow crash connect(). + """ + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + weaviate_auth_mock.expect_request("/auth").respond_with_json( + { + "access_token": ACCESS_TOKEN, + "expires_in": 500, + "refresh_token": REFRESH_TOKEN, + } + ) + + # compare thread OBJECTS, not names: earlier sync tests leave stale TokenRefresh + # daemon threads alive, which would mask a regression in a name-set comparison + threads_before = set(threading.enumerate()) + tasks_before = asyncio.all_tasks() + async with weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=500 + ), + ) as client: + await client.collections.list_all() + new_thread_names = {t.name for t in set(threading.enumerate()) - threads_before} + assert "TokenRefresh" not in new_thread_names + assert "eventLoop" not in new_thread_names + refresh_tasks = [ + t for t in asyncio.all_tasks() - tasks_before if "token_refresh" in repr(t.get_coro()) + ] + assert len(refresh_tasks) == 1 # the refresher runs as an asyncio task instead + # ... and close() must cancel it AND await it: done as soon as close() returns + assert refresh_tasks[0].done() + + +def test_async_close_from_a_different_event_loop( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """close() must work when it runs on a different loop than connect(). + + Sync-first apps wrap each async step in its own asyncio.run(), so the refresh task + belongs to a loop that is already closed by the time close() runs. close() must skip + waiting for that task instead of raising out of asyncio. + """ + weaviate_auth_mock.expect_request("/auth").respond_with_json( + {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN} + ) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + + client = weaviate.use_async_with_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=500 + ), + ) + asyncio.run(client.connect()) # the refresh task is created on this run's loop + task = getattr(client._connection, "_ConnectionBase__token_refresh_task") # noqa: B009 + assert task is not None and task.cancelled() # asyncio.run cancelled it at teardown + + asyncio.run(client.close()) # a second loop; the task's loop is gone + assert not client.is_connected() + + +def test_sync_reconnect_leaves_exactly_one_refresher_thread( + weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server +) -> None: + """close() must end the daemon thread promptly, and a later connect() must not revive it. + + The thread used to re-read the connection's shutdown event on every loop, so after + close()+connect() it picked up the NEW (unset) event and kept refreshing next to the + new thread — two refreshers per client. + """ + weaviate_auth_mock.expect_request("/auth").respond_with_json( + {"access_token": ACCESS_TOKEN, "expires_in": 500, "refresh_token": REFRESH_TOKEN} + ) + weaviate_auth_mock.expect_request( + "/v1/schema", headers={"Authorization": "Bearer " + ACCESS_TOKEN} + ).respond_with_json({"classes": []}) + + def refreshers() -> List[threading.Thread]: + return [t for t in set(threading.enumerate()) - threads_before if t.name == "TokenRefresh"] + + threads_before = set(threading.enumerate()) + client = weaviate.connect_to_local( + host=MOCK_IP, + port=MOCK_PORT, + grpc_port=MOCK_PORT_GRPC, + auth_credentials=weaviate.auth.AuthBearerToken( + ACCESS_TOKEN, refresh_token=REFRESH_TOKEN, expires_in=500 + ), + ) + (first,) = refreshers() + client.close() + first.join(timeout=2) + assert not first.is_alive() # not asleep until the next (470s away) wake-up + + client.connect() + client.collections.list_all() + alive = [t for t in refreshers() if t.is_alive()] + assert len(alive) == 1 and alive[0] is not first + client.close() + alive[0].join(timeout=2) + assert not alive[0].is_alive() + + @pytest.mark.parametrize("header_name", ["Authorization", "authorization"]) def test_auth_header_priority( recwarn, weaviate_auth_mock: HTTPServer, start_grpc_server: grpc.Server, header_name: str diff --git a/weaviate/connect/event_loop.py b/weaviate/connect/event_loop.py deleted file mode 100644 index 938cb2b56..000000000 --- a/weaviate/connect/event_loop.py +++ /dev/null @@ -1,136 +0,0 @@ -import asyncio -import os -import threading -import time -from concurrent.futures import Future -from typing import Any, Callable, Coroutine, Dict, Generic, Optional, TypeVar, cast - -from typing_extensions import ParamSpec - -from weaviate.exceptions import WeaviateClosedClientError - -P = ParamSpec("P") -T = TypeVar("T") - - -class _Future(Future, Generic[T]): - def result(self, timeout: Optional[float] = None) -> T: - return cast(T, super().result(timeout)) - - -class _EventLoop: - def __init__(self, loop: Optional[asyncio.AbstractEventLoop] = None) -> None: - self.loop = loop - - def start(self) -> None: - if self.loop is not None: - return - self.loop = self.__start_new_event_loop() - _EventLoop.patch_exception_handler(self.loop) - - def run_until_complete( - self, f: Callable[P, Coroutine[Any, Any, T]], *args: P.args, **kwargs: P.kwargs - ) -> T: - """This method runs the provided coroutine in a blocking manner by scheduling its execution in an event loop running in a parallel thread. - - The result of the coroutine is returned, either when the coroutine completes or raises an exception. - """ - if self.loop is None or self.loop.is_closed(): - raise WeaviateClosedClientError() - fut = asyncio.run_coroutine_threadsafe(f(*args, **kwargs), self.loop) - return fut.result() - - def schedule( - self, f: Callable[P, Coroutine[Any, Any, T]], *args: P.args, **kwargs: P.kwargs - ) -> _Future[T]: - """This method schedules the provided coroutine for execution in the event loop running in a parallel thread. - - The coroutine will be executed asynchronously in the background. - """ - if self.loop is None or self.loop.is_closed(): - raise WeaviateClosedClientError() - return cast(_Future[T], asyncio.run_coroutine_threadsafe(f(*args, **kwargs), self.loop)) - - def shutdown(self) -> None: - if self.loop is None: - return - self.loop.call_soon_threadsafe(self.loop.stop) - - @staticmethod - def __run_event_loop(loop: asyncio.AbstractEventLoop) -> None: - try: - loop.run_forever() - finally: - # This is entered when loop.stop is scheduled from the main thread - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.close() - - @staticmethod - def __start_new_event_loop() -> asyncio.AbstractEventLoop: - loop = asyncio.new_event_loop() - - event_loop = threading.Thread( - target=_EventLoop.__run_event_loop, - daemon=True, - args=(loop,), - name="eventLoop", - ) - event_loop.start() - - while not loop.is_running(): - time.sleep(0.01) - - return loop - - @staticmethod - def patch_exception_handler(loop: asyncio.AbstractEventLoop) -> None: - """This patches the asyncio exception handler. - - It ignores the `BlockingIOError: [Errno 35] Resource temporarily unavailable` error - that is emitted by `aio.grpc` when multiple event loops are used in separate threads. This error is not actually an implementation/call error, - it's just a problem with grpc's cython implementation of `aio.Channel.__init__` whereby a `socket.recv(1)` call only works on the first call with - all subsequent calls to `aio.Channel.__init__` throwing the above error. - - This call within the `aio.Channel.__init__` method does not affect the functionality of the library and can be safely ignored. - - Context: - - https://github.com/grpc/grpc/issues/25364 - - https://github.com/grpc/grpc/pull/36096 - """ - - def exception_handler(loop: asyncio.AbstractEventLoop, context: Dict[str, Any]) -> None: - if "exception" in context: - if type( - context["exception"] - ).__name__ == "BlockingIOError" and "Resource temporarily unavailable" in str( - context["exception"] - ): - return - loop.default_exception_handler(context) - - loop.set_exception_handler(exception_handler) - - def __del__(self) -> None: - self.shutdown() - - -class _EventLoopSingleton: - _instances: Optional[Dict[int, _EventLoop]] = None - - @classmethod - def get_instance(cls) -> _EventLoop: - pid = os.getpid() - if cls._instances is not None and pid in cls._instances: - return cls._instances[pid] - if cls._instances is None: - cls._instances = {} - instance = _EventLoop() - instance.start() - cls._instances[pid] = instance - return instance - - def __del__(self) -> None: - if self._instances is not None: - for instance in self._instances.values(): - instance.shutdown() - self._instances = None diff --git a/weaviate/connect/v4.py b/weaviate/connect/v4.py index d70bfc77b..be82b47a3 100644 --- a/weaviate/connect/v4.py +++ b/weaviate/connect/v4.py @@ -62,7 +62,6 @@ JSONPayload, _get_proxies, ) -from weaviate.connect.event_loop import _EventLoopSingleton from weaviate.connect.integrations import _IntegrationConfig from weaviate.embedded import EmbeddedV4 from weaviate.exceptions import ( @@ -154,6 +153,8 @@ def __init__( self._connected = False self._skip_init_checks = skip_init_checks self._grpc_config = grpc_config + self._shutdown_background_event: Optional[Event] = None + self.__token_refresh_task: Optional["asyncio.Task[None]"] = None client_type = "sync" if isinstance(self, ConnectionSync) else "async" embedded_suffix = "-embedded" if self.embedded_db is not None else "" @@ -531,79 +532,99 @@ def _create_background_token_refresh(self, _auth: Optional[_Auth] = None) -> Non if "refresh_token" not in self._client.token and _auth is None: return - # make an event loop sidecar thread for running async token refreshing - event_loop = ( - _EventLoopSingleton.get_instance() - if isinstance(self._client, AsyncOAuth2Client) - else None - ) - - expires_in: int = self._client.token.get( - "expires_in", 60 - ) # use 1minute as token lifetime if not supplied - self._shutdown_background_event = Event() - - def refresh_token() -> None: - if isinstance(self._client, AsyncOAuth2Client): - assert event_loop is not None - self._client.token = event_loop.run_until_complete( - self._client.refresh_token, - url=self._client.metadata["token_endpoint"], - ) - elif isinstance(self._client, OAuth2Client): - self._client.token = self._client.refresh_token( - url=self._client.metadata["token_endpoint"] - ) - - def refresh_session() -> None: - assert _auth is not None - if isinstance(self._client, AsyncOAuth2Client): - assert event_loop is not None - new_session = event_loop.run_until_complete( - _auth.aresult, result=_auth.get_auth_session() - ) - self._client.token = event_loop.run_until_complete(new_session.fetch_token) - elif isinstance(self._client, OAuth2Client): - new_session = _auth.result(_auth.get_auth_session()) - self._client.token = new_session.fetch_token() - - def update_refresh_time() -> int: - assert isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)) - return self._client.token.get("expires_in", 60) - 30 + # stop the refresher from an earlier connect(), if any + self._cancel_background_token_refresh() + + # refresh 30s before the token expires (assume 1 minute if the token does not say); + # the loops below always wait at least 1s + refresh_in: int = self._client.token.get("expires_in", 60) - 30 + # the refresher keeps its own event: after close() + connect() the old refresher + # must stop on this one instead of running on with the new one + shutdown = Event() + self._shutdown_background_event = shutdown + + if isinstance(self._client, AsyncOAuth2Client): + # async client: refresh in an asyncio task on the current loop, not in a thread + # (threads cannot start under WASM/Pyodide) + self.__token_refresh_task = asyncio.get_running_loop().create_task( + self.__periodic_token_refresh_async(refresh_in, _auth, shutdown) + ) + return def periodic_refresh_token(refresh_time: int, _auth: Optional[_Auth]) -> None: - while ( - self._shutdown_background_event is not None - and not self._shutdown_background_event.is_set() - ): - # use refresh token when available - time.sleep(max(refresh_time, 1)) + # wait on the event instead of sleeping, so close() can end the thread right away + while not shutdown.wait(timeout=max(refresh_time, 1)): try: - if self._client is None: + # use one client for the whole round: close()/connect() may replace + # self._client while a refresh is running + client = self._client + if not isinstance(client, OAuth2Client): continue - elif ( - isinstance(self._client, (OAuth2Client, AsyncOAuth2Client)) - and "refresh_token" in self._client.token - ): - refresh_token() + if "refresh_token" in client.token: + client.token = client.refresh_token(url=client.metadata["token_endpoint"]) else: - # client credentials usually does not contain a refresh token => get a new token using the - # saved credentials - refresh_session() - refresh_time = update_refresh_time() - except HTTPError as exc: - # retry again after one second, might be an unstable connection + # client credentials usually does not contain a refresh token => get a + # new token using the saved credentials + assert _auth is not None + new_session = _auth.result(_auth.get_auth_session()) + client.token = new_session.fetch_token() + refresh_time = client.token.get("expires_in", 60) - 30 + except Exception as exc: + # retry in one second; any error must keep the refresher alive, not only + # network errors refresh_time = 1 _Warnings.token_refresh_failed(exc) demon = Thread( target=periodic_refresh_token, - args=(expires_in, _auth), + args=(refresh_in, _auth), daemon=True, name="TokenRefresh", ) demon.start() + def _cancel_background_token_refresh(self) -> Optional["asyncio.Task[None]"]: + """Stop the token refresher: set the shutdown event (sync thread) and cancel the async task. + + Returns the cancelled task, if any, so close() can wait for it to finish. + """ + if self._shutdown_background_event is not None: + self._shutdown_background_event.set() + task, self.__token_refresh_task = self.__token_refresh_task, None + if task is not None: + try: + task.cancel() + except RuntimeError: + # the task's loop is closed; the task can never run again + pass + return task + + async def __periodic_token_refresh_async( + self, refresh_time: int, _auth: Optional[_Auth], shutdown: Event + ) -> None: + """Async version of ``periodic_refresh_token``, run as a task; close() cancels it.""" + while not shutdown.is_set(): + await asyncio.sleep(max(refresh_time, 1)) + try: + client = self._client + if not isinstance(client, AsyncOAuth2Client): + continue + if "refresh_token" in client.token: + client.token = await client.refresh_token(url=client.metadata["token_endpoint"]) + else: + # client credentials usually does not contain a refresh token => get a + # new token using the saved credentials + assert _auth is not None + new_session = await _Auth.aresult(_auth.get_auth_session()) + client.token = await new_session.fetch_token() + refresh_time = client.token.get("expires_in", 60) - 30 + except asyncio.CancelledError: + raise + except Exception as exc: + # retry in one second; any error must keep the refresher alive + refresh_time = 1 + _Warnings.token_refresh_failed(exc) + def __get_latest_headers(self) -> Dict[str, str]: if "authorization" in self._headers: return self._headers @@ -711,9 +732,17 @@ def exc(e: Exception) -> None: def close(self, colour: executor.Colour) -> executor.Result[None]: if self.embedded_db is not None: self.embedded_db.stop() + refresh_task = self._cancel_background_token_refresh() if colour == "async": async def execute() -> None: + if ( + refresh_task is not None + and refresh_task.get_loop() is asyncio.get_running_loop() + ): + # wait for the task to finish before the client is closed; a task from + # another loop cannot be awaited here and is done or orphaned with its loop + await asyncio.gather(refresh_task, return_exceptions=True) if self._client is not None: assert isinstance(self._client, AsyncClient) await self._client.aclose()