From 34085d6e0bffec2bf273905d6432568298d37935 Mon Sep 17 00:00:00 2001 From: TimAndy Date: Thu, 29 Jan 2026 14:46:13 +0800 Subject: [PATCH 1/5] fix test cases --- internal/trace/exporter_test.go | 3 ++- internal/trace/queue_manager_test.go | 2 +- internal/trace/span_test.go | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) 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/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_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{}), } From 85bf3847c487d06235b01370dc01a07befbec689 Mon Sep 17 00:00:00 2001 From: TimAndy Date: Thu, 29 Jan 2026 14:46:15 +0800 Subject: [PATCH 2/5] move Span interface to package span --- span.go => internal/span/span.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) rename span.go => internal/span/span.go (99%) diff --git a/span.go b/internal/span/span.go similarity index 99% rename from span.go rename to internal/span/span.go index f9a4044..6563a00 100644 --- a/span.go +++ b/internal/span/span.go @@ -1,7 +1,7 @@ // Copyright (c) 2025 Bytedance Ltd. and/or its affiliates // SPDX-License-Identifier: MIT -package cozeloop +package span import ( "context" From bf8c7f5b82bb27144efad256dece497fc8ed1056 Mon Sep 17 00:00:00 2001 From: TimAndy Date: Thu, 29 Jan 2026 14:46:18 +0800 Subject: [PATCH 3/5] add Span and SpanContext interfaces for span management --- span.go | 12 ++++++++++++ 1 file changed, 12 insertions(+) create mode 100644 span.go diff --git a/span.go b/span.go new file mode 100644 index 0000000..db01c93 --- /dev/null +++ b/span.go @@ -0,0 +1,12 @@ +// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates +// SPDX-License-Identifier: MIT + +package cozeloop + +import "github.com/coze-dev/cozeloop-go/internal/span" + +// Span is the interface for span. +type Span = span.Span + +// SpanContext is the interface for span Baggage transfer. +type SpanContext = span.SpanContext From 6f66e959c3f04220c246aad10bff2e70abc3cf40 Mon Sep 17 00:00:00 2001 From: TimAndy Date: Thu, 29 Jan 2026 14:46:20 +0800 Subject: [PATCH 4/5] implement span interface compliance for noopSpan and SpanContext --- internal/trace/noop_span.go | 3 +++ internal/trace/span.go | 5 +++++ 2 files changed, 8 insertions(+) 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/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 From 9d351571b99e4306296e8ba556fa1e3e4245716c Mon Sep 17 00:00:00 2001 From: TimAndy Date: Thu, 29 Jan 2026 14:46:22 +0800 Subject: [PATCH 5/5] add NoSample option for span creation and update related tests --- client_test.go | 26 ++++++++++++++++++++++++ internal/prompt/prompt_hub.go | 5 +++-- internal/trace/trace.go | 18 ++++++++++++++++- internal/trace/trace_test.go | 37 +++++++++++++++++++++++++++++++++++ trace.go | 8 ++++++++ 5 files changed, 91 insertions(+), 3 deletions(-) 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/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/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 + } +}