init: influx + udp

Author: gistart

.gitignore

@@ -0,0 +1,17 @@
+*.swo
+*.swp
+*.pyc
+*.pyo
+*.log
+*.lock
+.env*
+.venv*
+.pypyenv*
+.pytest_*
+.mypy_cache
+*.egg-info
+__pycache__
+.vscode*
+
+!.gitkeep
+!/.gitignore

prometheus_push_client/__init__.py

@@ -0,0 +1,15 @@
+from .clients.bg import ThreadBGClient, AsyncioBGClient
+# TODO: clients.fg
+from .decorators import (
+    influx_udp_async,
+    graphite_udp_async,
+    influx_udp_thread,
+    graphite_udp_thread,
+)
+from .formats.influx import InfluxFormat
+from .formats.graphite import GraphiteFormat
+from .formats.prometheus import PrometheusFormat
+from .metrics import Counter, Gauge, Summary, Histogram, Info, Enum
+from .registry import PUSH_REGISTRY
+from .transports.udp import SyncUdpTransport, AioUdpTransport
+# TODO: transports.http

prometheus_push_client/clients/__init__.py

@@ -0,0 +1,84 @@
+import asyncio
+import threading
+import time
+
+
+class BackgroundClient:
+
+    def __init__(self, format, transport, period=15.0, *args, **kwargs):
+        self.format = format
+        self.transport = transport
+        self.period = float(period)
+        self.stop_event = None
+
+        self._period_step = 0.5  # check event every 0.5 seconds
+
+        super().__init__(*args, **kwargs)
+
+    def sleep_steps_iter(self, period):
+        n_full_steps, last_step = divmod(period, self._period_step)
+        for _ in range(int(n_full_steps)):
+            yield self._period_step
+        if last_step > 0:
+            yield last_step
+
+
+class ThreadBGClient(BackgroundClient, threading.Thread):
+    def start(self):
+        self.stop_event = threading.Event()
+        self.transport.start()
+        super().start()
+
+    def stop(self):
+        self.stop_event.set()
+        self.join()
+        self.transport.stop()
+
+    def run(self):
+        period = self.period
+
+        while not self.stop_event.is_set():
+            for step in self.sleep_steps_iter(period):
+                time.sleep(step)
+                if self.stop_event.is_set():
+                    break
+
+            ts_start = time.time()
+            samples_iter = self.format.iter_samples()
+            try:
+                self.transport.push_all(samples_iter)
+            except Exception:
+                pass  # TODO log?
+            period = self.period - (time.time() - ts_start)
+
+
+class AsyncioBGClient(BackgroundClient):
+    async def start(self):
+        self.stop_event = asyncio.Event()
+        await self.transport.start()
+        self._runner = asyncio.create_task(self.run())
+
+    async def stop(self):
+        self.stop_event.set()
+        await asyncio.wait([self._runner])  # TODO timeout, catch errors
+        self.transport.stop()
+
+    async def run(self):
+        period = self.period
+
+        while not self.stop_event.is_set():
+            for step in self.sleep_steps_iter(period):
+                try:
+                    await asyncio.sleep(step)
+                except asyncio.CancelledError:
+                    self.stop_event.set()
+                if self.stop_event.is_set():
+                    break
+
+            ts_start = time.time()
+            samples_iter = self.format.iter_samples()
+            try:
+                await self.transport.push_all(samples_iter)
+            except Exception:
+                pass  # TODO: log?
+            period = self.period - (time.time() - ts_start)

prometheus_push_client/clients/bg.py

@@ -0,0 +1,6 @@
+import asyncio
+
+if hasattr(asyncio, "get_running_loop"):
+    get_running_loop = asyncio.get_running_loop
+else:  # < 3.7, may be unsafe outside coroutines
+    get_running_loop = asyncio.get_event_loop

prometheus_push_client/compat.py

@@ -0,0 +1,65 @@
+import functools
+
+import prometheus_push_client as ppc
+
+
+def _thread_bg(client):
+    def _wrap_func(func):
+        @functools.wraps(func)
+        def _wrap_call(*args, **kwargs):
+            client.start()
+            try:
+                return func(*args, **kwargs)
+            finally:
+                client.stop()
+        return _wrap_call
+    return _wrap_func
+
+
+def _async_bg(client):
+    def _wrap_func(func):
+        @functools.wraps(func)
+        async def _wrap_call(*args, **kwargs):
+            await client.start()
+            try:
+                return await func(*args, **kwargs)
+            finally:
+                await client.stop()
+        return _wrap_call
+    return _wrap_func
+
+
+def influx_udp_async(host, port, period=15.0):
+    client = ppc.AsyncioBGClient(
+        format=ppc.InfluxFormat(),
+        transport=ppc.AioUdpTransport(host, port),
+        period=period,
+    )
+    return _async_bg(client)
+
+
+def influx_udp_thread(host, port, period=15.0):
+    client = ppc.ThreadBGClient(
+        format=ppc.InfluxFormat(),
+        transport=ppc.SyncUdpTransport(host, port),
+        period=period,
+    )
+    return _thread_bg(client)
+
+
+def graphite_udp_async(host, port, period=15.0):
+    client = ppc.AsyncioBGClient(
+        format=ppc.GraphiteFormat(),
+        transport=ppc.AioUdpTransport(host, port),
+        period=period,
+    )
+    return _async_bg(client)
+
+
+def graphite_udp_thread(host, port, period=15.0):
+    client = ppc.ThreadBGClient(
+        format=ppc.GraphiteFormat(),
+        transport=ppc.SyncUdpTransport(host, port),
+        period=period,
+    )
+    return _thread_bg(client)

prometheus_push_client/decorators.py

@@ -0,0 +1,23 @@
+import time
+
+from prometheus_push_client.registry import PUSH_REGISTRY
+
+
+class BaseFormat:
+    def __init__(self, registry=PUSH_REGISTRY):
+        self.registry = registry
+
+    def generate_latest(self):
+        samples = self.iter_samples()
+        return b"\n".join(samples)
+
+    def iter_samples(self):
+        for metric in self.registry.collect():
+            for sample in metric.samples:
+                line = self.format_sample(sample, metric)
+                if isinstance(line, str):
+                    line = line.encode("utf-8")
+                yield line
+
+    def format_sample(self, metric, sample):
+        raise NotImplementedError()

prometheus_push_client/formats/__init__.py

@@ -0,0 +1,43 @@
+from prometheus_push_client.formats.base import BaseFormat
+
+
+class GraphiteFormat(BaseFormat):
+
+    FMT_DATAPOINT = "{measurement}:{value}|{dtype}"  # no sampling?
+
+    DTYPES = {
+        "counter": "c",
+        "gauge": "g",
+        "info": "s",
+        "enum": "s", # TODO?
+        # histogram and summary are always counters?
+        "histogram": "c",
+        "summary": "c",
+    }
+
+    def format_sample(self, sample, metric):
+        measurement_name = sample.name
+        if sample.labels:
+            tags = " ".join(f"{k}={v}" for k,v in sample.labels.items())
+            measurement_name = f"{measurement_name},{tags}"
+
+        dtype = self.DTYPES.get(metric.type)
+        value = sample.value
+
+        # https://prometheus.io/docs/practices/naming/#base-units
+        if metric.type == "counter":  # TODO: other types
+            if metric.unit == "seconds":
+                dtype = "ms"
+                value *= 1000
+
+        if dtype is None:
+            raise NotImplementedError(f"cant convert type:{metric.type} unit:{metric.unit}")
+
+        if isinstance(value, float):
+            value = int(round(value))  # TODO
+
+        return self.FMT_DATAPOINT.format(
+            measurement=measurement_name,
+            value=value,
+            dtype=dtype,
+        )

prometheus_push_client/formats/base.py

@@ -0,0 +1,43 @@
+from prometheus_push_client.formats.base import BaseFormat
+
+
+class InfluxFormat(BaseFormat):
+
+    FMT_SAMPLE = "{sample_name}{tag_set} {measurement_name}={value}{timestamp}"
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.validate_metric_names(self.registry)
+
+    def validate_metric_names(self, registry):
+        names_not_supported = []
+        for metric_name in self.registry._names_to_collectors.keys():
+            if "_" not in metric_name:
+                names_not_supported.append(metric_name)
+
+        if names_not_supported:
+            raise ValueError(
+                "Following metrics can't be exported to Influx. Add `namespace`, "
+                "`subsystem`, `unit`, or just come up with 'two_words' name: "
+                f"{names_not_supported}"
+            )
+
+    def format_sample(self, sample, *_):
+        sample_name, measurement_name = sample.name.rsplit("_", 1)
+
+        tags = ""
+        if sample.labels:
+            tags = ",".join(f"{k}={v}" for k,v in sample.labels.items())
+            tags = f",{tags}"
+
+        ts = ""
+        if sample.timestamp:
+            ts = f" {int(sample.timestamp * 1e9)}"  # to nanoseconds
+
+        return self.FMT_SAMPLE.format(
+            sample_name=sample_name,
+            tag_set=tags,
+            measurement_name=measurement_name,
+            value=sample.value,
+            timestamp=ts,
+        )