support FinishEventProcessor and some Span Set Method

Author: JqRrt

client.go

@@ -35,7 +35,7 @@ type Client interface {
 
 	// GetWorkspaceID return workspace id
 	GetWorkspaceID() string
-	// Close Close the client. Should be called before program exit.
+	// Close close the client. Should be called before program exit.
 	Close(ctx context.Context)
 }
 
@@ -62,6 +62,8 @@ type options struct {
 	promptCacheRefreshInterval time.Duration
 	promptTrace                bool
 	exporter                   trace.Exporter
+	traceFinishEventProcessor  func(ctx context.Context, info *FinishEventInfo)
+	traceTagTruncateConf       *TagTruncateConf
 }
 
 func (o *options) MD5() string {
@@ -133,10 +135,19 @@ func NewClient(opts ...Option) (Client, error) {
 			Timeout:       options.timeout,
 			UploadTimeout: options.uploadTimeout,
 		})
+	traceFinishEventProcessor := trace.DefaultFinishEventProcessor
+	if options.traceFinishEventProcessor != nil {
+		traceFinishEventProcessor = func(ctx context.Context, info *consts.FinishEventInfo) {
+			trace.DefaultFinishEventProcessor(ctx, info)
+			options.traceFinishEventProcessor(ctx, (*FinishEventInfo)(info))
+		}
+	}
 	c.traceProvider = trace.NewTraceProvider(httpClient, trace.Options{
-		WorkspaceID:      options.workspaceID,
-		UltraLargeReport: options.ultraLargeReport,
-		Exporter:         options.exporter,
+		WorkspaceID:               options.workspaceID,
+		UltraLargeReport:          options.ultraLargeReport,
+		Exporter:                  options.exporter,
+		TraceFinishEventProcessor: traceFinishEventProcessor,
+		TraceTagTruncateConf:      (*trace.TagTruncateConf)(options.traceTagTruncateConf),
 	})
 	c.promptProvider = prompt.NewPromptProvider(httpClient, c.traceProvider, prompt.Options{
 		WorkspaceID:                options.workspaceID,
@@ -247,6 +258,18 @@ func WithExporter(e trace.Exporter) Option {
 	}
 }
 
+func WithTraceFinishEventProcessor(f func(ctx context.Context, info *FinishEventInfo)) Option {
+	return func(p *options) {
+		p.traceFinishEventProcessor = f
+	}
+}
+
+func WithTraceTagTruncateConf(conf *TagTruncateConf) Option {
+	return func(p *options) {
+		p.traceTagTruncateConf = conf
+	}
+}
+
 // GetWorkspaceID return space id
 func GetWorkspaceID() string {
 	return getDefaultClient().GetWorkspaceID()
@@ -350,13 +373,27 @@ func buildAuth(opts options) (httpclient.Auth, error) {
 	return nil, ErrAuthInfoRequired
 }
 
+func SetDefaultClient(client Client) {
+	defaultClientLock.Lock()
+	defer defaultClientLock.Unlock()
+	defaultClient = client
+}
+
 func getDefaultClient() Client {
+	if defaultClient != nil {
+		return defaultClient
+	}
 	once.Do(func() {
 		var err error
-		defaultClient, err = NewClient()
+		client, err := NewClient()
 		if err != nil {
+			defaultClientLock.Lock()
 			defaultClient = &NoopClient{newClientError: err}
+			defaultClientLock.Unlock()
 		} else {
+			defaultClientLock.Lock()
+			defaultClient = client
+			defaultClientLock.Unlock()
 			sigChan := make(chan os.Signal, 1)
 			signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
 			go func() {
@@ -366,7 +403,9 @@ func getDefaultClient() Client {
 
 				logger.CtxInfof(ctx, "Received signal: %v, starting graceful shutdown...", sig)
 				defaultClient.Close(ctx)
+				defaultClientLock.Lock()
 				defaultClient = &NoopClient{newClientError: consts.ErrClientClosed}
+				defaultClientLock.Unlock()
 				logger.CtxInfof(ctx, "Graceful shutdown finished.")
 				os.Exit(0)
 			}()
@@ -376,9 +415,10 @@ func getDefaultClient() Client {
 }
 
 var (
-	defaultClient Client
-	once          sync.Once
-	clientCache   sync.Map // client cache to avoid creating multiple clients with the same options
+	defaultClient     Client
+	defaultClientLock sync.RWMutex
+	once              sync.Once
+	clientCache       sync.Map // client cache to avoid creating multiple clients with the same options
 )
 
 type loopClient struct {
@@ -426,7 +466,7 @@ func (c *loopClient) PromptFormat(ctx context.Context, loopPrompt *entity.Prompt
 
 func (c *loopClient) StartSpan(ctx context.Context, name, spanType string, opts ...StartSpanOption) (context.Context, Span) {
 	if c.closed {
-		return ctx, defaultNoopSpan
+		return ctx, DefaultNoopSpan
 	}
 	config := trace.StartSpanOptions{}
 	for _, opt := range opts {
@@ -438,25 +478,25 @@ func (c *loopClient) StartSpan(ctx context.Context, name, spanType string, opts
 	ctx, span, err := c.traceProvider.StartSpan(ctx, name, spanType, config)
 	if err != nil {
 		logger.CtxWarnf(ctx, "start span failed, return noop span. %v", err)
-		return ctx, defaultNoopSpan
+		return ctx, DefaultNoopSpan
 	}
 	return ctx, span
 }
 
 func (c *loopClient) GetSpanFromContext(ctx context.Context) Span {
 	if c.closed {
-		return defaultNoopSpan
+		return DefaultNoopSpan
 	}
 	span := c.traceProvider.GetSpanFromContext(ctx)
 	if span == nil {
-		return defaultNoopSpan
+		return DefaultNoopSpan
 	}
 	return span
 }
 
 func (c *loopClient) GetSpanFromHeader(ctx context.Context, header map[string]string) SpanContext {
 	if c.closed {
-		return defaultNoopSpan
+		return DefaultNoopSpan
 	}
 	return c.traceProvider.GetSpanFromHeader(ctx, header)
 }

const.go

@@ -5,6 +5,7 @@ package cozeloop
 
 import (
 	"github.com/coze-dev/cozeloop-go/internal/consts"
+	"github.com/coze-dev/cozeloop-go/internal/trace"
 )
 
 const (
@@ -19,3 +20,17 @@ const (
 	// ComBaseURL = consts.ComBaseURL
 	CnBaseURL = consts.CnBaseURL
 )
+
+// SpanFinishEvent finish inner event
+type SpanFinishEvent consts.SpanFinishEvent
+
+const (
+	SpanFinishEventSpanQueueEntryRate = SpanFinishEvent(consts.SpanFinishEventSpanQueueEntryRate)
+	SpanFinishEventFileQueueEntryRate = SpanFinishEvent(consts.SpanFinishEventFileQueueEntryRate)
+	SpanFinishEventFlushSpanRate      = SpanFinishEvent(consts.SpanFinishEventFlushSpanRate)
+	SpanFinishEventFlushFileRate      = SpanFinishEvent(consts.SpanFinishEventFlushFileRate)
+)
+
+type FinishEventInfo consts.FinishEventInfo
+
+type TagTruncateConf trace.TagTruncateConf

entity/export.go

@@ -1,18 +1,20 @@
 package entity
 
 type UploadSpan struct {
-	StartedATMicros  int64              `json:"started_at_micros"`
+	StartedATMicros  int64              `json:"started_at_micros"` // start time in microseconds
+	LogID            string             `json:"log_id"`            // the custom log id, identify different query.
 	SpanID           string             `json:"span_id"`
 	ParentID         string             `json:"parent_id"`
 	TraceID          string             `json:"trace_id"`
-	DurationMicros   int64              `json:"duration_micros"`
-	WorkspaceID      string             `json:"workspace_id"`
+	DurationMicros   int64              `json:"duration_micros"` // duration in microseconds
+	ServiceName      string             `json:"service_name"`    // the custom service name, identify different services.
+	WorkspaceID      string             `json:"workspace_id"`    // cozeloop workspace id
 	SpanName         string             `json:"span_name"`
 	SpanType         string             `json:"span_type"`
 	StatusCode       int32              `json:"status_code"`
 	Input            string             `json:"input"`
 	Output           string             `json:"output"`
-	ObjectStorage    string             `json:"object_storage"`
+	ObjectStorage    string             `json:"object_storage"` // file object storage
 	SystemTagsString map[string]string  `json:"system_tags_string"`
 	SystemTagsLong   map[string]int64   `json:"system_tags_long"`
 	SystemTagsDouble map[string]float64 `json:"system_tags_double"`

internal/consts/finish_event.go

@@ -0,0 +1,23 @@
+package consts
+
+type SpanFinishEvent string
+
+const (
+	SpanFinishEventSpanQueueEntryRate SpanFinishEvent = "queue_manager.span_entry.rate"
+	SpanFinishEventFileQueueEntryRate SpanFinishEvent = "queue_manager.file_entry.rate"
+
+	SpanFinishEventFlushSpanRate SpanFinishEvent = "exporter.span_flush.rate"
+	SpanFinishEventFlushFileRate SpanFinishEvent = "exporter.file_flush.rate"
+)
+
+type FinishEventInfo struct {
+	EventType   SpanFinishEvent
+	IsEventFail bool
+	ItemNum     int // maybe multiple span is processed in one event
+	DetailMsg   string
+	ExtraParams *FinishEventInfoExtra
+}
+
+type FinishEventInfoExtra struct {
+	IsRootSpan bool
+}

internal/httpclient/client.go

@@ -203,6 +203,8 @@ func setAuthorizationHeader(ctx context.Context, request *http.Request, auth Aut
 }
 
 func parseResponse(ctx context.Context, url string, response *http.Response, resp OpenAPIResponse) error {
+	defer response.Body.Close()
+
 	logID := response.Header.Get(consts.LogIDHeader)
 	respBody, err := io.ReadAll(response.Body)
 	if err != nil {

internal/trace/exporter.go

@@ -52,12 +52,10 @@ func (e *SpanExporter) ExportFiles(ctx context.Context, files []*entity.UploadFi
 		resp := httpclient.BaseResponse{}
 		err := e.client.UploadFile(ctx, pathUploadFile, file.TosKey, bytes.NewReader([]byte(file.Data)), map[string]string{"workspace_id": file.SpaceID}, &resp)
 		if err != nil {
-			logger.CtxErrorf(ctx, "export files[%s] fail, err:[%v], retry later", file.TosKey, err)
-			return err
+			return consts.NewError(fmt.Sprintf("export files[%s] fail", file.TosKey)).Wrap(err)
 		}
 		if resp.GetCode() != 0 { // todo: some err code do not need retry
-			logger.CtxErrorf(ctx, "export files[%s] fail, code:[%v], msg:[%v] retry later", file.TosKey, resp.GetCode(), resp.GetMsg())
-			return consts.ErrRemoteService
+			return consts.NewError(fmt.Sprintf("export files[%s] fail, code:[%v], msg:[%v] retry later", file.TosKey, resp.GetCode(), resp.GetMsg()))
 		}
 		logger.CtxDebugf(ctx, "uploadFile end, file name: %s", file.Name)
 	}
@@ -74,12 +72,10 @@ func (e *SpanExporter) ExportSpans(ctx context.Context, ss []*entity.UploadSpan)
 	resp := httpclient.BaseResponse{}
 	err = e.client.Post(ctx, pathIngestTrace, UploadSpanData{ss}, &resp)
 	if err != nil {
-		logger.CtxErrorf(ctx, "export spans fail, span count: [%d], err:[%v]", len(ss), err)
-		return err
+		return consts.NewError(fmt.Sprintf("export spans fail, span count: [%d]", len(ss))).Wrap(err)
 	}
 	if resp.GetCode() != 0 { // todo: some err code do not need retry
-		logger.CtxErrorf(ctx, "export spans fail, span count: [%d], code:[%v], msg:[%v]", len(ss), resp.GetCode(), resp.GetMsg())
-		return consts.ErrRemoteService
+		return consts.NewError(fmt.Sprintf("export spans fail, span count: [%d], code:[%v], msg:[%v]", len(ss), resp.GetCode(), resp.GetMsg()))
 	}
 
 	return
@@ -107,10 +103,12 @@ func transferToUploadSpanAndFile(ctx context.Context, spans []*Span) ([]*entity.
 		systemTagStrM, systemTagLongM, systemTagDoubleM, _ := parseTag(span.SystemTagMap, true)
 		resSpan = append(resSpan, &entity.UploadSpan{
 			StartedATMicros:  span.GetStartTime().UnixMicro(),
+			LogID:            span.GetLogID(),
 			SpanID:           span.GetSpanID(),
 			ParentID:         span.GetParentID(),
 			TraceID:          span.GetTraceID(),
 			DurationMicros:   span.GetDuration(),
+			ServiceName:      span.GetServiceName(),
 			WorkspaceID:      span.GetSpaceID(),
 			SpanName:         span.GetSpanName(),
 			SpanType:         span.GetSpanType(),

internal/trace/noop_span.go

@@ -34,6 +34,9 @@ func (n noopSpan) SetInputTokens(ctx context.Context, inputTokens int)
 func (n noopSpan) SetOutputTokens(ctx context.Context, outputTokens int)                 {}
 func (n noopSpan) SetStartTimeFirstResp(ctx context.Context, startTimeFirstResp int64)   {}
 func (n noopSpan) SetRuntime(ctx context.Context, runtime tracespec.Runtime)             {}
+func (n noopSpan) SetServiceName(ctx context.Context, serviceName string)                {}
+func (n noopSpan) SetLogID(ctx context.Context, logID string)                            {}
+func (n noopSpan) SetFinishTime(finishTime time.Time)                                    {}
 
 // implement of Span
 func (n noopSpan) SetTags(ctx context.Context, tagKVs map[string]interface{})     {}

internal/trace/queue_manager.go

@@ -13,14 +13,23 @@ package trace
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"sync/atomic"
 	"time"
 
+	"github.com/coze-dev/cozeloop-go/internal/consts"
 	"github.com/coze-dev/cozeloop-go/internal/logger"
 	"github.com/coze-dev/cozeloop-go/internal/util"
 )
 
+const (
+	queueNameSpan      = "span"
+	queueNameSpanRetry = "span_retry"
+	queueNameFile      = "file"
+	queueNameFileRetry = "file_retry"
+)
+
 type exportFunc func(ctx context.Context, s []interface{})
 
 // QueueManager is a queue that batches spans and exports them
@@ -37,7 +46,8 @@ type batchQueueManagerOptions struct {
 	maxExportBatchLength   int
 	maxExportBatchByteSize int
 
-	exportFunc exportFunc
+	exportFunc           exportFunc
+	finishEventProcessor func(ctx context.Context, info *consts.FinishEventInfo)
 }
 
 func newBatchQueueManager(o batchQueueManagerOptions) *BatchQueueManager {
@@ -181,18 +191,41 @@ func (b *BatchQueueManager) Enqueue(ctx context.Context, sd interface{}, byteSiz
 	if atomic.LoadInt32(&b.stopped) != 0 {
 		return
 	}
-
+	var extraParams *consts.FinishEventInfoExtra
+	var eventType = consts.SpanFinishEventFileQueueEntryRate
+	var errMsg string
+	var isFail bool
 	select {
 	case b.queue <- sd:
 		b.batchMutex.Lock()
 		b.batchByteSize += byteSize
 		b.batchMutex.Unlock()
-		//logger.CtxDebugf(ctx, "%s queue length: %d", b.o.queueName, len(b.queue))
-		return
 	default: // queue is full, not block, drop
-		logger.CtxErrorf(ctx, "%s queue is full, dropped item", b.o.queueName)
+		errMsg = fmt.Sprintf("%s queue is full, dropped item", b.o.queueName)
+		isFail = true
 		atomic.AddUint32(&b.dropped, 1)
 	}
+
+	switch b.o.queueName {
+	case queueNameSpan, queueNameSpanRetry:
+		eventType = consts.SpanFinishEventSpanQueueEntryRate
+		span, ok := sd.(*Span)
+		if ok {
+			extraParams = &consts.FinishEventInfoExtra{
+				IsRootSpan: span.IsRootSpan(),
+			}
+		}
+	default:
+	}
+	if b.o.finishEventProcessor != nil {
+		b.o.finishEventProcessor(ctx, &consts.FinishEventInfo{
+			EventType:   eventType,
+			IsEventFail: isFail,
+			ItemNum:     1,
+			DetailMsg:   errMsg,
+			ExtraParams: extraParams,
+		})
+	}
 	return
 }