diff --git a/client_test.go b/client_test.go index bed985f..88d21c3 100644 --- a/client_test.go +++ b/client_test.go @@ -4,6 +4,7 @@ package cozeloop import ( + "context" "testing" . "github.com/smartystreets/goconvey/convey" @@ -22,3 +23,28 @@ func TestNewClient(t *testing.T) { So(client1, ShouldNotEqual, client3) }) } + +func TestNoSample(t *testing.T) { + Convey("Test NoSample spans", t, func() { + client, err := NewClient(WithWorkspaceID("test-workspace"), WithAPIToken("test-token")) + So(err, ShouldBeNil) + So(client, ShouldNotBeNil) + + ctx := context.Background() + + // 1. Create a normal span + normalCtx, normalSpan := client.StartSpan(ctx, "normal-span", "test-type") + So(normalSpan, ShouldNotBeNil) + So(normalSpan.GetSpanID(), ShouldNotBeEmpty) + + // 2. Create a NoSample span using WithNoSample() + noSampleCtx, noSampleSpan := client.StartSpan(normalCtx, "nosample-span", "test-type", WithNoSample()) + So(noSampleSpan, ShouldNotBeNil) + So(noSampleSpan.GetSpanID(), ShouldBeEmpty) // NoopSpan has empty SpanID + + // 3. Create a child span from NoSample span - should also be NoopSpan + _, childSpan := client.StartSpan(noSampleCtx, "child-span", "test-type") + So(childSpan, ShouldNotBeNil) + So(childSpan.GetSpanID(), ShouldBeEmpty) // Child of NoopSpan is also NoopSpan + }) +} diff --git a/internal/prompt/prompt_hub.go b/internal/prompt/prompt_hub.go index a43e0ea..af393ff 100644 --- a/internal/prompt/prompt_hub.go +++ b/internal/prompt/prompt_hub.go @@ -9,6 +9,7 @@ import ( "io" "time" + "github.com/coze-dev/cozeloop-go/internal/span" "github.com/valyala/fasttemplate" "github.com/coze-dev/cozeloop-go/entity" @@ -60,7 +61,7 @@ func NewPromptProvider(httpClient *httpclient.Client, traceProvider *trace.Provi func (p *Provider) GetPrompt(ctx context.Context, param GetPromptParam, options GetPromptOptions) (prompt *entity.Prompt, err error) { if p.config.PromptTrace && p.traceProvider != nil { - var promptHubSpan *trace.Span + var promptHubSpan span.Span var spanErr error ctx, promptHubSpan, spanErr = p.traceProvider.StartSpan(ctx, consts.TracePromptHubSpanName, tracespec.VPromptHubSpanType, trace.StartSpanOptions{Scene: tracespec.VScenePromptHub}) @@ -135,7 +136,7 @@ func (p *Provider) PromptFormat(ctx context.Context, prompt *entity.Prompt, vari return nil, nil } if p.config.PromptTrace && p.traceProvider != nil { - var promptTemplateSpan *trace.Span + var promptTemplateSpan span.Span var spanErr error ctx, promptTemplateSpan, spanErr = p.traceProvider.StartSpan(ctx, consts.TracePromptTemplateSpanName, tracespec.VPromptTemplateSpanType, trace.StartSpanOptions{Scene: tracespec.VScenePromptTemplate}) diff --git a/internal/span/span.go b/internal/span/span.go new file mode 100644 index 0000000..6563a00 --- /dev/null +++ b/internal/span/span.go @@ -0,0 +1,140 @@ +// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +// SPDX-License-Identifier: MIT + +package span + +import ( + "context" + "time" + + "github.com/coze-dev/cozeloop-go/entity" + "github.com/coze-dev/cozeloop-go/spec/tracespec" +) + +// Span is the interface for span. +type Span interface { + SpanContext + commonSpanSetter + + // SetTags sets business custom tags. + SetTags(ctx context.Context, tagKVs map[string]interface{}) + + // SetBaggage sets tags and also passes these tags to other downstream spans (assuming + // the user uses ToHeader and FromHeader to handle header passing between services). + SetBaggage(ctx context.Context, baggageItems map[string]string) + + // Finish The span will be reported only after an explicit call to Finish. + // Under the hood, it is actually placed in an asynchronous queue waiting to be reported. + Finish(ctx context.Context) + + // GetStartTime returns the start time of the Span. + GetStartTime() time.Time + + // ToHeader Convert the span to headers. Used for cross-process correlation. + ToHeader() (map[string]string, error) +} + +// Set system-defined fields +type commonSpanSetter interface { + // SetInput key: `input` + // Input information. The input will be serialized into a JSON string. + // You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec + // Or you can use any struct you like. + SetInput(ctx context.Context, input interface{}) + + // SetOutput key: `output` + // Output information. The output will be serialized into a JSON string. + // You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec + // Or you can use any struct you like. + SetOutput(ctx context.Context, output interface{}) + + // SetError key: `error` + // Set error message. + SetError(ctx context.Context, err error) + + // SetStatusCode key: `status_code` + // Set status code. A non-zero code is considered an exception. + SetStatusCode(ctx context.Context, code int) + + // SetUserID key: `user_id` + // Set user id. + SetUserID(ctx context.Context, userID string) + SetUserIDBaggage(ctx context.Context, userID string) + + // SetMessageID key: `message_id` + // Set message id. + SetMessageID(ctx context.Context, messageID string) + SetMessageIDBaggage(ctx context.Context, messageID string) + + // SetThreadID key: `thread_id` + // Set thread id. + // thread_id is used to correlate multiple requests and is passed in by the business. + SetThreadID(ctx context.Context, threadID string) + SetThreadIDBaggage(ctx context.Context, threadID string) + + // SetPrompt key: `prompt + // Associated with PromptKey and PromptVersion, it will write two tags: prompt_key and prompt_version. + // SetPrompt is used to set the PromptKey and PromptVersion to tag. + SetPrompt(ctx context.Context, prompt entity.Prompt) + + // SetModelProvider key: `model_provider` + // The provider of the LLM, such as OpenAI, etc. + SetModelProvider(ctx context.Context, modelProvider string) + + // SetModelName key: `model_name` + // The Name of the LLM model, such as: gpt-4-1106-preview. + SetModelName(ctx context.Context, modelName string) + + // SetModelCallOptions key: `call_options` + // The call options of the LLM, such as temperature, max_tokens, etc. + // The recommended standard format is CallOption of spec package + SetModelCallOptions(ctx context.Context, callOptions interface{}) + + // SetInputTokens key: `input_tokens` + // The usage of input tokens. When the value of input_tokens is set, + // It will be automatically summed with output_tokens to calculate the tokens tag. + SetInputTokens(ctx context.Context, inputTokens int) + + // SetOutputTokens key: `output_tokens` + // The usage of output tokens. When the value of output_tokens is set, + // It will be automatically summed with input_tokens to calculate the tokens tag. + SetOutputTokens(ctx context.Context, outputTokens int) + + // SetStartTimeFirstResp key: `start_time_first_resp` + // Timestamp of the first packet return from LLM, unit: microseconds. + // When `start_time_first_resp` is set, a tag named `latency_first_resp` calculated + // based on the span's StartTime will be added, meaning the latency for the first packet. + SetStartTimeFirstResp(ctx context.Context, startTimeFirstResp int64) + + // SetRuntime key: `runtime` + // The runtime of the LLM, such as language, library, scene, etc. + // The recommended standard format is Runtime of spec package + SetRuntime(ctx context.Context, runtime tracespec.Runtime) + + // SetServiceName + // set the custom service name, identify different services. + SetServiceName(ctx context.Context, serviceName string) + + // SetLogID + // set the custom log id, identify different query. + SetLogID(ctx context.Context, logID string) + + // SetFinishTime + // Default is time.Now() when span Finish(). DO NOT set unless you do not use default time. + SetFinishTime(finishTime time.Time) + + // SetSystemTags + // set the system tags. DO NOT set unless you know what you are doing. + SetSystemTags(ctx context.Context, systemTags map[string]interface{}) + + // SetDeploymentEnv + // set the deployment env, identify custom env. + SetDeploymentEnv(ctx context.Context, deploymentEnv string) +} + +// SpanContext is the interface for span Baggage transfer. +type SpanContext interface { + GetSpanID() string + GetTraceID() string + GetBaggage() map[string]string +} diff --git a/internal/trace/exporter_test.go b/internal/trace/exporter_test.go index 2061e04..f072a08 100644 --- a/internal/trace/exporter_test.go +++ b/internal/trace/exporter_test.go @@ -8,13 +8,14 @@ import ( "testing" . "github.com/bytedance/mockey" + "github.com/coze-dev/cozeloop-go/entity" "github.com/coze-dev/cozeloop-go/internal/httpclient" . "github.com/smartystreets/goconvey/convey" ) func Test_ExportSpans(t *testing.T) { ctx := context.Background() - spans := []*UploadSpan{{}, {}} + spans := []*entity.UploadSpan{{}, {}} PatchConvey("Test transferToUploadSpanAndFile failed", t, func() { Mock((*httpclient.Client).Post).Return(nil).Build() diff --git a/internal/trace/noop_span.go b/internal/trace/noop_span.go index 3b0ff73..025850d 100644 --- a/internal/trace/noop_span.go +++ b/internal/trace/noop_span.go @@ -8,11 +8,14 @@ import ( "time" "github.com/coze-dev/cozeloop-go/entity" + "github.com/coze-dev/cozeloop-go/internal/span" "github.com/coze-dev/cozeloop-go/spec/tracespec" ) var DefaultNoopSpan = &noopSpan{} +var _ span.Span = noopSpan{} + type noopSpan struct{} // implement of commonSpanSetter diff --git a/internal/trace/queue_manager_test.go b/internal/trace/queue_manager_test.go index 7fff147..87bd41e 100644 --- a/internal/trace/queue_manager_test.go +++ b/internal/trace/queue_manager_test.go @@ -15,7 +15,7 @@ import ( func Test_GetBatchSpanProcessor(t *testing.T) { ctx := context.Background() httpClient := &httpclient.Client{} - spanQM := NewBatchSpanProcessor(nil, httpClient, nil, nil) + spanQM := NewBatchSpanProcessor(nil, httpClient, nil, nil, nil) PatchConvey("Test GetBatchSpanProcessor", t, func() { PatchConvey("Test with valid inputs", func() { diff --git a/internal/trace/span.go b/internal/trace/span.go index de6283a..2f98ba4 100644 --- a/internal/trace/span.go +++ b/internal/trace/span.go @@ -19,6 +19,7 @@ import ( "github.com/coze-dev/cozeloop-go/internal" "github.com/coze-dev/cozeloop-go/internal/consts" "github.com/coze-dev/cozeloop-go/internal/logger" + "github.com/coze-dev/cozeloop-go/internal/span" "github.com/coze-dev/cozeloop-go/internal/util" "github.com/coze-dev/cozeloop-go/spec/tracespec" ) @@ -28,6 +29,8 @@ const ( spanFinished = 1 ) +var _ span.SpanContext = (*SpanContext)(nil) + type SpanContext struct { SpanID string TraceID string @@ -46,6 +49,8 @@ func (s *SpanContext) GetBaggage() map[string]string { return s.Baggage } +var _ span.Span = (*Span)(nil) + type Span struct { // span context param SpanContext diff --git a/internal/trace/span_test.go b/internal/trace/span_test.go index b046ed0..255d300 100644 --- a/internal/trace/span_test.go +++ b/internal/trace/span_test.go @@ -129,7 +129,7 @@ func Test_Finish(t *testing.T) { httpClient := httpclient.NewClient("", nil, nil, nil) s := &Span{ isFinished: 0, - spanProcessor: NewBatchSpanProcessor(nil, httpClient, nil, nil), + spanProcessor: NewBatchSpanProcessor(nil, httpClient, nil, nil, nil), lock: sync.RWMutex{}, TagMap: make(map[string]interface{}), } diff --git a/internal/trace/trace.go b/internal/trace/trace.go index 6e7f806..d6750b3 100644 --- a/internal/trace/trace.go +++ b/internal/trace/trace.go @@ -11,6 +11,7 @@ import ( "github.com/coze-dev/cozeloop-go/internal/consts" "github.com/coze-dev/cozeloop-go/internal/httpclient" "github.com/coze-dev/cozeloop-go/internal/logger" + "github.com/coze-dev/cozeloop-go/internal/span" "github.com/coze-dev/cozeloop-go/internal/util" "github.com/coze-dev/cozeloop-go/spec/tracespec" ) @@ -41,6 +42,7 @@ type StartSpanOptions struct { StartNewTrace bool Scene string WorkspaceID string + NoSample bool // whether to disable sampling for this span and its child spans } type loopSpanKey struct{} @@ -71,7 +73,16 @@ func (t *Provider) GetOpts() *Options { return t.opt } -func (t *Provider) StartSpan(ctx context.Context, name, spanType string, opts StartSpanOptions) (context.Context, *Span, error) { +func (t *Provider) StartSpan(ctx context.Context, name, spanType string, opts StartSpanOptions) (context.Context, span.Span, error) { + if opts.NoSample { + ctx = context.WithValue(ctx, loopSpanKey{}, DefaultNoopSpan) + return ctx, DefaultNoopSpan, nil + } + + if !t.shouldSample(ctx) { + return ctx, DefaultNoopSpan, nil + } + // 0. check param if name == "" { name = "unknown" @@ -125,6 +136,11 @@ func (t *Provider) GetSpanFromHeader(ctx context.Context, header map[string]stri return FromHeader(ctx, header) } +func (t *Provider) shouldSample(ctx context.Context) bool { + parentSpan := ctx.Value(loopSpanKey{}) + return parentSpan != DefaultNoopSpan +} + func (t *Provider) startSpan(ctx context.Context, spanName string, spanType string, options StartSpanOptions) *Span { // 1. pack base data // get parentID from opt first, or set it to "0". diff --git a/internal/trace/trace_test.go b/internal/trace/trace_test.go index f3e0db9..f1ea788 100644 --- a/internal/trace/trace_test.go +++ b/internal/trace/trace_test.go @@ -76,3 +76,40 @@ func Test_GetSpanFromHeader(t *testing.T) { So(actual, ShouldEqual, expectedSpan) }) } + +func Test_NoSample(t *testing.T) { + PatchConvey("Test NoSample spans", t, func() { + provider := &Provider{ + httpClient: &httpclient.Client{}, + opt: &Options{ + WorkspaceID: "workspace-id", + UltraLargeReport: true, + }, + } + ctx := context.Background() + + // 1. Create a normal span + normalCtx, normalSpan, err := provider.StartSpan(ctx, "normal-span", "test-type", StartSpanOptions{ + StartTime: time.Now(), + }) + So(err, ShouldBeNil) + So(normalSpan, ShouldNotEqual, DefaultNoopSpan) + So(normalSpan, ShouldNotBeNil) + + // 2. Create a NoSample span + noSampleCtx, noSampleSpan, err := provider.StartSpan(normalCtx, "nosample-span", "test-type", StartSpanOptions{ + StartTime: time.Now(), + NoSample: true, + }) + So(err, ShouldBeNil) + So(noSampleSpan, ShouldEqual, DefaultNoopSpan) + + // 3. Create a child span from NoSample span - should also be NoopSpan + childCtx, childSpan, err := provider.StartSpan(noSampleCtx, "child-span", "test-type", StartSpanOptions{ + StartTime: time.Now(), + }) + So(err, ShouldBeNil) + So(childSpan, ShouldEqual, DefaultNoopSpan) + So(childCtx, ShouldNotBeNil) + }) +} diff --git a/span.go b/span.go index f9a4044..db01c93 100644 --- a/span.go +++ b/span.go @@ -3,138 +3,10 @@ package cozeloop -import ( - "context" - "time" - - "github.com/coze-dev/cozeloop-go/entity" - "github.com/coze-dev/cozeloop-go/spec/tracespec" -) +import "github.com/coze-dev/cozeloop-go/internal/span" // Span is the interface for span. -type Span interface { - SpanContext - commonSpanSetter - - // SetTags sets business custom tags. - SetTags(ctx context.Context, tagKVs map[string]interface{}) - - // SetBaggage sets tags and also passes these tags to other downstream spans (assuming - // the user uses ToHeader and FromHeader to handle header passing between services). - SetBaggage(ctx context.Context, baggageItems map[string]string) - - // Finish The span will be reported only after an explicit call to Finish. - // Under the hood, it is actually placed in an asynchronous queue waiting to be reported. - Finish(ctx context.Context) - - // GetStartTime returns the start time of the Span. - GetStartTime() time.Time - - // ToHeader Convert the span to headers. Used for cross-process correlation. - ToHeader() (map[string]string, error) -} - -// Set system-defined fields -type commonSpanSetter interface { - // SetInput key: `input` - // Input information. The input will be serialized into a JSON string. - // You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec - // Or you can use any struct you like. - SetInput(ctx context.Context, input interface{}) - - // SetOutput key: `output` - // Output information. The output will be serialized into a JSON string. - // You can find recommended specification in https://github.com/coze-dev/cozeloop-go/tree/main/spec/tracespec - // Or you can use any struct you like. - SetOutput(ctx context.Context, output interface{}) - - // SetError key: `error` - // Set error message. - SetError(ctx context.Context, err error) - - // SetStatusCode key: `status_code` - // Set status code. A non-zero code is considered an exception. - SetStatusCode(ctx context.Context, code int) - - // SetUserID key: `user_id` - // Set user id. - SetUserID(ctx context.Context, userID string) - SetUserIDBaggage(ctx context.Context, userID string) - - // SetMessageID key: `message_id` - // Set message id. - SetMessageID(ctx context.Context, messageID string) - SetMessageIDBaggage(ctx context.Context, messageID string) - - // SetThreadID key: `thread_id` - // Set thread id. - // thread_id is used to correlate multiple requests and is passed in by the business. - SetThreadID(ctx context.Context, threadID string) - SetThreadIDBaggage(ctx context.Context, threadID string) - - // SetPrompt key: `prompt - // Associated with PromptKey and PromptVersion, it will write two tags: prompt_key and prompt_version. - // SetPrompt is used to set the PromptKey and PromptVersion to tag. - SetPrompt(ctx context.Context, prompt entity.Prompt) - - // SetModelProvider key: `model_provider` - // The provider of the LLM, such as OpenAI, etc. - SetModelProvider(ctx context.Context, modelProvider string) - - // SetModelName key: `model_name` - // The Name of the LLM model, such as: gpt-4-1106-preview. - SetModelName(ctx context.Context, modelName string) - - // SetModelCallOptions key: `call_options` - // The call options of the LLM, such as temperature, max_tokens, etc. - // The recommended standard format is CallOption of spec package - SetModelCallOptions(ctx context.Context, callOptions interface{}) - - // SetInputTokens key: `input_tokens` - // The usage of input tokens. When the value of input_tokens is set, - // It will be automatically summed with output_tokens to calculate the tokens tag. - SetInputTokens(ctx context.Context, inputTokens int) - - // SetOutputTokens key: `output_tokens` - // The usage of output tokens. When the value of output_tokens is set, - // It will be automatically summed with input_tokens to calculate the tokens tag. - SetOutputTokens(ctx context.Context, outputTokens int) - - // SetStartTimeFirstResp key: `start_time_first_resp` - // Timestamp of the first packet return from LLM, unit: microseconds. - // When `start_time_first_resp` is set, a tag named `latency_first_resp` calculated - // based on the span's StartTime will be added, meaning the latency for the first packet. - SetStartTimeFirstResp(ctx context.Context, startTimeFirstResp int64) - - // SetRuntime key: `runtime` - // The runtime of the LLM, such as language, library, scene, etc. - // The recommended standard format is Runtime of spec package - SetRuntime(ctx context.Context, runtime tracespec.Runtime) - - // SetServiceName - // set the custom service name, identify different services. - SetServiceName(ctx context.Context, serviceName string) - - // SetLogID - // set the custom log id, identify different query. - SetLogID(ctx context.Context, logID string) - - // SetFinishTime - // Default is time.Now() when span Finish(). DO NOT set unless you do not use default time. - SetFinishTime(finishTime time.Time) - - // SetSystemTags - // set the system tags. DO NOT set unless you know what you are doing. - SetSystemTags(ctx context.Context, systemTags map[string]interface{}) - - // SetDeploymentEnv - // set the deployment env, identify custom env. - SetDeploymentEnv(ctx context.Context, deploymentEnv string) -} +type Span = span.Span // SpanContext is the interface for span Baggage transfer. -type SpanContext interface { - GetSpanID() string - GetTraceID() string - GetBaggage() map[string]string -} +type SpanContext = span.SpanContext diff --git a/trace.go b/trace.go index f948896..acf9f2a 100644 --- a/trace.go +++ b/trace.go @@ -81,3 +81,11 @@ func WithSpanID(spanID string) StartSpanOption { ops.SpanID = spanID } } + +// WithNoSample Set the span and its child span to not be sampled. +// This field is optional. If not specified, the sampling decision will be made by the parent span. +func WithNoSample() StartSpanOption { + return func(ops *startSpanOptions) { + ops.NoSample = true + } +}