Merge branch 'batched-telemetry' into 'enterprise'

MCP & UI: async batched Mixpanel telemetry (TG-1076)

See merge request dkinternal/testgen/dataops-testgen!571

Author: ci bot

testgen/common/mixpanel_service.py

@@ -1,7 +1,11 @@
+import atexit
 import functools
 import json
 import logging
+import queue
 import ssl
+import threading
+import time
 import uuid
 from base64 import b64encode
 from functools import cached_property, wraps
@@ -18,6 +22,13 @@
 
 LOG = logging.getLogger("testgen")
 
+_BATCH_SIZE = 50           # Mixpanel /track array cap
+_FLUSH_INTERVAL_SEC = 10   # Time-based flush
+_QUEUE_MAX_SIZE = 1000     # Memory cap before drop-on-overflow
+_DRAIN_TIMEOUT_SEC = 5     # Bounded shutdown drain
+
+_SHUTDOWN = object()       # sentinel enqueued by drain() to stop the worker
+
 
 def safe_method(method):
     @wraps(method)
@@ -33,6 +44,14 @@ def wrapped(*args, **kwargs):
 
 class MixpanelService(Singleton):
 
+    def __init__(self) -> None:
+        self._queue: queue.Queue | None = None
+        self._worker: threading.Thread | None = None
+        self._worker_lock = threading.Lock()
+        self._started = False
+        self._stopped = False
+        atexit.register(self.drain)
+
     @cached_property
     @with_database_session
     def instance_id(self):
@@ -54,33 +73,106 @@ def _hash_value(self, value: bytes | str, digest_size: int = 8) -> str:
 
     @safe_method
     def send_event(self, event_name, include_usage=False, **properties):
-        self._track(event_name, include_usage=include_usage, **properties)
+        self._enqueue(self._build_event(event_name, include_usage=include_usage, **properties))
 
     def send_feedback(self, **properties):
-        # User-submitted feedback is content the user explicitly chose to share
-        # so it is not gated by the TG_ANALYTICS opt-out.
+        # User-submitted feedback is content the user explicitly chose to share,
+        # so it is not gated by the TG_ANALYTICS opt-out. It is a foreground action
+        # posted synchronously — never enqueued — so it never starts the worker.
         try:
-            self._track("feedback", **properties)
+            self.send_mp_request("track?ip=1", self._build_event("feedback", **properties))
         except Exception:
             LOG.exception("Error sending feedback")
 
-    def _track(self, event_name, include_usage=False, **properties):
+    def _build_event(self, event_name, include_usage=False, **properties) -> dict:
         properties.setdefault("instance_id", self.instance_id)
         properties.setdefault("edition", settings.DOCKER_HUB_REPOSITORY)
         properties.setdefault("version", settings.VERSION)
         properties.setdefault("username", session.auth.user_display if session.auth else None)
         properties.setdefault("distinct_id", self.get_distinct_id(properties["username"]))
+        properties.setdefault("time", int(time.time()))
         if include_usage:
             properties.update(self.get_usage())
 
-        track_payload = {
+        return {
             "event": event_name,
             "properties": {
                 "token": settings.MIXPANEL_TOKEN,
                 **properties,
-            }
+            },
         }
-        self.send_mp_request("track?ip=1", track_payload)
+
+    def _ensure_worker(self) -> None:
+        if self._started:
+            return
+        with self._worker_lock:
+            if self._started:
+                return
+            self._queue = queue.Queue(maxsize=_QUEUE_MAX_SIZE)
+            self._worker = threading.Thread(target=self._worker_loop, name="mixpanel-flush", daemon=True)
+            self._worker.start()
+            self._started = True
+
+    def _enqueue(self, event: dict) -> None:
+        if self._stopped:
+            LOG.warning("analytics worker stopped; dropping event")
+            return
+        self._ensure_worker()
+        try:
+            self._queue.put_nowait(event)
+        except queue.Full:
+            LOG.warning("analytics queue full; dropping event")
+
+    def _worker_loop(self) -> None:
+        while True:
+            batch, stop = self._next_batch()
+            if batch:
+                self._flush(batch)
+            if stop:
+                return
+
+    def _next_batch(self) -> tuple[list[dict], bool]:
+        """Block up to _FLUSH_INTERVAL_SEC for the first event, then drain up to
+        _BATCH_SIZE without blocking. Returns (batch, stop)."""
+        try:
+            first = self._queue.get(timeout=_FLUSH_INTERVAL_SEC)
+        except queue.Empty:
+            return [], False
+        if first is _SHUTDOWN:
+            return [], True
+        batch = [first]
+        while len(batch) < _BATCH_SIZE:
+            try:
+                event = self._queue.get_nowait()
+            except queue.Empty:
+                break
+            if event is _SHUTDOWN:
+                return batch, True
+            batch.append(event)
+        return batch, False
+
+    def _flush(self, events: list[dict]) -> None:
+        for start in range(0, len(events), _BATCH_SIZE):
+            chunk = events[start:start + _BATCH_SIZE]
+            try:
+                self.send_mp_request("track?ip=1", chunk)
+            except Exception:
+                LOG.exception("Failed to flush analytics batch")
+
+    def drain(self) -> None:
+        """Flush queued events and stop the worker. Idempotent; bounded by
+        _DRAIN_TIMEOUT_SEC. Called by the atexit hook and the server lifespan."""
+        if not self._started or self._stopped:
+            return
+        self._stopped = True
+        try:
+            self._queue.put_nowait(_SHUTDOWN)
+        except queue.Full:
+            LOG.warning("analytics queue full at shutdown; in-flight events may be dropped")
+        if self._worker is not None:
+            self._worker.join(timeout=_DRAIN_TIMEOUT_SEC)
+            if self._worker.is_alive():
+                LOG.warning("analytics drain timed out after %ss", _DRAIN_TIMEOUT_SEC)
 
     def get_ssl_context(self):
         ssl_context = ssl.create_default_context()

testgen/mcp/permissions.py

@@ -80,6 +80,11 @@ def set_mcp_token(token: str | None) -> None:
     _mcp_token.set(token)
 
 
+def get_mcp_username() -> str | None:
+    """Return the authenticated username for the current MCP request (or None)."""
+    return _mcp_username.get()
+
+
 def get_authorized_mcp_user() -> User:
     """Get the authenticated and authorized User for the current MCP request.