fix

Author: JqRrt

cozeloop/_client.py

@@ -19,7 +19,7 @@
 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.model.model import FinishEventInfo, TagTruncateConf, QueueConf
 from cozeloop.internal.trace.trace import default_finish_event_processor
 from cozeloop.span import SpanContext, Span
 
@@ -71,6 +71,7 @@ def new_client(
         trace_finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None,
         tag_truncate_conf: Optional[TagTruncateConf] = None,
         api_base_path: Optional[APIBasePath] = None,
+        queue_conf: Optional[QueueConf] = None,
 ) -> Client:
     cache_key = _generate_cache_key( # all args are used to generate cache key
         api_base_url,
@@ -89,6 +90,7 @@ def new_client(
         trace_finish_event_processor,
         tag_truncate_conf,
         api_base_path,
+        queue_conf,
     )
 
     with _cache_lock:
@@ -113,6 +115,7 @@ def new_client(
             trace_finish_event_processor=trace_finish_event_processor,
             tag_truncate_conf=tag_truncate_conf,
             api_base_path=api_base_path,
+            queue_conf=queue_conf,
         )
         _client_cache[cache_key] = client
         return client
@@ -143,6 +146,7 @@ def __init__(
             trace_finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None,
             tag_truncate_conf: Optional[TagTruncateConf] = None,
             api_base_path: Optional[APIBasePath] = None,
+            queue_conf: Optional[QueueConf] = 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)
@@ -201,6 +205,7 @@ def combined_processor(event_info: FinishEventInfo):
             tag_truncate_conf=tag_truncate_conf,
             span_upload_path=span_upload_path,
             file_upload_path=file_upload_path,
+            queue_conf=queue_conf,
         )
         self._prompt_provider = PromptProvider(
             workspace_id=workspace_id,

cozeloop/internal/trace/model/model.py

@@ -33,10 +33,12 @@ class UploadType(str, Enum):
     "exporter.file_flush.rate"
 ]
 
+
 @dataclass
 class FinishEventInfoExtra:
     is_root_span: bool
 
+
 @dataclass
 class FinishEventInfo:
     event_type: SpanFinishEvent
@@ -45,7 +47,14 @@ class FinishEventInfo:
     detail_msg: str
     extra_params: Optional[FinishEventInfoExtra] = None
 
+
 @dataclass
 class TagTruncateConf:
     normal_field_max_byte: int
     input_output_field_max_byte: int
+
+
+@dataclass
+class QueueConf:
+    span_queue_length: int
+    span_max_export_batch_length: int

cozeloop/internal/trace/span_processor.py

@@ -13,7 +13,7 @@
 from multipart import file_path
 
 from cozeloop.internal.trace.exporter import *
-from cozeloop.internal.trace.model.model import FinishEventInfo
+from cozeloop.internal.trace.model.model import FinishEventInfo, QueueConf
 from cozeloop.internal.trace.queue_manager import BatchQueueManager, BatchQueueManagerOptions, QUEUE_NAME_FILE_RETRY, \
     QUEUE_NAME_FILE, QUEUE_NAME_SPAN_RETRY, QUEUE_NAME_SPAN
 from cozeloop.internal.trace.span import Span
@@ -49,6 +49,7 @@ def __init__(
             client,
             upload_path: UploadPath = None,
             finish_event_processor: Optional[Callable[[FinishEventInfo], None]] = None,
+            queue_conf: Optional[QueueConf] = None,
     ):
         span_upload_path = PATH_INGEST_TRACE
         file_upload_path = PATH_UPLOAD_FILE
@@ -65,6 +66,14 @@ def __init__(
             )
         )
 
+        span_queue_length = DEFAULT_MAX_QUEUE_LENGTH
+        span_export_batch_size = DEFAULT_MAX_EXPORT_BATCH_LENGTH
+        if queue_conf:
+            if queue_conf.span_queue_length > 0:
+                span_queue_length = queue_conf.span_queue_length
+            if queue_conf.span_max_export_batch_length > 0:
+                span_export_batch_size = queue_conf.span_max_export_batch_length
+
         self.file_retry_qm = BatchQueueManager(
             BatchQueueManagerOptions(
                 queue_name=QUEUE_NAME_FILE_RETRY,
@@ -105,8 +114,8 @@ def __init__(
             BatchQueueManagerOptions(
                 queue_name=QUEUE_NAME_SPAN,
                 batch_timeout=DEFAULT_SCHEDULE_DELAY,
-                max_queue_length=DEFAULT_MAX_QUEUE_LENGTH,
-                max_export_batch_length=DEFAULT_MAX_EXPORT_BATCH_LENGTH,
+                max_queue_length=span_queue_length,
+                max_export_batch_length=span_export_batch_size,
                 max_export_batch_byte_size=DEFAULT_MAX_EXPORT_BATCH_BYTE_SIZE,
                 export_func=self._new_export_spans_func(self.exporter, self.span_retry_qm, self.file_qm, finish_event_processor),
                 finish_event_processor=finish_event_processor,

cozeloop/internal/trace/trace.py

@@ -6,7 +6,7 @@
 from typing import Dict, Optional, Callable
 
 from cozeloop.internal.trace.exporter import UploadPath
-from cozeloop.internal.trace.model.model import FinishEventInfo, TagTruncateConf
+from cozeloop.internal.trace.model.model import FinishEventInfo, TagTruncateConf, QueueConf
 from cozeloop.spec.tracespec import Runtime, RUNTIME_
 from cozeloop.internal import consts
 from cozeloop.internal.trace.span import from_header, Span, SpanContext, \
@@ -41,6 +41,7 @@ def __init__(
             tag_truncate_conf: Optional[TagTruncateConf] = None,
             span_upload_path: str = None,
             file_upload_path: str = None,
+            queue_conf: Optional[QueueConf] = None,
     ):
         self.workspace_id = workspace_id
         self.ultra_large_report = ultra_large_report
@@ -54,6 +55,7 @@ def __init__(
             http_client,
             upload_path,
             finish_event_processor,
+            queue_conf,
         )
         self.tag_truncate_conf = tag_truncate_conf