Re-use existing scrape sessions after scrape timeout

Dockerfile

@@ -1,6 +1,6 @@
 ARG KATSDPDOCKERBASE_REGISTRY=quay.io/ska-sa
 
-FROM $KATSDPDOCKERBASE_REGISTRY/docker-base-build as build
+FROM $KATSDPDOCKERBASE_REGISTRY/docker-base-build AS build
 
 # Enable Python 3 venv
 ENV PATH="$PATH_PYTHON3" VIRTUAL_ENV="$VIRTUAL_ENV_PYTHON3"

src/switch_exporter/cache.py

@@ -108,12 +108,12 @@ def __init__(self, item_cls: Callable[['Cache[_T]', _T], Item[_T]],
         self.item_cls = item_cls
         self.timeout = timeout
 
-    def get(self, key: _T) -> Item[_T]:
+    def get(self, key: _T, *init_params) -> Item[_T]:
         """Obtain an item from the cache, creating it if necessary."""
         try:
             item = self._items[key]
         except KeyError:
-            item = self.item_cls(self, key)
+            item = self.item_cls(self, key, *init_params)
             logging.info('Created %r', item)
             self._items[key] = item
         return item

src/switch_exporter/scraper.py

@@ -0,0 +1,145 @@
+
+import asyncio
+from collections.abc import Coroutine
+from typing import Optional, List
+import time
+from typing_extensions import override
+import prometheus_client
+import logging
+
+from .cache import Cache, Item
+from .switch import Switch
+logger = logging.getLogger(__name__)
+
+
+class ValidationError(Exception):
+    pass
+
+
+class Scraper(Item):
+    def __init__(
+        self,
+        cache: Cache,
+        key: str,
+        switch: Switch,
+        enable_timing_metrics: bool = True,
+    ) -> None:
+        self.key = key
+        super().__init__(cache, key)
+        self.enable_timing_metrics = enable_timing_metrics
+        self.switch = switch
+        self._lock = asyncio.Lock()
+        # TODO: Use a TaskGroup instead of a list of tasks to robustly handle the async context.
+        self.tasks = []
+        self.done = asyncio.Event()  # Set to True when the scraper is done scraping.
+        self.done.set()  # Initially set to True to indicate that the scraper should be started.
+        self.registry = prometheus_client.CollectorRegistry()
+        self._error = None
+        # if we timout too many times, we should raise an error and reset scraper
+        self.timeout_counter = 0
+
+    async def timed(
+        self,
+        coroutine: Coroutine,
+        timing_gauge: prometheus_client.Gauge,
+        hostname: str,
+    ) -> None:
+        start_time = time.perf_counter()
+        await coroutine
+        end_time = time.perf_counter()
+        duration = end_time - start_time
+        if self.enable_timing_metrics:
+            timing_gauge.labels(hostname, coroutine.__name__).set(duration)
+
+    async def wait_for_scraper(self) -> None:
+        """Wait until collector tasks finish and update the registry.
+        Sets the self.done event when the scraper is done.
+
+        Must not raise: this runs as a background task so that a timed-out
+        caller does not prevent ``done`` from being set.
+        """
+        self._error = None
+        try:
+            done, _ = await asyncio.wait(self.tasks)
+            exceptions = []
+            for task in done:
+                try:
+                    task.result()
+                except Exception as e:
+                    logger.error('Error during scraping metrics: %s', task.get_name())
+                    exceptions.append(e)
+            if exceptions:
+                self._error = Exception(
+                    "Error during scraping metrics: " + ', '.join([str(e) for e in exceptions])
+                )
+        except Exception as e:
+            self._error = e
+        finally:
+            self.done.set()
+
+    async def await_scraper_done(self, timeout: float) -> prometheus_client.CollectorRegistry:
+        if self._error is not None:
+            raise self._error
+        try:
+            await asyncio.wait_for(self.done.wait(), timeout=timeout)
+        except asyncio.TimeoutError:
+            self.timeout_counter += 1
+            if self.timeout_counter > 10:
+                raise RuntimeError(f'Timeout handling {self.key} metrics too many times')
+            raise asyncio.TimeoutError(f'Timeout handling {self.key} metrics')
+        except Exception as e:
+            raise e
+
+        self.timeout_counter = 0
+        return self.registry
+
+    async def scrape(
+        self,
+        timeout: float,
+        collectors: Optional[List[str]],
+    ) -> prometheus_client.CollectorRegistry:
+        """Obtain the metrics from the switch"""
+        start_time = time.perf_counter()
+
+        await self.switch.refresh_port_info()
+        temp_registry = prometheus_client.CollectorRegistry()
+        timing_gauge = prometheus_client.Gauge(
+            'switch_coroutine_duration_seconds', 'duration of the coroutine',
+            labelnames=('hostname', 'coroutine'),
+            registry=temp_registry,
+        )
+        if collectors is None:
+            scraper_fns = list(self.switch.collectors.values())
+        else:
+            scraper_fns = []
+            for collector in collectors:
+                try:
+                    scraper_fns.append(self.switch.collectors[collector])
+                except KeyError as e:
+                    raise ValidationError(f'Unknown collector: {collector}') from e
+
+        async with self._lock:
+            scrape_timeout = timeout - (time.perf_counter() - start_time)
+            new_scrape = self.done.is_set()
+            if new_scrape:
+                scrapers = [fn(temp_registry) for fn in scraper_fns]
+                self.tasks = [
+                    asyncio.create_task(
+                        self.timed(s, timing_gauge, self.switch.hostname),
+                        name=s.__name__ + f'({self.key})'
+                    )
+                    for s in scrapers
+                ]
+                self.registry = temp_registry
+                self.done.clear()
+
+        if new_scrape:
+            asyncio.create_task(self.wait_for_scraper(), name=f'wait_for_scraper({self.key})')
+
+        return await self.await_scraper_done(scrape_timeout)
+
+    @override
+    async def close(self) -> None:
+        await self.switch.close()
+        self.done.set()
+        self.timeout_counter = 0

src/switch_exporter/server.py

@@ -3,12 +3,14 @@
 import functools
 import argparse
 import logging
+from typing import Callable
 
 import katsdpservices
 from aiohttp import web
 import prometheus_client
 
-from .switch import Switch, ValidationError
+from .scraper import Scraper, ValidationError
+from .switch import Switch
 from .cache import Cache
 
 
@@ -25,16 +27,16 @@ async def get_metrics(request: web.Request) -> web.Response:
 
     collect = request.query.getall('collect', None)
     cache = request.app['cache']
-    switch = cache.get(target)
+    scraper = cache.get(target + ':' + ','.join(collect) if collect else target, target)
     timeout = request.app['scrape_timeout']
     try:
         timeout = int(request.query.get('scrape_timeout', timeout))
     except ValueError:
         logger.exception('Invalid scrape_timeout value')
         raise web.HTTPBadRequest(text='scrape_timeout must be an integer') from None
     try:
-        with switch:
-            counters = await switch.scrape(timeout, collect)
+        with scraper:
+            counters = await scraper.scrape(timeout, collect)
     except asyncio.CancelledError:
         raise
     except asyncio.TimeoutError:
@@ -47,24 +49,31 @@ async def get_metrics(request: web.Request) -> web.Response:
         raise web.HTTPBadRequest(text=str(e)) from None
     except Exception as exc:
         # Possibly a failed connection, so reset it
-        logger.exception('Exception during scrape, resetting switch')
-        switch.destroy()
+        logger.exception('Exception during scrape, resetting scraper')
+        scraper.destroy()
         raise web.HTTPInternalServerError(text='Scrape failed: ' + str(exc)) from None
     else:
         content = prometheus_client.generate_latest(counters).decode()
         return web.Response(text=content)
 
 
+def scraper_factory(switch_factory: Callable) -> Callable[[Cache, str], Scraper]:
+    def scraper(cache: Cache, key: str, target: str) -> Scraper:
+        switch = switch_factory(target)
+        return Scraper(cache, key, switch)
+    return scraper
+
+
 async def make_app(args: argparse.Namespace, loop: asyncio.AbstractEventLoop) -> web.Application:
     app = web.Application(loop=loop)
-    factory = functools.partial(
+    switch_factory = functools.partial(
         Switch,
         username=args.username,
         password=args.password,
         keyfile=args.keyfile,
         lldp_timeout=args.lldp_timeout,
     )
-    app['cache'] = Cache(factory, args.connection_timeout)
+    app['cache'] = Cache(scraper_factory(switch_factory), args.connection_timeout)
     app['scrape_timeout'] = args.scrape_timeout
     app.router.add_get('/metrics', get_metrics)
     return app