support FinishEventProcessor and some Span Set Method

Author: JqRrt

cozeloop/_client.py

@@ -7,7 +7,9 @@
 import os
 import threading
 from datetime import datetime
-from typing import Dict, Any, List, Optional
+from typing import Dict, Any, List, Optional, Callable
+
+import httpx
 
 from cozeloop.client import Client
 from cozeloop._noop import NOOP_SPAN, _NoopClient
@@ -17,6 +19,8 @@
 from cozeloop.internal.httpclient import Auth
 from cozeloop.internal.prompt import PromptProvider
 from cozeloop.internal.trace import TraceProvider
+from cozeloop.internal.trace.model.model import FinishEventInfo, TagTruncateConf
+from cozeloop.internal.trace.trace import default_finish_event_processor
 from cozeloop.span import SpanContext, Span
 
 logger = logging.getLogger(__name__)
@@ -54,8 +58,11 @@ def new_client(
         prompt_cache_max_count: int = consts.DEFAULT_PROMPT_CACHE_MAX_COUNT,
         prompt_cache_refresh_interval: int = consts.DEFAULT_PROMPT_CACHE_REFRESH_INTERVAL,
         prompt_trace: bool = False,
+        http_client: Optional[httpx.Client] = None,
+        trace_finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None,
+        tag_truncate_conf: Optional[TagTruncateConf] = None,
 ) -> Client:
-    cache_key = _generate_cache_key(
+    cache_key = _generate_cache_key( # all args are used to generate cache key
         api_base_url,
         workspace_id,
         api_token,
@@ -67,7 +74,10 @@ def new_client(
         ultra_large_report,
         prompt_cache_max_count,
         prompt_cache_refresh_interval,
-        prompt_trace
+        prompt_trace,
+        http_client,
+        trace_finish_event_processor,
+        tag_truncate_conf,
     )
 
     with _cache_lock:
@@ -88,6 +98,8 @@ def new_client(
             prompt_cache_max_count=prompt_cache_max_count,
             prompt_cache_refresh_interval=prompt_cache_refresh_interval,
             prompt_trace=prompt_trace,
+            arg_http_client=http_client,
+            trace_finish_event_processor=trace_finish_event_processor,
         )
         _client_cache[cache_key] = client
         return client
@@ -113,7 +125,10 @@ def __init__(
             ultra_large_report: bool = False,
             prompt_cache_max_count: int = consts.DEFAULT_PROMPT_CACHE_MAX_COUNT,
             prompt_cache_refresh_interval: int = consts.DEFAULT_PROMPT_CACHE_REFRESH_INTERVAL,
-            prompt_trace: bool = False
+            prompt_trace: bool = False,
+            arg_http_client: Optional[httpx.Client] = None,
+            trace_finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None,
+            tag_truncate_conf: Optional[TagTruncateConf] = None,
     ):
         workspace_id = self._get_from_env(workspace_id, ENV_WORKSPACE_ID)
         api_base_url = self._get_from_env(api_base_url, ENV_API_BASE_URL)
@@ -136,6 +151,8 @@ def __init__(
 
         self._workspace_id = workspace_id
         inner_client = httpclient.HTTPClient()
+        if arg_http_client:
+            inner_client = arg_http_client
         auth = self._build_auth(
             api_base_url=api_base_url,
             http_client=inner_client,
@@ -151,10 +168,18 @@ def __init__(
             timeout=timeout,
             upload_timeout=upload_timeout
         )
+        finish_pro = default_finish_event_processor
+        if trace_finish_event_processor:
+            def combined_processor(event_info: FinishEventInfo):
+                default_finish_event_processor(event_info)
+                trace_finish_event_processor(event_info)
+            finish_pro = combined_processor
         self._trace_provider = TraceProvider(
             http_client=http_client,
             workspace_id=workspace_id,
-            ultra_large_report=ultra_large_report
+            ultra_large_report=ultra_large_report,
+            finish_event_processor=finish_pro,
+            tag_truncate_conf=tag_truncate_conf,
         )
         self._prompt_provider = PromptProvider(
             workspace_id=workspace_id,
@@ -234,7 +259,7 @@ def start_span(
             else:
                 return self._trace_provider.start_span(name=name, span_type=span_type, start_time=start_time,
                                                        parent_span_id=child_of.span_id, trace_id=child_of.trace_id,
-                                                       baggage=child_of.baggage, start_new_trace=start_new_trace)
+                                                       baggage=child_of.baggage(), start_new_trace=start_new_trace)
         except Exception as e:
             logger.warning(f"Start span failed, returning noop span. Error: {e}")
             return NOOP_SPAN

cozeloop/internal/httpclient/client.py

@@ -43,10 +43,10 @@ def _set_headers(self, headers: Optional[Dict[str, str]] = None) -> Dict[str, st
             res.update(headers)
         res[consts.AUTHORIZE_HEADER] = f"Bearer {self.auth.token}"
 
-        tt_env = os.getenv("x-tt-env")
+        tt_env = os.getenv("x_tt_env")
         if tt_env:
             res["x-tt-env"] = tt_env
-        ppe_env = os.getenv("x-use-ppe")
+        ppe_env = os.getenv("x_use_ppe")
         if ppe_env:
             res["x-use-ppe"] = "1"
 

cozeloop/internal/trace/__init__.py

@@ -6,4 +6,3 @@
 from .trace import TraceProvider
 
 from .trace import Span
-

cozeloop/internal/trace/exporter.py

@@ -18,10 +18,10 @@
 logger = logging.getLogger(__name__)
 
 class Exporter:
-    def export_spans(self, ctx: dict, spans: List['UploadSpan']) -> bool:
+    def export_spans(self, ctx: dict, spans: List['UploadSpan']) -> (bool, str):
         raise NotImplementedError
 
-    def export_files(self, ctx: dict, files: List['UploadFile']) -> bool:
+    def export_files(self, ctx: dict, files: List['UploadFile']) -> (bool, str):
         raise NotImplementedError
 
 
@@ -40,7 +40,7 @@ class SpanExporter(Exporter):
     def __init__(self, client: Client):
         self.client = client
 
-    def export_files(self, ctx: dict, files: List['UploadFile']) -> bool:
+    def export_files(self, ctx: dict, files: List['UploadFile']) -> (bool, str):
         for file in files:
             if not file:
                 continue
@@ -55,18 +55,14 @@ def export_files(self, ctx: dict, files: List['UploadFile']) -> bool:
                     {"workspace_id": file.space_id},
                 )
                 if resp.code != 0: # todo: some err code do not need retry
-                    logger.error(f"export files[{file.tos_key}] fail, code:[{resp.code}], msg:[{resp.msg}]")
-                    return False
+                    return False, f"export files[{file.tos_key}] fail, code:[{resp.code}], msg:[{resp.msg}]"
             except Exception as e:
-                logger.error(f"export files[{file.tos_key}] fail, err:[{e}], file.name:[{file.name}]")
-                return False
+                return False, f"export files[{file.tos_key}] fail, err:[{e}], file.name:[{file.name}]"
 
             logger.debug(f"uploadFile end, file name: {file.name}")
-        return True
-
-    def export_spans(self, ctx: dict, spans: List['UploadSpan']) -> bool:
-        logger.debug(f"export spans, spans count: {len(spans)}")
+        return True, ""
 
+    def export_spans(self, ctx: dict, spans: List['UploadSpan']) -> (bool, str):
         if not spans:
             return True
 
@@ -77,13 +73,11 @@ def export_spans(self, ctx: dict, spans: List['UploadSpan']) -> bool:
                 UploadSpanData(spans=spans),
             )
             if resp.code != 0: # todo: some err code do not need retry
-                logger.error(f"export spans fail, code:[{resp.code}], msg:[{resp.msg}]")
-                return False
+                return False, f"export spans fail, code:[{resp.code}], msg:[{resp.msg}]"
         except Exception as e:
-            logger.error(f"export spans fail, err:[{e}]")
-            return False
+            return False, f"export spans fail, err:[{e}]"
 
-        return True
+        return True, ""
 
 
 class UploadSpanData(BaseModel):
@@ -92,10 +86,12 @@ class UploadSpanData(BaseModel):
 
 class UploadSpan(BaseModel):
     started_at_micros: int
+    log_id: str
     span_id: str
     parent_id: str
     trace_id: str
     duration_micros: int
+    service_name: str
     workspace_id: str
     span_name: str
     span_type: str
@@ -137,10 +133,12 @@ def transfer_to_upload_span_and_file(spans: List['Span']) -> (List[UploadSpan],
 
         res_span.append(UploadSpan(
             started_at_micros=int(span.start_time.timestamp() * 1_000_000),
+            log_id=span.log_id,
             span_id=span.span_id,
             parent_id=span.parent_span_id,
             trace_id=span.trace_id,
             duration_micros=span.get_duration(),
+            service_name=span.service_name,
             workspace_id=span.get_space_id(),
             span_name=span.get_span_name(),
             span_type=span.get_span_type(),

cozeloop/internal/trace/model/model.py

@@ -4,7 +4,8 @@
 from enum import Enum
 
 from pydantic import BaseModel
-from typing import List, Optional
+from typing import List, Optional, Literal
+from pydantic.dataclasses import dataclass
 
 
 class ObjectStorage(BaseModel):
@@ -23,3 +24,28 @@ class Attachment(BaseModel):
 class UploadType(str, Enum):
     LONG = 1
     MULTI_MODALITY = 2
+
+
+SpanFinishEvent = Literal[
+    "queue_manager.span_entry.rate",
+    "queue_manager.file_entry.rate",
+    "exporter.span_flush.rate",
+    "exporter.file_flush.rate"
+]
+
+@dataclass
+class FinishEventInfoExtra:
+    is_root_span: bool
+
+@dataclass
+class FinishEventInfo:
+    event_type: SpanFinishEvent
+    is_event_fail: bool
+    item_num: int  # maybe multiple span is processed in one event
+    detail_msg: str
+    extra_params: Optional[FinishEventInfoExtra] = None
+
+@dataclass
+class TagTruncateConf:
+    normal_field_max_byte: int
+    input_output_field_max_byte: int

cozeloop/internal/trace/noop_span.py

@@ -110,6 +110,12 @@ def set_start_time_first_resp(self, start_time_first_resp: int) -> None:
     def set_runtime(self, runtime: Runtime) -> None:
         pass
 
+    def set_service_name(self, service_name: str) -> None:
+        pass
+
+    def set_log_id(self, log_id: str) -> None:
+        pass
+
     def __enter__(self):
         return self
 

cozeloop/internal/trace/queue_manager.py

@@ -12,13 +12,20 @@
 import logging
 import queue
 import threading
-from typing import Any, Callable, List
+from typing import Any, Callable, List, Optional
 from queue import Queue
 
 from pydantic import BaseModel
 
+from cozeloop.internal.trace.model.model import FinishEventInfo, FinishEventInfoExtra
+
 logger = logging.getLogger(__name__)
 
+QUEUE_NAME_SPAN = "span"
+QUEUE_NAME_SPAN_RETRY = "span_retry"
+QUEUE_NAME_FILE = "file"
+QUEUE_NAME_FILE_RETRY = "file_retry"
+
 class QueueManager:
     def enqueue(self, s: Any, byte_size: int):
         raise NotImplementedError
@@ -37,6 +44,7 @@ class BatchQueueManagerOptions(BaseModel):
     max_export_batch_length: int
     max_export_batch_byte_size: int
     export_func: Callable[[dict, List[Any]], None]
+    finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None
 
 
 class BatchQueueManager(QueueManager):
@@ -112,8 +120,6 @@ def _do_export(self):
         self._do_export_batch()
 
     def _do_export_batch(self):
-        logger.debug(
-            f"{self.options.queue_name} queue _do_export_batch, len: {len(self.batch)}")
         with self.batch_lock:
             if self.batch:
                 if self.export_func:
@@ -125,20 +131,40 @@ def enqueue(self, item: Any, byte_size: int):
         if self.stop_event.is_set():
             return
 
+        is_fail = False
+        detail_msg = ""
         try:
             self.queue.put_nowait(item)
             if self.queue.qsize() >= self.options.max_queue_length:
                 with self.condition:
                     self.condition.notify()
-            logger.debug(f"{self.options.queue_name} enqueue, queue length: {self.queue.qsize()}")
+            detail_msg = f"{self.options.queue_name} enqueue, queue length: {self.queue.qsize()}"
         except queue.Full:
-            logger.error(
-                f"{self.options.queue_name} queue is full, dropped span")
+            is_fail = True
+            detail_msg = f"{self.options.queue_name} queue is full, dropped span"
             self.dropped += 1
         else:
             with self.batch_lock:
                 self.batch_byte_size += byte_size
 
+        event_typ = "queue_manager.file_entry.rate"
+        extra_params = FinishEventInfoExtra(is_root_span=False)
+        if self.options.queue_name == QUEUE_NAME_SPAN or self.options.queue_name == QUEUE_NAME_SPAN_RETRY:
+            event_typ = "queue_manager.span_entry.rate"
+            if item.is_root_span():
+                extra_params = FinishEventInfoExtra(
+                    is_root_span=True,
+                )
+
+        if self.options and self.options.finish_event_processor:
+            self.options.finish_event_processor(FinishEventInfo(
+                event_type=event_typ,
+                is_event_fail=is_fail,
+                item_num=1,
+                detail_msg=detail_msg,
+                extra_params=extra_params
+            ))
+
     def shutdown(self) -> bool:
         if self.stop_event.is_set():
             return True