mirror of
https://github.com/hanzoai/telemetry.git
synced 2026-10-30 20:01:49 +00:00
Sindy LiandGitHub
4a8a8c9d84
<!--Ex. Fixing a bug - Describe the bug and how this fixes the issue. Ex. Adding a feature - Explain what this achieves.--> #### Description This PR provides a solution to propagate context values across batching. This is particularly useful in scenarios where context carries important metadata—such as authentication tokens. **New API** A new API has been added to support this functionality: ``` mergeCtx func(ctx1 context.Context, ctx2 context.Context) context.Context ``` By supplying a custom mergeCtx function, users can control how context values are preserved or combined—for example, by selectively copying specific keys, merging metadata, or keeping the larger value. Note: `context.Context` returned from the above function should not override `Value()` function. **Default Behavior** If mergeCtx is not provided (i.e., left as nil), the batching system defaults to a conservative behavior: no context values are retained from the original incoming contexts. Users who need context propagation must explicitly opt in by defining and supplying a mergeCtx function. <!-- Issue number if applicable --> #### Link to tracking issue https://github.com/open-telemetry/opentelemetry-collector/issues/13320 <!--Describe what testing was performed and which tests were added.--> #### Testing * `batch_context_test.go` checks the `mergeCtx` does not interfere with span link * `partition_batcher_test.go` checks `mergeCtx=nil` works properly <!--Describe the documentation added.--> #### Documentation <!--Please delete paragraphs that you did not use before submitting.-->
73 lines
2.5 KiB
Go
73 lines
2.5 KiB
Go
// Copyright The OpenTelemetry Authors
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
package queuebatch
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/require"
|
|
"go.opentelemetry.io/otel/trace"
|
|
|
|
"go.opentelemetry.io/collector/component/componenttest"
|
|
)
|
|
|
|
type testTimestampKeyType int
|
|
|
|
const testTimestampKey testTimestampKeyType = iota
|
|
|
|
// mergeCtxFunc corresponds to user specified mergeCtx function in the batcher settings.
|
|
// This specific merge Context function keeps the greater of timestamps from two contexts.
|
|
func mergeCtxFunc(ctx1, ctx2 context.Context) context.Context {
|
|
timestamp1 := ctx1.Value(testTimestampKey)
|
|
timestamp2 := ctx2.Value(testTimestampKey)
|
|
if timestamp1 != nil && timestamp2 != nil {
|
|
if timestamp1.(int) > timestamp2.(int) {
|
|
return context.WithValue(context.Background(), testTimestampKey, timestamp1)
|
|
}
|
|
return context.WithValue(context.Background(), testTimestampKey, timestamp2)
|
|
}
|
|
if timestamp1 != nil {
|
|
return context.WithValue(context.Background(), testTimestampKey, timestamp1)
|
|
}
|
|
return context.WithValue(context.Background(), testTimestampKey, timestamp2)
|
|
}
|
|
|
|
// mergeContextHelper performs the same operation done during batching.
|
|
func mergeContextHelper(ctx1, ctx2 context.Context) context.Context {
|
|
return contextWithMergedLinks(mergeCtxFunc(ctx1, ctx2), ctx1, ctx2)
|
|
}
|
|
|
|
func TestBatchContextLink(t *testing.T) {
|
|
tracerProvider := componenttest.NewTelemetry().NewTelemetrySettings().TracerProvider
|
|
tracer := tracerProvider.Tracer("go.opentelemetry.io/collector/exporter/exporterhelper")
|
|
|
|
ctx1 := context.Background()
|
|
|
|
ctx2, span2 := tracer.Start(ctx1, "span2")
|
|
defer span2.End()
|
|
|
|
ctx3, span3 := tracer.Start(ctx1, "span3")
|
|
defer span3.End()
|
|
|
|
ctx4, span4 := tracer.Start(ctx1, "span4")
|
|
defer span4.End()
|
|
|
|
batchContext := mergeContextHelper(ctx2, ctx3)
|
|
batchContext = mergeContextHelper(batchContext, ctx4)
|
|
|
|
actualLinks := LinksFromContext(batchContext)
|
|
require.Len(t, actualLinks, 3)
|
|
require.Equal(t, trace.SpanContextFromContext(ctx2), actualLinks[0].SpanContext)
|
|
require.Equal(t, trace.SpanContextFromContext(ctx3), actualLinks[1].SpanContext)
|
|
require.Equal(t, trace.SpanContextFromContext(ctx4), actualLinks[2].SpanContext)
|
|
}
|
|
|
|
func TestMergedContext_GetValue(t *testing.T) {
|
|
ctx1 := context.WithValue(context.Background(), testTimestampKey, 1234)
|
|
ctx2 := context.WithValue(context.Background(), testTimestampKey, 2345)
|
|
batchContext := mergeContextHelper(ctx1, ctx2)
|
|
require.Equal(t, 2345, batchContext.Value(testTimestampKey))
|
|
}
|