Добавлен сбор телеметрии на ВМ2

VM2_services/codebase/services/message-safety/app/worker.py

@@ -5,6 +5,8 @@ import socket
 import uuid
 from datetime import UTC, datetime, timedelta
 
+from opentelemetry import trace
+
 from app.adapters import S3VersionReader
 from app.config import validate_config
 from app.contracts import Attachment
@@ -19,6 +21,15 @@ from app.file_pipeline import (
 )
 from app.repository import Repository
 from app.settings import BootstrapSettings
+from app.telemetry import (
+    init_telemetry,
+    record_dependency,
+    record_worker,
+    shutdown_telemetry,
+    task_link,
+)
+
+tracer = trace.get_tracer("message-safety.worker")
 
 
 class Worker:
@@ -34,14 +45,30 @@ class Worker:
         self.owner = f"{socket.gethostname()}:{uuid.uuid4()}"
 
     async def once(self) -> bool:
-        task = await self.repository.claim(self.owner)
+        with tracer.start_as_current_span("message_safety.worker.claim"):
+            task = await self.repository.claim(self.owner)
         if not task:
             return False
+        link = task_link(task)
+        with tracer.start_as_current_span(
+            "message_safety.worker.process",
+            links=[link] if link else (),
+            attributes={
+                "messaging.operation.type": "process",
+                "message_safety.attempt": task.attempt_count,
+            },
+        ):
+            await self._process(task)
+        return True
+
+    async def _process(self, task) -> None:
+        task_age = (datetime.now(UTC) - task.created_at).total_seconds()
         row = await self.repository.config_version(task.config_version)
         rules, detector, digest = validate_config(row.config, self.artifacts)
         if digest != row.config_sha256:
             await self.repository.retry_or_fail(task, row.config["task"]["max_attempts"])
-            return True
+            record_worker("config_error", task_age)
+            return
         stop = asyncio.Event()
         heartbeat = asyncio.create_task(
             self._heartbeat(
@@ -65,18 +92,22 @@ class Worker:
             body, _ = await collect_and_hash(
                 self.reader, attachment, max_size=row.config["file_policy"]["max_size_bytes"]
             )
+            record_dependency("s3", "get_object", "success")
             rule = detect_format(body, attachment.mime_type)
             if not rule:
                 malware = await self.antivirus.scan(one_chunk(body))
+                record_dependency("clamav", "scan", "success")
                 rule = "file.malware_detected" if malware else None
-            finished = await self.repository.finish(
-                task.id,
-                self.owner,
-                task.lease_generation,
-                allow=rule is None,
-                rule_id=rule or "safety.all_checks_passed",
-            )
+            with tracer.start_as_current_span("message_safety.worker.finalize"):
+                finished = await self.repository.finish(
+                    task.id,
+                    self.owner,
+                    task.lease_generation,
+                    allow=rule is None,
+                    rule_id=rule or "safety.all_checks_passed",
+                )
             if finished:
+                record_worker("allow" if rule is None else "deny", task_age)
                 now = datetime.now(UTC)
                 await self.repository.put_file_cache(
                     FileVerdictCache(
@@ -111,19 +142,23 @@ class Worker:
                     )
                 )
         except ObjectChanged:
-            await self.repository.finish(
-                task.id,
-                self.owner,
-                task.lease_generation,
-                allow=False,
-                rule_id="file.object_changed",
-            )
+            record_dependency("s3", "get_object", "object_changed")
+            with tracer.start_as_current_span("message_safety.worker.finalize"):
+                await self.repository.finish(
+                    task.id,
+                    self.owner,
+                    task.lease_generation,
+                    allow=False,
+                    rule_id="file.object_changed",
+                )
+            record_worker("deny", task_age)
         except DependencyFailure:
+            record_dependency("file_pipeline", "scan", "error")
             await self.repository.retry_or_fail(task, row.config["task"]["max_attempts"])
+            record_worker("retry_or_fail", task_age)
         finally:
             stop.set()
             await heartbeat
-        return True
 
     async def _heartbeat(
         self, task_id, generation: int, interval: int, lease: int, stop: asyncio.Event
@@ -144,25 +179,29 @@ class Worker:
 
 async def serve() -> None:
     settings = BootstrapSettings()
-    assert settings.database_url and settings.s3_access_key and settings.s3_secret_key
-    engine, sessions = engine_and_sessions(settings.database_url.get_secret_value())
-    worker = Worker(
-        Repository(sessions),
-        S3VersionReader(
-            settings.s3_endpoint_url,
-            settings.s3_bucket,
-            settings.s3_access_key.get_secret_value(),
-            settings.s3_secret_key.get_secret_value(),
-        ),
-        ClamAvInstream(settings.clamav_host, settings.clamav_port),
-        settings.artifacts_dir,
-    )
+    init_telemetry("worker")
+    engine = None
     try:
+        assert settings.database_url and settings.s3_access_key and settings.s3_secret_key
+        engine, sessions = engine_and_sessions(settings.database_url.get_secret_value())
+        worker = Worker(
+            Repository(sessions),
+            S3VersionReader(
+                settings.s3_endpoint_url,
+                settings.s3_bucket,
+                settings.s3_access_key.get_secret_value(),
+                settings.s3_secret_key.get_secret_value(),
+            ),
+            ClamAvInstream(settings.clamav_host, settings.clamav_port),
+            settings.artifacts_dir,
+        )
         async with asyncio.TaskGroup() as group:
             for _ in range(settings.worker_concurrency):
                 group.create_task(worker.loop())
     finally:
-        await engine.dispose()
+        if engine is not None:
+            await engine.dispose()
+        shutdown_telemetry()
 
 
 def run() -> None: