add update kube token task (#301)

zubenkoivan · web-flow · commit 99c0f4688e04 · 2023-07-18T20:50:04.000+03:00
diff --git a/platform_container_runtime/config.py b/platform_container_runtime/config.py
@@ -27,6 +27,7 @@ class KubeConfig:
     client_key_path: Optional[str] = None
     token: Optional[str] = None
     token_path: Optional[str] = None
+    token_update_interval_s: int = 300
     conn_force_close: bool = False
     conn_timeout_s: int = 300
     read_timeout_s: int = 100
diff --git a/platform_container_runtime/kube_client.py b/platform_container_runtime/kube_client.py
@@ -1,6 +1,8 @@
+import asyncio
 import logging
 import ssl
 from collections.abc import Sequence
+from contextlib import suppress
 from dataclasses import dataclass, field
 from pathlib import Path
 from types import TracebackType
@@ -13,7 +15,7 @@
 logger = logging.getLogger(__name__)
 
 
-class KubeClientUnautorized(Exception):
+class KubeClientUnauthorized(Exception):
     pass
 
 
@@ -60,6 +62,7 @@ def __init__(
         self._token = config.token
         self._trace_configs = trace_configs
         self._client: Optional[aiohttp.ClientSession] = None
+        self._token_updater_task: Optional[asyncio.Task[None]] = None
 
     def _create_ssl_context(self) -> Optional[ssl.SSLContext]:
         if self._config.url.scheme != "https":
@@ -76,7 +79,7 @@ def _create_ssl_context(self) -> Optional[ssl.SSLContext]:
         return ssl_context
 
     async def __aenter__(self) -> "KubeClient":
-        self._client = await self._create_http_client()
+        await self._init()
         return self
 
     async def __aexit__(
@@ -87,61 +90,64 @@ async def __aexit__(
     ) -> None:
         await self.aclose()
 
-    async def _create_http_client(self) -> aiohttp.ClientSession:
+    async def _init(self) -> None:
         connector = aiohttp.TCPConnector(
             limit=self._config.conn_pool_size,
             force_close=self._config.conn_force_close,
             ssl=self._create_ssl_context(),
         )
-        if self._config.auth_type == KubeClientAuthType.TOKEN:
-            token = self._token
-            if not token:
-                assert self._config.token_path is not None
-                token = Path(self._config.token_path).read_text()
-            headers = {"Authorization": "Bearer " + token}
-        else:
-            headers = {}
+        if self._config.token_path:
+            self._token = Path(self._config.token_path).read_text()
+            self._token_updater_task = asyncio.create_task(self._start_token_updater())
         timeout = aiohttp.ClientTimeout(
             connect=self._config.conn_timeout_s, total=self._config.read_timeout_s
         )
-        return aiohttp.ClientSession(
+        self._client = aiohttp.ClientSession(
             connector=connector,
             timeout=timeout,
-            headers=headers,
             trace_configs=self._trace_configs,
         )
 
-    async def _reload_http_client(self) -> None:
-        if self._client:
-            await self._client.close()
-        self._token = None
-        self._client = await self._create_http_client()
-
-    async def init_if_needed(self) -> None:
-        if not self._client or self._client.closed:
-            await self._reload_http_client()
+    async def _start_token_updater(self) -> None:
+        if not self._config.token_path:
+            return
+        while True:
+            try:
+                token = Path(self._config.token_path).read_text()
+                if token != self._token:
+                    self._token = token
+                    logger.info("Kube token was refreshed")
+            except asyncio.CancelledError:
+                raise
+            except Exception as exc:
+                logger.exception("Failed to update kube token: %s", exc)
+            await asyncio.sleep(self._config.token_update_interval_s)
 
     async def aclose(self) -> None:
-        assert self._client
-        await self._client.close()
-
-    async def request(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
-        await self.init_if_needed()
-        assert self._client, "client is not intialized"
-        doing_retry = kwargs.pop("doing_retry", False)
-
-        async with self._client.request(*args, **kwargs) as resp:
+        if self._client:
+            await self._client.close()
+            self._client = None
+        if self._token_updater_task:
+            self._token_updater_task.cancel()
+            with suppress(asyncio.CancelledError):
+                await self._token_updater_task
+            self._token_updater_task = None
+
+    def _create_headers(
+        self, headers: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        headers = dict(headers) if headers else {}
+        if self._config.auth_type == KubeClientAuthType.TOKEN and self._token:
+            headers["Authorization"] = "Bearer " + self._token
+        return headers
+
+    async def _request(self, *args: Any, **kwargs: Any) -> dict[str, Any]:
+        headers = self._create_headers(kwargs.pop("headers", None))
+        assert self._client, "client is not initialized"
+        async with self._client.request(*args, headers=headers, **kwargs) as resp:
             resp_payload = await resp.json()
-        try:
             self._raise_for_status(resp_payload)
             return resp_payload
-        except KubeClientUnautorized:
-            if doing_retry:
-                raise
-            # K8s SA's token might be stale, need to refresh it and retry
-            await self._reload_http_client()
-            kwargs["doing_retry"] = True
-            return await self.request(*args, **kwargs)
 
     def _raise_for_status(self, payload: dict[str, Any]) -> None:
         kind = payload["kind"]
@@ -150,18 +156,18 @@ def _raise_for_status(self, payload: dict[str, Any]) -> None:
                 return
             code = payload.get("code")
             if code == 401:
-                raise KubeClientUnautorized(payload)
+                raise KubeClientUnauthorized(payload)
             raise KubeClientException(payload)
 
     async def get_nodes(self) -> Sequence[Node]:
-        payload = await self.request(
+        payload = await self._request(
             method="get", url=self._config.url / "api/v1/nodes"
         )
         assert payload["kind"] == "NodeList"
         return [Node.from_payload(p) for p in payload["items"]]
 
     async def get_node(self, name: str) -> Node:
-        payload = await self.request(
+        payload = await self._request(
             method="get", url=self._config.url / "api/v1/nodes" / name
         )
         assert payload["kind"] == "Node"
diff --git a/tests/integration/test_kube_client.py b/tests/integration/test_kube_client.py
@@ -1,5 +1,82 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import tempfile
+from collections.abc import AsyncIterator, Iterator
+from pathlib import Path
+from typing import Any
+
+import aiohttp
+import aiohttp.web
+import pytest
+from yarl import URL
+
+from platform_container_runtime.config import KubeClientAuthType, KubeConfig
 from platform_container_runtime.kube_client import KubeClient
 
+from .conftest import create_local_app_server
+
+
+class TestKubeClientTokenUpdater:
+    @pytest.fixture
+    async def kube_app(self) -> aiohttp.web.Application:
+        async def _get_nodes(request: aiohttp.web.Request) -> aiohttp.web.Response:
+            auth = request.headers["Authorization"]
+            token = auth.split()[-1]
+            app["token"]["value"] = token
+            return aiohttp.web.json_response({"kind": "NodeList", "items": []})
+
+        app = aiohttp.web.Application()
+        app["token"] = {"value": ""}
+        app.router.add_routes([aiohttp.web.get("/api/v1/nodes", _get_nodes)])
+        return app
+
+    @pytest.fixture
+    async def kube_server(
+        self, kube_app: aiohttp.web.Application, unused_tcp_port_factory: Any
+    ) -> AsyncIterator[str]:
+        async with create_local_app_server(
+            kube_app, port=unused_tcp_port_factory()
+        ) as address:
+            yield f"http://{address.host}:{address.port}"
+
+    @pytest.fixture
+    def kube_token_path(self) -> Iterator[str]:
+        _, path = tempfile.mkstemp()
+        Path(path).write_text("token-1")
+        yield path
+        os.remove(path)
+
+    @pytest.fixture
+    async def kube_client(
+        self, kube_server: str, kube_token_path: str
+    ) -> AsyncIterator[KubeClient]:
+        async with KubeClient(
+            config=KubeConfig(
+                url=URL(kube_server),
+                auth_type=KubeClientAuthType.TOKEN,
+                token_path=kube_token_path,
+                token_update_interval_s=1,
+            )
+        ) as client:
+            yield client
+
+    async def test_token_periodically_updated(
+        self,
+        kube_app: aiohttp.web.Application,
+        kube_client: KubeClient,
+        kube_token_path: str,
+    ) -> None:
+        await kube_client.get_nodes()
+        assert kube_app["token"]["value"] == "token-1"
+
+        Path(kube_token_path).write_text("token-2")
+        await asyncio.sleep(2)
+
+        await kube_client.get_nodes()
+        assert kube_app["token"]["value"] == "token-2"
+
 
 class TestKubeClient:
     async def test_get_node(self, kube_client: KubeClient) -> None:
