taskflow-meter 1.0.0__py3-none-any.whl

taskflow_meter/transports/amqp.py

@@ -0,0 +1,204 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Carry events over a message broker, using kombu.
+
+kombu rather than a raw AMQP client because taskflow already depends on
+it for its worker-based engine, so a deployment running those flows has
+it already.
+
+This is the transport for the collector deployment: the processes
+running flows publish, one collector consumes and writes to a shared
+datasource, and the API workers read that with ``poll = false``.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+from collections.abc import Sequence
+from typing import Any
+
+import kombu
+
+from taskflow_meter.events import Event
+from taskflow_meter.transports.base import Publisher
+from taskflow_meter.transports.base import Subscriber
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_EXCHANGE = "taskflow-meter"
+DEFAULT_QUEUE = "taskflow-meter-events"
+DEFAULT_ROUTING_KEY = "events"
+DEFAULT_TIMEOUT = 1.0
+
+
+def build_exchange(name: str = DEFAULT_EXCHANGE) -> kombu.Exchange:
+    return kombu.Exchange(name, type="direct", durable=True)
+
+
+def build_queue(
+    name: str = DEFAULT_QUEUE,
+    *,
+    exchange: str = DEFAULT_EXCHANGE,
+    routing_key: str = DEFAULT_ROUTING_KEY,
+) -> kombu.Queue:
+    """Durable, so events survive a collector restart.
+
+    A monitoring queue that discards what it could not deliver is a
+    monitoring queue that lies about the gap.
+    """
+    return kombu.Queue(
+        name,
+        exchange=build_exchange(exchange),
+        routing_key=routing_key,
+        durable=True,
+    )
+
+
+class AMQPTransport(Publisher):
+    """Publishes batches to an exchange."""
+
+    name = "amqp"
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        exchange: str = DEFAULT_EXCHANGE,
+        queue: str = DEFAULT_QUEUE,
+        routing_key: str = DEFAULT_ROUTING_KEY,
+        connection: kombu.Connection | None = None,
+    ) -> None:
+        self.url = url
+        self.routing_key = routing_key
+        self._exchange = build_exchange(exchange)
+        self._queue = build_queue(
+            queue, exchange=exchange, routing_key=routing_key
+        )
+        self._connection = connection
+        self._owns_connection = connection is None
+
+    @property
+    def connection(self) -> kombu.Connection:
+        if self._connection is None:
+            self._connection = kombu.Connection(self.url)
+        return self._connection
+
+    def stop(self) -> None:
+        if self._owns_connection and self._connection is not None:
+            self._connection.release()
+            self._connection = None
+
+    def publish(self, events: Sequence[Event]) -> None:
+        payload = {"events": [event.to_dict() for event in events]}
+        producer = self.connection.Producer(serializer="json")
+        producer.publish(
+            payload,
+            exchange=self._exchange,
+            routing_key=self.routing_key,
+            # Declared on every publish: the exchange and queue may not
+            # exist yet if a flow starts before the collector does, and
+            # events published into nothing are simply lost.
+            declare=[self._queue],
+            retry=True,
+        )
+
+
+class AMQPSubscriber(Subscriber):
+    """Consumes batches from a queue."""
+
+    name = "amqp"
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        exchange: str = DEFAULT_EXCHANGE,
+        queue: str = DEFAULT_QUEUE,
+        routing_key: str = DEFAULT_ROUTING_KEY,
+        connection: kombu.Connection | None = None,
+    ) -> None:
+        self.url = url
+        self._queue = build_queue(
+            queue, exchange=exchange, routing_key=routing_key
+        )
+        self._connection = connection
+        self._owns_connection = connection is None
+
+    @property
+    def connection(self) -> kombu.Connection:
+        if self._connection is None:
+            self._connection = kombu.Connection(self.url)
+        return self._connection
+
+    def stop(self) -> None:
+        if self._owns_connection and self._connection is not None:
+            self._connection.release()
+            self._connection = None
+
+    def consume(
+        self,
+        handler: Callable[[Sequence[Event]], None],
+        *,
+        timeout: float | None = DEFAULT_TIMEOUT,
+    ) -> int:
+        """Deliver what has arrived, and return how many events."""
+        received = 0
+        give_up = False
+
+        def on_message(body: Any, message: Any) -> None:
+            nonlocal received, give_up
+            try:
+                events = _decode(body)
+            except Exception:
+                # A message we cannot read will never become readable.
+                # Rejecting it beats redelivering it forever.
+                LOG.exception("discarding an unreadable message")
+                message.reject()
+                return
+            try:
+                handler(events)
+            except Exception:
+                # Leave it on the queue: the handler may recover, and
+                # dropping monitoring data silently is the one outcome
+                # worse than redelivering it.  Then stop draining --
+                # a requeued message is redelivered immediately, so
+                # carrying on would spin on it until the handler
+                # recovered, which it cannot while we never return.
+                LOG.exception("handler failed; leaving the message")
+                message.requeue()
+                give_up = True
+                return
+            received += len(events)
+            message.ack()
+
+        with self.connection.Consumer(
+            [self._queue], callbacks=[on_message], accept=["json"]
+        ):
+            while not give_up:
+                try:
+                    self.connection.drain_events(timeout=timeout)
+                except TimeoutError:
+                    break
+                except OSError:  # pragma: no cover - broker went away
+                    LOG.exception("lost the broker connection")
+                    break
+        return received
+
+
+def _decode(body: Any) -> list[Event]:
+    if isinstance(body, str):
+        import json
+
+        body = json.loads(body)
+    return [Event.from_dict(item) for item in body["events"]]

taskflow_meter/transports/base.py

@@ -0,0 +1,105 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Where events go once they leave the process that produced them.
+
+A publisher is handed batches, never single events: the pipeline drains
+its queue in one go, so a transport that can amortise a round trip gets
+the chance to.
+"""
+
+from __future__ import annotations
+
+import abc
+from collections.abc import Callable
+from collections.abc import Sequence
+from types import TracebackType
+from typing import Self
+
+from taskflow_meter.events import Event
+
+
+class Publisher(abc.ABC):
+    """Sends events somewhere.
+
+    Implementations may raise: the pipeline that drives them counts and
+    logs failures rather than letting one reach the flow being watched.
+    """
+
+    #: Stevedore plugin name, set by subclasses.
+    name: str = ""
+
+    def start(self) -> None:  # noqa: B027 - optional hook, not abstract
+        """Acquire whatever the transport needs.  Idempotent."""
+
+    def stop(self) -> None:  # noqa: B027 - optional hook, not abstract
+        """Release it again.  Idempotent, and safe to call unstarted."""
+
+    @abc.abstractmethod
+    def publish(self, events: Sequence[Event]) -> None:
+        """Send a batch.  Never called with an empty one."""
+
+    def __enter__(self) -> Self:
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.stop()
+
+
+class Subscriber(abc.ABC):
+    """Receives events somebody else published.
+
+    The collector's end of a transport: one process consumes what the
+    flows emitted and writes it to a datasource that any number of API
+    workers read.
+    """
+
+    #: Stevedore plugin name, set by subclasses.
+    name: str = ""
+
+    def start(self) -> None:  # noqa: B027 - optional hook, not abstract
+        """Acquire whatever the transport needs.  Idempotent."""
+
+    def stop(self) -> None:  # noqa: B027 - optional hook, not abstract
+        """Release it again.  Idempotent, and safe to call unstarted."""
+
+    @abc.abstractmethod
+    def consume(
+        self,
+        handler: Callable[[Sequence[Event]], None],
+        *,
+        timeout: float | None = None,
+    ) -> int:
+        """Deliver whatever has arrived, and return how many events.
+
+        Returns when nothing more arrives within ``timeout``, so the
+        caller owns the loop -- and can shut it down between batches
+        rather than being trapped inside one.
+        """
+
+    def __enter__(self) -> Self:
+        self.start()
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        self.stop()

taskflow_meter/transports/http.py

@@ -0,0 +1,97 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""Post batches of events to a webhook.
+
+Uses :mod:`urllib.request` rather than a HTTP client library: this is
+one POST of a JSON array, and a monitoring sidecar is a poor reason to
+put another dependency into somebody's service.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+import time
+import urllib.error
+import urllib.request
+from collections.abc import Sequence
+
+from taskflow_meter.events import Event
+from taskflow_meter.transports.base import Publisher
+
+LOG = logging.getLogger(__name__)
+
+DEFAULT_TIMEOUT = 5.0
+DEFAULT_RETRIES = 2
+DEFAULT_BACKOFF = 0.2
+
+#: Statuses worth trying again.  A 4xx means the request was wrong and
+#: will be just as wrong the second time.
+RETRYABLE_STATUSES = frozenset({408, 425, 429, 500, 502, 503, 504})
+
+
+class HTTPTransport(Publisher):
+    """POSTs ``{"events": [...]}`` to a URL."""
+
+    name = "http"
+
+    def __init__(
+        self,
+        url: str,
+        *,
+        timeout: float = DEFAULT_TIMEOUT,
+        retries: int = DEFAULT_RETRIES,
+        backoff: float = DEFAULT_BACKOFF,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        if retries < 0:
+            msg = "retries cannot be negative"
+            raise ValueError(msg)
+        self.url = url
+        self.timeout = timeout
+        self.retries = retries
+        self.backoff = backoff
+        self.headers = {"content-type": "application/json", **(headers or {})}
+
+    def publish(self, events: Sequence[Event]) -> None:
+        body = json.dumps(
+            {"events": [event.to_dict() for event in events]},
+            separators=(",", ":"),
+        ).encode()
+
+        last: Exception | None = None
+        for attempt in range(self.retries + 1):
+            try:
+                self._post(body)
+            except urllib.error.HTTPError as exc:
+                last = exc
+                if exc.code not in RETRYABLE_STATUSES:
+                    # Retrying a rejected request just rejects it again.
+                    raise
+            except OSError as exc:
+                last = exc
+            else:
+                return
+
+            if attempt < self.retries:
+                time.sleep(self.backoff * (2**attempt))
+
+        assert last is not None
+        raise last
+
+    def _post(self, body: bytes) -> None:
+        request = urllib.request.Request(
+            self.url, data=body, headers=self.headers, method="POST"
+        )
+        with urllib.request.urlopen(request, timeout=self.timeout) as response:
+            LOG.debug("published to %s: %s", self.url, response.status)

taskflow_meter/transports/memory.py

@@ -0,0 +1,83 @@
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+"""In-process transports: keep the events, or fold them into a store."""
+
+from __future__ import annotations
+
+import threading
+from collections import deque
+from collections.abc import Sequence
+
+from taskflow_meter.datasource.base import WritableDataSource
+from taskflow_meter.events import Event
+from taskflow_meter.transports.base import Publisher
+
+#: Events retained before the oldest are dropped.
+DEFAULT_MAX_EVENTS = 10000
+
+
+class MemoryTransport(Publisher):
+    """Holds events in a bounded buffer for something else to drain.
+
+    Useful for tests and for handing a batch to another thread in the
+    same process.  Bounded because an undrained buffer is a leak with a
+    slower fuse than a crash.
+    """
+
+    name = "memory"
+
+    def __init__(self, *, max_events: int = DEFAULT_MAX_EVENTS) -> None:
+        if max_events < 1:
+            msg = "max_events must be at least 1"
+            raise ValueError(msg)
+        self._events: deque[Event] = deque(maxlen=max_events)
+        self._lock = threading.Lock()
+        self.dropped = 0
+
+    def publish(self, events: Sequence[Event]) -> None:
+        with self._lock:
+            room = self._events.maxlen or 0
+            overflow = max(0, len(self._events) + len(events) - room)
+            self.dropped += overflow
+            self._events.extend(events)
+
+    def drain(self) -> tuple[Event, ...]:
+        """Take everything buffered so far."""
+        with self._lock:
+            taken = tuple(self._events)
+            self._events.clear()
+            return taken
+
+    def peek(self) -> tuple[Event, ...]:
+        with self._lock:
+            return tuple(self._events)
+
+    def __len__(self) -> int:
+        with self._lock:
+            return len(self._events)
+
+
+class DataSourcePublisher(Publisher):
+    """Folds events straight into a writable datasource.
+
+    The whole in-process path: a listener produces events and the API
+    can serve them from the same process, with nothing on the wire.
+    """
+
+    name = "datasource"
+
+    def __init__(self, sink: WritableDataSource) -> None:
+        self.sink = sink
+
+    def publish(self, events: Sequence[Event]) -> None:
+        self.sink.apply_many(events)