|
1 | 1 | import asyncio |
2 | 2 | import logging |
3 | 3 | from abc import ABC, abstractmethod |
4 | | -from collections.abc import AsyncIterator, Callable, Sequence |
| 4 | +from collections.abc import AsyncIterator, Callable |
5 | 5 | from contextlib import asynccontextmanager |
| 6 | +from typing import Self |
6 | 7 |
|
7 | 8 | from aiorwlock import RWLock |
| 9 | +from apolo_events_client import ( |
| 10 | + AbstractEventsClient, |
| 11 | + EventType, |
| 12 | + FilterItem, |
| 13 | + RecvEvent, |
| 14 | + StreamType, |
| 15 | +) |
8 | 16 |
|
9 | 17 | from .cluster_config import ClusterConfig |
10 | | -from .config import Config |
11 | 18 | from .config_client import ConfigClient |
12 | 19 | from .orchestrator.base import Orchestrator |
13 | | -from .utils.update_notifier import Notifier |
14 | 20 |
|
15 | 21 | logger = logging.getLogger(__name__) |
16 | 22 |
|
@@ -59,84 +65,102 @@ def orchestrator(self) -> Orchestrator: # pragma: no cover |
59 | 65 |
|
60 | 66 |
|
61 | 67 | class ClusterUpdater: |
| 68 | + _CONFIG_STREAM = StreamType("platform-config") |
| 69 | + _CLUSTER_ADD_EVENT = EventType("cluster-add") |
| 70 | + _CLUSTER_UPDATE_EVENT = EventType("cluster-update") |
| 71 | + _CLUSTER_REMOVE_EVENT = EventType("cluster-remove") |
| 72 | + |
62 | 73 | def __init__( |
63 | 74 | self, |
64 | | - notifier: Notifier, |
| 75 | + events_client: AbstractEventsClient, |
65 | 76 | cluster_registry: "ClusterConfigRegistry", |
66 | | - config: Config, |
67 | 77 | config_client: ConfigClient, |
68 | 78 | ): |
69 | | - self._loop = asyncio.get_event_loop() |
70 | | - self._notifier = notifier |
| 79 | + self._events_client = events_client |
71 | 80 | self._cluster_registry = cluster_registry |
72 | | - self._config = config |
73 | 81 | self._config_client = config_client |
74 | 82 |
|
75 | 83 | self._is_active: asyncio.Future[None] | None = None |
76 | 84 | self._task: asyncio.Future[None] | None = None |
77 | 85 |
|
78 | | - async def start(self) -> None: |
79 | | - logger.info("Starting Cluster Updater") |
80 | | - await self._init_task() |
81 | | - |
82 | | - async def __aenter__(self) -> "ClusterUpdater": |
83 | | - await self.start() |
| 86 | + async def __aenter__(self) -> Self: |
| 87 | + logger.info("Subscribe for %r", self._CONFIG_STREAM) |
| 88 | + await self._events_client.subscribe_group( |
| 89 | + self._CONFIG_STREAM, |
| 90 | + self._on_event, |
| 91 | + filters=[ |
| 92 | + FilterItem( |
| 93 | + event_types=frozenset( |
| 94 | + [ |
| 95 | + self._CLUSTER_ADD_EVENT, |
| 96 | + self._CLUSTER_UPDATE_EVENT, |
| 97 | + self._CLUSTER_REMOVE_EVENT, |
| 98 | + ] |
| 99 | + ) |
| 100 | + ) |
| 101 | + ], |
| 102 | + ) |
| 103 | + logger.info("Subscribed") |
84 | 104 | return self |
85 | 105 |
|
86 | | - async def __aexit__(self, *args: object) -> None: |
87 | | - await self.stop() |
88 | | - |
89 | | - async def _init_task(self) -> None: |
90 | | - assert not self._is_active |
91 | | - assert not self._task |
92 | | - |
93 | | - self._is_active = self._loop.create_future() |
94 | | - self._task = asyncio.ensure_future(self._run()) |
95 | | - # forcing execution of the newly created task |
96 | | - await asyncio.sleep(0) |
97 | | - |
98 | | - async def stop(self) -> None: |
99 | | - logger.info("Stopping Cluster Updater") |
100 | | - assert self._is_active is not None |
101 | | - self._is_active.set_result(None) |
102 | | - |
103 | | - assert self._task |
104 | | - await self._task |
105 | | - |
106 | | - self._task = None |
107 | | - self._is_active = None |
108 | | - |
109 | | - async def _run(self) -> None: |
110 | | - assert self._is_active is not None |
111 | | - |
112 | | - def _listener() -> None: |
113 | | - self._loop.create_task(self._do_update()) |
| 106 | + async def __aexit__(self, exc_typ: object, exc_val: object, exc_tb: object) -> None: |
| 107 | + pass |
114 | 108 |
|
115 | | - async with self._notifier.listen_to_updates(_listener): |
116 | | - await self._is_active |
| 109 | + async def _on_event(self, ev: RecvEvent) -> None: |
| 110 | + assert ev.cluster, "event cluster is required" |
117 | 111 |
|
118 | | - async def _do_update(self) -> None: |
119 | | - cluster_configs = await self._config_client.get_clusters() |
120 | | - cluster_registry = self._cluster_registry |
121 | | - for cluster_config in cluster_configs: |
122 | | - await cluster_registry.replace(cluster_config) |
123 | | - await cluster_registry.cleanup(cluster_configs) |
| 112 | + if ( |
| 113 | + ev.event_type == self._CLUSTER_UPDATE_EVENT |
| 114 | + or ev.event_type == self._CLUSTER_ADD_EVENT |
| 115 | + ): |
| 116 | + cluster_config = await self._config_client.get_cluster(ev.cluster) |
| 117 | + if cluster_config: |
| 118 | + await self._cluster_registry.replace(cluster_config) |
| 119 | + else: |
| 120 | + logger.warning("Cluster %r not found", ev.cluster) |
| 121 | + if ev.event_type == self._CLUSTER_REMOVE_EVENT: |
| 122 | + self._cluster_registry.remove(ev.cluster) |
124 | 123 |
|
125 | 124 |
|
126 | 125 | class SingleClusterUpdater: |
| 126 | + _CONFIG_STREAM = StreamType("platform-config") |
| 127 | + _CLUSTER_UPDATE_EVENT = EventType("cluster-update") |
| 128 | + |
127 | 129 | def __init__( |
128 | 130 | self, |
129 | | - cluster_holder: "ClusterHolder", |
| 131 | + events_client: AbstractEventsClient, |
130 | 132 | config_client: ConfigClient, |
| 133 | + cluster_holder: "ClusterHolder", |
131 | 134 | cluster_name: str, |
132 | 135 | ): |
133 | | - self._loop = asyncio.get_event_loop() |
134 | | - self._cluster_holder = cluster_holder |
| 136 | + self._events_client = events_client |
135 | 137 | self._config_client = config_client |
| 138 | + self._cluster_holder = cluster_holder |
136 | 139 | self._cluster_name = cluster_name |
137 | 140 |
|
138 | 141 | self.disable_updates_for_test = False |
139 | 142 |
|
| 143 | + async def __aenter__(self) -> Self: |
| 144 | + logger.info("Subscribe for %r", self._CONFIG_STREAM) |
| 145 | + await self._events_client.subscribe_group( |
| 146 | + self._CONFIG_STREAM, |
| 147 | + self._on_event, |
| 148 | + filters=[ |
| 149 | + FilterItem( |
| 150 | + event_types=frozenset([self._CLUSTER_UPDATE_EVENT]), |
| 151 | + clusters=frozenset([self._cluster_name]), |
| 152 | + ) |
| 153 | + ], |
| 154 | + ) |
| 155 | + logger.info("Subscribed") |
| 156 | + return self |
| 157 | + |
| 158 | + async def __aexit__(self, exc_typ: object, exc_val: object, exc_tb: object) -> None: |
| 159 | + pass |
| 160 | + |
| 161 | + async def _on_event(self, _: RecvEvent) -> None: |
| 162 | + await self.do_update() |
| 163 | + |
140 | 164 | async def do_update(self) -> None: |
141 | 165 | if self.disable_updates_for_test: |
142 | 166 | return |
@@ -233,14 +257,3 @@ def remove(self, name: str) -> ClusterConfig: |
233 | 257 | if not record: |
234 | 258 | raise ClusterNotFound.create(name) |
235 | 259 | return record |
236 | | - |
237 | | - async def cleanup(self, keep_clusters: Sequence[ClusterConfig]) -> None: |
238 | | - all_cluster_names = set(self._records.keys()) |
239 | | - keep_clusters_with_names = { |
240 | | - cluster_config.name for cluster_config in keep_clusters |
241 | | - } |
242 | | - for cluster_for_removal in all_cluster_names - keep_clusters_with_names: |
243 | | - try: |
244 | | - self.remove(cluster_for_removal) |
245 | | - except ClusterNotFound: |
246 | | - pass |
|
0 commit comments