start new transaction support (#3)

coralogix · Nov 14, 2024 · b12cb0c · b12cb0c
1 parent 1e9f5ea
commit b12cb0c
Show file tree

Hide file tree

Showing 4 changed files with 215 additions and 1 deletion.
diff --git a/go.mod b/go.mod
@@ -18,3 +18,4 @@ require (
 	golang.org/x/sys v0.0.0-20220919091848-fb04ddd9f9c8 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 )
+
diff --git a/sampler/coralogix-integration-sampler_test.go b/sampler/coralogix-integration-sampler_test.go
@@ -0,0 +1,75 @@
+package sampler
+
+import (
+	"github.com/stretchr/testify/assert"
+	traceSdk "go.opentelemetry.io/otel/sdk/trace"
+	traceCore "go.opentelemetry.io/otel/trace"
+	"testing"
+)
+
+import (
+	"context"
+)
+
+const (
+	spanName = "spanName"
+)
+
+func TestCoralogixSampler_ShouldSample_flow(t *testing.T) {
+	t.Run("When_alwaysSampler_Should_AppendAttributesAndState", func(t *testing.T) {
+		alwaysSampler := traceSdk.AlwaysSample()
+		coralogixSampler := NewCoralogixSampler(alwaysSampler)
+
+		// Act
+
+		tracer := traceSdk.NewTracerProvider(traceSdk.WithSampler(coralogixSampler)).Tracer("test")
+		_, span := tracer.Start(context.Background(), "parent")
+		testAttribute(t, span, "parent")
+		span.End()
+
+	})
+
+	t.Run("When_alwaysSampler_Should_AppendAttributesAndState", func(t *testing.T) {
+		alwaysSampler := traceSdk.AlwaysSample()
+		coralogixSampler := NewCoralogixSampler(alwaysSampler)
+
+		// Act
+
+		tracer := traceSdk.NewTracerProvider(traceSdk.WithSampler(coralogixSampler)).Tracer("test")
+		ctx, span := tracer.Start(context.Background(), "parent")
+		ctxFlow1, spanFlow1 := tracer.Start(ctx, "flow1")
+		StartNewTransaction(spanFlow1, "flow1")
+		_, spanSubFlow1 := tracer.Start(ctxFlow1, "subFlow1")
+
+		testAttribute(t, spanSubFlow1, "flow1")
+		ctwFlo2, spanFlow2 := tracer.Start(ctx, "flow2")
+		StartNewTransaction(spanFlow2, "flow2")
+		_, spanSubFlow2 := tracer.Start(ctwFlo2, "subFlow2")
+		testAttribute(t, spanSubFlow2, "flow2")
+
+		attributes := span.(traceSdk.ReadWriteSpan).Attributes()
+		found := false
+		for _, attribute := range attributes {
+			if attribute.Key == "cgx.transaction" {
+				assert.Equal(t, attribute.Value.AsString(), "parent")
+				found = true
+			}
+		}
+		assert.True(t, found)
+		span.End()
+
+	})
+
+}
+
+func testAttribute(t *testing.T, span traceCore.Span, attributeVal string) {
+	attributes := span.(traceSdk.ReadWriteSpan).Attributes()
+	found := false
+	for _, attribute := range attributes {
+		if attribute.Key == "cgx.transaction" {
+			assert.Equal(t, attributeVal, attribute.Value.AsString())
+			found = true
+		}
+	}
+	assert.True(t, found)
+}
diff --git a/sampler/coralogix-sampler.go b/sampler/coralogix-sampler.go
@@ -70,6 +70,26 @@ func (s *CoralogixSampler) generateNewTraceState(ctx context.Context, name strin
 	parentTraceState := samplingResult.Tracestate
 
 	if !parentSpanContext.IsRemote() && parentTraceState.Get(TransactionIdentifierTraceState) != "" {
+		span := traceCore.SpanFromContext(ctx)
+		if span != nil {
+			readWriteSpan, ok := span.(traceSdk.ReadWriteSpan)
+			if ok {
+				attributes := readWriteSpan.Attributes()
+				if attributes != nil {
+					for _, attribute := range attributes {
+						if attribute.Key == TransactionIdentifier {
+							parentTraceState, err := parentTraceState.Insert(TransactionIdentifierTraceState, attribute.Value.AsString())
+							if err == nil {
+								return parentTraceState
+							}
+						}
+					}
+				}
+
+			}
+		}
+
+		/**/
 		return parentTraceState
 	}
 
@@ -94,3 +114,8 @@ func (s *CoralogixSampler) getParentSpanContext(ctx context.Context) traceCore.S
 	}
 	return traceCore.SpanContext{}
 }
+func StartNewTransaction(span traceCore.Span, flow string) traceCore.Span {
+	span.SetAttributes(attribute.String(TransactionIdentifier, flow))
+	span.SetAttributes(attribute.Bool(TransactionIdentifierRoot, true))
+	return span
+}
diff --git a/sampler/coralogix-sampler_test.go b/sampler/coralogix-sampler_test.go
@@ -13,7 +13,7 @@ import (
 )
 
 const (
-	spanName = "spanName"
+	spanNames = "spanName"
 )
 
 func TestCoralogixSampler_ShouldSample(t *testing.T) {
@@ -79,6 +79,43 @@ func TestCoralogixSampler_ShouldSample(t *testing.T) {
 		})
 	})
 
+	t.Run("When_ParentContextExistsAndNotRemote_ShouldCopyParentTraceState", func(t *testing.T) {
+		alwaysSampler := traceSdk.AlwaysSample()
+		coralogixSampler := NewCoralogixSampler(alwaysSampler)
+
+		traceState := traceCore.TraceState{}
+
+		traceState, _ = traceState.Insert(TransactionIdentifierTraceState, "fatherSpanName")
+		traceState, _ = traceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		parentSpan := traceCore.NewSpanContext(traceCore.SpanContextConfig{
+			TraceID:    traceCore.TraceID{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F},
+			SpanID:     traceCore.SpanID{0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8},
+			TraceFlags: traceCore.FlagsSampled,
+			TraceState: traceState,
+			Remote:     false,
+		})
+		parentCtx := traceCore.ContextWithSpanContext(context.Background(), parentSpan)
+
+		// Act
+		parameters := traceSdk.SamplingParameters{
+			ParentContext: parentCtx,
+			Name:          spanName,
+			Attributes:    []attribute.KeyValue{},
+		}
+		result := coralogixSampler.ShouldSample(parameters)
+
+		expectedAttributes := []attribute.KeyValue{
+			attribute.String(TransactionIdentifier, "fatherSpanName"),
+			attribute.String(DistributedTransactionIdentifier, "fatherSpanName"),
+		}
+		expectedTraceState := traceCore.TraceState{}
+		expectedTraceState, _ = traceState.Insert(TransactionIdentifierTraceState, "fatherSpanName")
+		expectedTraceState, _ = traceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		assert.ElementsMatch(t, expectedAttributes, result.Attributes)
+		assert.Equal(t, expectedTraceState, result.Tracestate)
+	})
 	t.Run("When_ParentContextExistsAndNotRemote_ShouldCopyParentTraceState", func(t *testing.T) {
 		alwaysSampler := traceSdk.AlwaysSample()
 		coralogixSampler := NewCoralogixSampler(alwaysSampler)
@@ -195,4 +232,80 @@ func TestCoralogixSampler_ShouldSample(t *testing.T) {
 		assert.Equal(t, expectedTraceState.Get(TransactionIdentifierTraceState), result.Tracestate.Get(TransactionIdentifierTraceState))
 		assert.Equal(t, expectedTraceState.Get(DistributedTransactionIdentifierTraceState), result.Tracestate.Get(DistributedTransactionIdentifierTraceState))
 	})
+	t.Run("When_ParentContextExistsAndRemote_ShouldCopyParentTraceState", func(t *testing.T) {
+		alwaysSampler := traceSdk.AlwaysSample()
+		coralogixSampler := NewCoralogixSampler(alwaysSampler)
+
+		traceState := traceCore.TraceState{}
+
+		traceState, _ = traceState.Insert(TransactionIdentifierTraceState, "fatherSpanName")
+		traceState, _ = traceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		parentSpan := traceCore.NewSpanContext(traceCore.SpanContextConfig{
+			TraceID:    traceCore.TraceID{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F},
+			SpanID:     traceCore.SpanID{0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8},
+			TraceFlags: traceCore.FlagsSampled,
+			TraceState: traceState,
+			Remote:     true,
+		})
+		parentCtx := traceCore.ContextWithSpanContext(context.Background(), parentSpan)
+
+		// Act
+		parameters := traceSdk.SamplingParameters{
+			ParentContext: parentCtx,
+			Name:          spanName,
+			Attributes:    []attribute.KeyValue{},
+		}
+		result := coralogixSampler.ShouldSample(parameters)
+
+		expectedAttributes := []attribute.KeyValue{
+			attribute.String(TransactionIdentifier, "spanName"),
+			attribute.String(DistributedTransactionIdentifier, "fatherSpanName"),
+			attribute.Bool(TransactionIdentifierRoot, true),
+		}
+		expectedTraceState := traceCore.TraceState{}
+		expectedTraceState, _ = expectedTraceState.Insert(TransactionIdentifierTraceState, spanName)
+		expectedTraceState, _ = expectedTraceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		assert.ElementsMatch(t, expectedAttributes, result.Attributes)
+		assert.Equal(t, expectedTraceState.Get(TransactionIdentifierTraceState), result.Tracestate.Get(TransactionIdentifierTraceState))
+		assert.Equal(t, expectedTraceState.Get(DistributedTransactionIdentifierTraceState), result.Tracestate.Get(DistributedTransactionIdentifierTraceState))
+	})
+	t.Run("When_SetFlow_ParentContextExistsAndNotRemote_ShouldCopyParentTraceState", func(t *testing.T) {
+		alwaysSampler := traceSdk.AlwaysSample()
+		coralogixSampler := NewCoralogixSampler(alwaysSampler)
+
+		traceState := traceCore.TraceState{}
+
+		traceState, _ = traceState.Insert(TransactionIdentifierTraceState, "fatherSpanName")
+		traceState, _ = traceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		parentSpan := traceCore.NewSpanContext(traceCore.SpanContextConfig{
+			TraceID:    traceCore.TraceID{0x00, 0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09, 0x0A, 0x0B, 0x0C, 0x0D, 0x0E, 0x0F},
+			SpanID:     traceCore.SpanID{0xFF, 0xFE, 0xFD, 0xFC, 0xFB, 0xFA, 0xF9, 0xF8},
+			TraceFlags: traceCore.FlagsSampled,
+			TraceState: traceState,
+			Remote:     false,
+		})
+		parentCtx := traceCore.ContextWithSpanContext(context.Background(), parentSpan)
+
+		// Act
+		parameters := traceSdk.SamplingParameters{
+			ParentContext: parentCtx,
+			Name:          spanName,
+			Attributes:    []attribute.KeyValue{},
+		}
+		result := coralogixSampler.ShouldSample(parameters)
+
+		expectedAttributes := []attribute.KeyValue{
+			attribute.String(TransactionIdentifier, "fatherSpanName"),
+			attribute.String(DistributedTransactionIdentifier, "fatherSpanName"),
+		}
+		expectedTraceState := traceCore.TraceState{}
+		expectedTraceState, _ = traceState.Insert(TransactionIdentifierTraceState, "fatherSpanName")
+		expectedTraceState, _ = traceState.Insert(DistributedTransactionIdentifierTraceState, "fatherSpanName")
+
+		assert.ElementsMatch(t, expectedAttributes, result.Attributes)
+		assert.Equal(t, expectedTraceState, result.Tracestate)
+	})
 }