temporal-flowrensic-python-sdk 0.6.63__py3-none-any.whl

temporal_flowrensic/_sink.py

@@ -0,0 +1,574 @@
+"""The record buffer every sender of the SDK uses: records wait in the worker's
+memory and leave in batches on a thread of their own, so recording never waits
+on Flowrensic."""
+
+from __future__ import annotations
+
+import os
+import socket
+import threading
+import time
+import uuid
+from collections import deque
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from typing import Any, Generic, Optional, TypeVar
+
+from connectrpc.code import Code
+from connectrpc.errors import ConnectError
+from google.protobuf.message import Message
+
+from temporal_flowrensic._gen.flowrensic.v1.activity_metrics_pb2 import (
+    ActivityAttempt,
+    IngestActivityEventsRequest,
+    WorkerBufferReport,
+)
+from temporal_flowrensic._gen.flowrensic.v1.services_connect import IngestServiceClientSync
+
+__all__ = [
+    "AttemptSink",
+    "BufferTarget",
+    "RecordSink",
+    "SendId",
+    "SinkCounts",
+    "approx_size",
+    "create_attempt_sink",
+    "create_record_buffer",
+]
+
+T = TypeVar("T")
+
+CURRENT_ENV_HEADER = "Flowrensic-Current-Env"
+ENV_ACCESS_KEY_HEADER = "Flowrensic-Env-Access-Key"
+CURRENT_EPHEMERAL_RUN_HEADER = "Flowrensic-Current-Ephemeral-Run"
+
+DEFAULT_BATCH_SIZE = 64
+DEFAULT_FLUSH_INTERVAL = 1.0
+DEFAULT_TIMEOUT_MS = 5000
+DEFAULT_MAX_WAITING = 10_000
+DEFAULT_MAX_BATCH_SIZE = 500
+DEFAULT_MAX_WAITING_BYTES = 16 * 1024 * 1024
+
+
+@dataclass(frozen=True)
+class SinkCounts:
+    """What a sink lost and holds, over every env."""
+
+    dropped: int
+    """Records dropped since the sink was made: past a limit, in a send that was refused, or left at close."""
+    timed_out: int
+    """Sends given up at their timeout since the sink was made."""
+    waiting: int
+    """Records waiting to be sent now."""
+    waiting_bytes: int
+    """The size of the records waiting now, in bytes, as approx_size gives it."""
+
+
+@dataclass(frozen=True)
+class SendId:
+    """Which send a call is: the buffer that sends it and the number of the
+    send. A send sent again keeps its number, so Flowrensic stores it once."""
+
+    sender: str
+    send_number: int
+
+
+class BufferTarget(Generic[T]):
+    """What a buffer needs from the sink that owns it."""
+
+    def send(self, records: list[T], timeout_ms: int, id: SendId) -> None:
+        """Sends one batch, and raises ConnectError when it fails."""
+        raise NotImplementedError
+
+    def env_of(self, record: T) -> str:
+        """The env a record is counted under."""
+        return ""
+
+    def report(self, reports: list[WorkerBufferReport], timeout_ms: int, id: SendId) -> None:
+        """Sends the counts. A target that does not send them raises
+        NotImplementedError, and the counts stay in the worker."""
+        raise NotImplementedError
+
+
+def approx_size(value: Any, depth: int = 0) -> int:
+    """The size of a record in the worker's memory, worked out from its fields:
+    two bytes per character of each text, the length of each byte array, eight
+    bytes for each other value, and 32 bytes for the record and each message or
+    list in it. Every field counts, a field left at its default too."""
+    if isinstance(value, str):
+        return len(value) * 2
+    if isinstance(value, (bytes, bytearray)):
+        return 32 + len(value)
+    if isinstance(value, Message):
+        if depth >= 8:
+            return 32
+        total = 32
+        for field in value.DESCRIPTOR.fields:
+            total += approx_size(getattr(value, field.name), depth + 1)
+        return total
+    if isinstance(value, Mapping):
+        if depth >= 8:
+            return 32
+        return 32 + sum(approx_size(one, depth + 1) for one in value.values())
+    if isinstance(value, Iterable):
+        if depth >= 8:
+            return 32
+        return 32 + sum(approx_size(one, depth + 1) for one in value)
+    return 8
+
+
+def _send_again(error: BaseException) -> bool:
+    """True when Flowrensic did not answer or answered try again later, so the batch is sent again."""
+    if not isinstance(error, ConnectError):
+        return False
+    return error.code in (Code.UNAVAILABLE, Code.DEADLINE_EXCEEDED, Code.RESOURCE_EXHAUSTED)
+
+
+def _add(counts: dict[str, int], env: str, n: int) -> None:
+    now = counts.get(env, 0) + n
+    if now == 0:
+        counts.pop(env, None)
+    else:
+        counts[env] = now
+
+
+@dataclass
+class _Entry(Generic[T]):
+    record: T
+    size: int
+
+
+@dataclass
+class _Kept(Generic[T]):
+    entries: list[_Entry[T]]
+    id: SendId
+    size: int
+
+
+@dataclass
+class _KeptReport:
+    reports: list[WorkerBufferReport]
+    id: SendId
+    dropped: dict[str, int]
+    timed_out: dict[str, int]
+    waiting: dict[str, int]
+
+
+class RecordSink(Generic[T]):
+    """Buffers records and sends them to a Flowrensic server in batches, on a
+    thread of its own. record never waits and never raises."""
+
+    def __init__(
+        self,
+        target: BufferTarget[T],
+        *,
+        batch_size: int = DEFAULT_BATCH_SIZE,
+        flush_interval: float = DEFAULT_FLUSH_INTERVAL,
+        timeout_ms: int = DEFAULT_TIMEOUT_MS,
+        max_waiting: int = DEFAULT_MAX_WAITING,
+        max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
+        max_waiting_bytes: int = DEFAULT_MAX_WAITING_BYTES,
+    ) -> None:
+        self._target = target
+        self._batch_size = batch_size
+        self._flush_interval = flush_interval
+        self._timeout_ms = timeout_ms
+        self._max_waiting = max_waiting
+        self._max_batch_size = max_batch_size
+        self._max_waiting_bytes = max_waiting_bytes
+
+        # Guards every field below. Held only for short work, never across a send.
+        self._lock = threading.Lock()
+        self._wake = threading.Condition(self._lock)
+        # One send or report at a time, by the thread or by flush.
+        self._sending = threading.Lock()
+
+        self._queue: deque[_Entry[T]] = deque()
+        self._queued_bytes = 0
+        self._kept: Optional[_Kept[T]] = None
+        self._kept_report: Optional[_KeptReport] = None
+        sender = str(uuid.uuid4())
+        self._records_sender = f"{sender}/records"
+        self._reports_sender = f"{sender}/reports"
+        self._send_number = 0
+        self._report_number = 0
+        self._waiting_by_env: dict[str, int] = {}
+        self._dropped: dict[str, int] = {}
+        self._timed_out: dict[str, int] = {}
+        self._reported_waiting: dict[str, int] = {}
+        self._dropped_total = 0
+        self._timed_out_total = 0
+        self._failing = False
+        self._woken = False
+        self._stopped = False
+
+        self._thread = threading.Thread(target=self._run, name="flowrensic-sink", daemon=True)
+        self._thread.start()
+
+    # What the worker calls.
+
+    def record(self, one: T) -> None:
+        """Puts one record in the buffer. Never waits on a send and never raises."""
+        try:
+            size = approx_size(one)
+            env = self._target.env_of(one)
+            with self._lock:
+                if self._stopped:
+                    _add(self._dropped, env, 1)
+                    self._dropped_total += 1
+                    return
+                self._queue.append(_Entry(one, size))
+                self._queued_bytes += size
+                _add(self._waiting_by_env, env, 1)
+                self._drop_oldest()
+                if self._waiting_count() >= self._batch_size and not self._failing:
+                    self._woken = True
+                    self._wake.notify()
+        except Exception:  # noqa: BLE001
+            return
+
+    def counts(self) -> SinkCounts:
+        with self._lock:
+            return SinkCounts(
+                dropped=self._dropped_total,
+                timed_out=self._timed_out_total,
+                waiting=self._waiting_count(),
+                waiting_bytes=self._waiting_bytes(),
+            )
+
+    @property
+    def batch_size(self) -> int:
+        """A send starts once this many records wait."""
+        return self._batch_size
+
+    @property
+    def flush_interval(self) -> float:
+        """Seconds between two ticks, each of which sends what waits."""
+        return self._flush_interval
+
+    @property
+    def timeout_ms(self) -> int:
+        """Milliseconds one send may take before it is given up."""
+        return self._timeout_ms
+
+    @property
+    def max_waiting(self) -> int:
+        """How many records may wait. Past it the oldest is dropped."""
+        return self._max_waiting
+
+    @property
+    def max_batch_size(self) -> int:
+        """How many records one send carries at most."""
+        return self._max_batch_size
+
+    @property
+    def max_waiting_bytes(self) -> int:
+        """How many bytes the records waiting may take, as approx_size gives it. Past it the oldest is dropped."""
+        return self._max_waiting_bytes
+
+    def flush_sync(self) -> None:
+        """Sends what waits and the counts, and returns once they landed or a
+        send went back. Blocks the calling thread."""
+        with self._sending:
+            tried = False
+            while True:
+                with self._lock:
+                    if self._waiting_count() == 0 or (self._failing and tried):
+                        break
+                    tried = True
+                    self._failing = False
+                self._send_while_answered()
+            self._report()
+
+    async def flush(self) -> None:
+        """Sends what waits and the counts, and waits for them to land, off the event loop."""
+        import asyncio
+
+        await asyncio.get_running_loop().run_in_executor(None, self.flush_sync)
+
+    def close_sync(self) -> None:
+        """Stops the thread, sends what waits once more, and drops what is left. Blocks the calling thread."""
+        with self._lock:
+            if self._stopped:
+                return
+            self._stopped = True
+            self._wake.notify()
+        self._thread.join()
+        self.flush_sync()
+        with self._sending:
+            with self._lock:
+                if self._kept is not None:
+                    for one in self._kept.entries:
+                        env = self._target.env_of(one.record)
+                        _add(self._waiting_by_env, env, -1)
+                        self._drop(env, 1)
+                    self._kept = None
+                for one in self._take(len(self._queue)):
+                    self._drop(self._target.env_of(one.record), 1)
+            self._report()
+
+    async def close(self) -> None:
+        """Sends what waits, stops, and drops what is left, off the event loop."""
+        import asyncio
+
+        await asyncio.get_running_loop().run_in_executor(None, self.close_sync)
+
+    # Inside, all with self._lock held unless said.
+
+    def _waiting_count(self) -> int:
+        return len(self._queue) + (len(self._kept.entries) if self._kept is not None else 0)
+
+    def _waiting_bytes(self) -> int:
+        return self._queued_bytes + (self._kept.size if self._kept is not None else 0)
+
+    def _drop(self, env: str, n: int) -> None:
+        _add(self._dropped, env, n)
+        self._dropped_total += n
+
+    def _take(self, count: int) -> list[_Entry[T]]:
+        out: list[_Entry[T]] = []
+        for _ in range(count):
+            if not self._queue:
+                break
+            one = self._queue.popleft()
+            _add(self._waiting_by_env, self._target.env_of(one.record), -1)
+            self._queued_bytes -= one.size
+            out.append(one)
+        return out
+
+    def _drop_oldest(self) -> None:
+        # Past either limit the oldest queued records go, never the kept send,
+        # which Flowrensic may already hold. The kept send does not count toward
+        # the count limit, and counts toward the memory limit.
+        excess = len(self._queue) - self._max_waiting
+        if excess > 0:
+            for one in self._take(excess):
+                self._drop(self._target.env_of(one.record), 1)
+        while self._queue and self._waiting_bytes() > self._max_waiting_bytes:
+            for one in self._take(1):
+                self._drop(self._target.env_of(one.record), 1)
+
+    def _next_send(self) -> Optional[_Kept[T]]:
+        if self._kept is not None:
+            out = self._kept
+            self._kept = None
+            for one in out.entries:
+                _add(self._waiting_by_env, self._target.env_of(one.record), -1)
+            return out
+        if not self._queue:
+            return None
+        entries = self._take(min(self._max_batch_size, len(self._queue)))
+        self._send_number += 1
+        return _Kept(entries, SendId(self._records_sender, self._send_number), sum(one.size for one in entries))
+
+    # Sending, with self._sending held and self._lock not held.
+
+    def _send_while_answered(self) -> None:
+        """Sends batch after batch while records wait, until a send goes back."""
+        while True:
+            with self._lock:
+                if self._failing:
+                    return
+                send = self._next_send()
+            if send is None:
+                return
+            records = [one.record for one in send.entries]
+            try:
+                self._target.send(records, self._timeout_ms, send.id)
+                with self._lock:
+                    self._failing = False
+            except BaseException as error:  # noqa: BLE001
+                envs = {self._target.env_of(one) for one in records}
+                with self._lock:
+                    if isinstance(error, ConnectError) and error.code == Code.DEADLINE_EXCEEDED:
+                        for env in envs:
+                            _add(self._timed_out, env, 1)
+                        self._timed_out_total += 1
+                    if _send_again(error):
+                        self._failing = True
+                        self._kept = send
+                        for one in send.entries:
+                            _add(self._waiting_by_env, self._target.env_of(one.record), 1)
+                        self._drop_oldest()
+                        return
+                    for one in records:
+                        self._drop(self._target.env_of(one), 1)
+
+    def _next_report(self) -> Optional[_KeptReport]:
+        envs = set(self._dropped) | set(self._timed_out) | set(self._waiting_by_env) | set(self._reported_waiting)
+        changed = any(
+            env in self._dropped
+            or env in self._timed_out
+            or self._waiting_by_env.get(env, 0) != self._reported_waiting.get(env, 0)
+            for env in envs
+        )
+        if not changed:
+            return None
+        dropped = dict(self._dropped)
+        timed_out = dict(self._timed_out)
+        waiting = dict(self._waiting_by_env)
+        reports = [
+            WorkerBufferReport(
+                env=env,
+                records_dropped=dropped.get(env, 0),
+                sends_timed_out=timed_out.get(env, 0),
+                records_waiting=waiting.get(env, 0),
+            )
+            for env in sorted(envs)
+        ]
+        self._report_number += 1
+        return _KeptReport(reports, SendId(self._reports_sender, self._report_number), dropped, timed_out, waiting)
+
+    def _report(self) -> None:
+        """At most one report at a time. A report that was not answered is sent
+        again as it was. Otherwise one goes only when a count changed since the
+        last report that was answered, and carries what was counted since."""
+        with self._lock:
+            if self._kept_report is None:
+                self._kept_report = self._next_report()
+            one = self._kept_report
+        if one is None:
+            return
+        try:
+            self._target.report(one.reports, self._timeout_ms, one.id)
+        except NotImplementedError:
+            with self._lock:
+                self._kept_report = None
+            return
+        except BaseException:  # noqa: BLE001
+            return
+        with self._lock:
+            for env, n in one.dropped.items():
+                _add(self._dropped, env, -n)
+            for env, n in one.timed_out.items():
+                _add(self._timed_out, env, -n)
+            self._reported_waiting = one.waiting
+            self._kept_report = None
+
+    def _run(self) -> None:
+        """The thread: a tick every flush interval sends and reports, and a full
+        batch sends at once while sends are answered."""
+        next_tick = time.monotonic() + self._flush_interval
+        while True:
+            with self._lock:
+                while not self._stopped and not self._woken and time.monotonic() < next_tick:
+                    self._wake.wait(max(0.0, next_tick - time.monotonic()))
+                if self._stopped:
+                    return
+                woken = self._woken
+                self._woken = False
+                tick = time.monotonic() >= next_tick
+                if tick:
+                    next_tick = time.monotonic() + self._flush_interval
+                    self._failing = False
+            try:
+                with self._sending:
+                    if tick or woken:
+                        self._send_while_answered()
+                    if tick:
+                        self._report()
+            except BaseException:  # noqa: BLE001
+                continue
+
+
+def create_record_buffer(target: BufferTarget[T], **options: Any) -> RecordSink[T]:
+    """A record buffer that sends through target, with the buffer options given."""
+    return RecordSink(target, **options)
+
+
+AttemptSink = RecordSink[ActivityAttempt]
+"""Buffers activity attempts and sends them to a Flowrensic server in batches."""
+
+
+def default_worker_id() -> str:
+    return f"{socket.gethostname() or 'unknown'}-{os.getpid()}"
+
+
+def client_headers(current_env: str, env_access_key: str, current_ephemeral_run: str) -> dict[str, str]:
+    """The headers every client call carries: the env, its access key and the ephemeral run. An empty value is not sent."""
+    headers: dict[str, str] = {}
+    if current_env:
+        headers[CURRENT_ENV_HEADER] = current_env
+    if env_access_key:
+        headers[ENV_ACCESS_KEY_HEADER] = env_access_key
+    if current_ephemeral_run:
+        headers[CURRENT_EPHEMERAL_RUN_HEADER] = current_ephemeral_run
+    return headers
+
+
+def ingest_client_sync(server_base_url: str) -> IngestServiceClientSync:
+    return IngestServiceClientSync(server_base_url.rstrip("/"), send_compression=None)
+
+
+class _AttemptTarget(BufferTarget[ActivityAttempt]):
+    def __init__(self, server_base_url: str, worker_id: str, headers: dict[str, str], worker_name: str = "") -> None:
+        self._ingest = ingest_client_sync(server_base_url)
+        self._worker_id = worker_id
+        self._worker_name = worker_name
+        self._headers = headers
+
+    def send(self, records: list[ActivityAttempt], timeout_ms: int, id: SendId) -> None:
+        self._ingest.ingest_activity_events(
+            IngestActivityEventsRequest(
+                events=records,
+                worker_id=self._worker_id,
+                worker_name=self._worker_name,
+                sender=id.sender,
+                send_number=id.send_number,
+            ),
+            headers=self._headers,
+            timeout_ms=max(1, timeout_ms),
+        )
+
+    def env_of(self, record: ActivityAttempt) -> str:
+        return record.env
+
+    def report(self, reports: list[WorkerBufferReport], timeout_ms: int, id: SendId) -> None:
+        self._ingest.ingest_activity_events(
+            IngestActivityEventsRequest(
+                reports=reports, worker_id=self._worker_id, sender=id.sender, send_number=id.send_number
+            ),
+            headers=self._headers,
+            timeout_ms=max(1, timeout_ms),
+        )
+
+
+def create_attempt_sink(
+    server_base_url: str,
+    *,
+    worker_id: str = "",
+    worker_name: str = "",
+    current_env: str = "",
+    env_access_key: str = "",
+    current_ephemeral_run: str = "",
+    batch_size: int = DEFAULT_BATCH_SIZE,
+    flush_interval: float = DEFAULT_FLUSH_INTERVAL,
+    timeout_ms: int = DEFAULT_TIMEOUT_MS,
+    max_waiting: int = DEFAULT_MAX_WAITING,
+    max_batch_size: int = DEFAULT_MAX_BATCH_SIZE,
+    max_waiting_bytes: int = DEFAULT_MAX_WAITING_BYTES,
+) -> AttemptSink:
+    """A sink of activity attempts that sends them to the Flowrensic server at
+    server_base_url with IngestActivityEvents. Every call carries current_env,
+    env_access_key and current_ephemeral_run in their headers. worker_id
+    names this worker; empty, the host name and the process id joined by a
+    dash. worker_name is the name of the worker, the same for every copy of
+    it, sent with each attempt; empty, the worker has no name. Times are in
+    seconds."""
+    if not server_base_url:
+        raise ValueError("flowrensic: server_base_url is required")
+    return RecordSink(
+        _AttemptTarget(
+            server_base_url,
+            worker_id or default_worker_id(),
+            client_headers(current_env, env_access_key, current_ephemeral_run),
+            worker_name,
+        ),
+        batch_size=batch_size,
+        flush_interval=flush_interval,
+        timeout_ms=timeout_ms,
+        max_waiting=max_waiting,
+        max_batch_size=max_batch_size,
+        max_waiting_bytes=max_waiting_bytes,
+    )
+