Support correlation (#91)

README.md

@@ -42,6 +42,11 @@ You can also create tracer with sampling rate.
 tracer, err := go2sky.NewTracer("example", go2sky.WithReporter(r), go2sky.WithSampler(0.5))
 ```
 
+Also could customize correlation context config.
+```go
+tracer, err := go2sky.NewTracer("example", go2sky.WithReporter(r), go2sky.WithSampler(0.5), go2sky.WithCorrelation(3, 128))
+```
+
 ## Create span
 
 To create a span in a trace, we used the `Tracer` to start a new span. We indicate this as the root span because of 
@@ -67,6 +72,22 @@ A sub span created as the children of root span links to its parent with `Contex
 subSpan, newCtx, err := tracer.CreateLocalSpan(ctx)
 ```
 
+## Get correlation
+
+Get custom data from tracing context.
+
+```go
+value := go2sky.GetCorrelation(ctx, key)
+```
+
+## Put correlation
+
+Put custom data to tracing context.
+
+```go
+success := go2sky.PutCorrelation(ctx, key, value)
+```
+
 ## End span
 
 We must end the spans so they becomes available for sending to the backend by a reporter.
@@ -109,16 +130,16 @@ upstream service.
 
 ```go
 //Extract context from HTTP request header `sw8`
-span, ctx, err := tracer.CreateEntrySpan(r.Context(), "/api/login", func() (string, error) {
-		return r.Header.Get("sw8"), nil
+span, ctx, err := tracer.CreateEntrySpan(r.Context(), "/api/login", func(key string) (string, error) {
+		return r.Header.Get(key), nil
 })
 
 // Some operation
 ...
 
 // Inject context into HTTP request header `sw8`
-span, err := tracer.CreateExitSpan(req.Context(), "/service/validate", "tomcat-service:8080", func(header string) error {
-		req.Header.Set(propagation.Header, header)
+span, err := tracer.CreateExitSpan(req.Context(), "/service/validate", "tomcat-service:8080", func(key, value string) error {
+		req.Header.Set(key, value)
 		return nil
 })
 ```

correlation.go

@@ -0,0 +1,84 @@
+// Licensed to SkyAPM org under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. SkyAPM org licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package go2sky
+
+import "context"
+
+type CorrelationConfig struct {
+	MaxKeyCount  int
+	MaxValueSize int
+}
+
+func WithCorrelation(keyCount, valueSize int) TracerOption {
+	return func(t *Tracer) {
+		t.correlation = &CorrelationConfig{
+			MaxKeyCount:  keyCount,
+			MaxValueSize: valueSize,
+		}
+	}
+}
+
+func PutCorrelation(ctx context.Context, key, value string) bool {
+	if key == "" {
+		return false
+	}
+
+	activeSpan := ctx.Value(ctxKeyInstance)
+	if activeSpan == nil {
+		return false
+	}
+
+	span, ok := activeSpan.(segmentSpan)
+	if !ok {
+		return false
+	}
+	correlationContext := span.context().CorrelationContext
+	// remove key
+	if value == "" {
+		delete(correlationContext, key)
+		return true
+	}
+	// out of max value size
+	if len(value) > span.tracer().correlation.MaxValueSize {
+		return false
+	}
+	// already exists key
+	if _, ok := correlationContext[key]; ok {
+		correlationContext[key] = value
+		return true
+	}
+	// out of max key count
+	if len(correlationContext) >= span.tracer().correlation.MaxKeyCount {
+		return false
+	}
+	span.context().CorrelationContext[key] = value
+	return true
+}
+
+func GetCorrelation(ctx context.Context, key string) string {
+	activeSpan := ctx.Value(ctxKeyInstance)
+	if activeSpan == nil {
+		return ""
+	}
+
+	span, ok := activeSpan.(segmentSpan)
+	if !ok {
+		return ""
+	}
+	return span.context().CorrelationContext[key]
+}

correlation_test.go

@@ -0,0 +1,206 @@
+// Licensed to SkyAPM org under one or more contributor
+// license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright
+// ownership. SkyAPM org licenses this file to you under
+// the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package go2sky_test
+
+import (
+	"context"
+	"log"
+	"reflect"
+	"testing"
+
+	"github.com/SkyAPM/go2sky"
+	"github.com/SkyAPM/go2sky/propagation"
+	"github.com/SkyAPM/go2sky/reporter"
+)
+
+const (
+	correlationTestKey   = "test-key"
+	correlationTestValue = "test-value"
+)
+
+func TestGetCorrelation_WithTracingContest(t *testing.T) {
+	verifyPutResult := func(ctx context.Context, key, value string, result bool, t *testing.T) {
+		if success := go2sky.PutCorrelation(ctx, key, value); success != result {
+			t.Errorf("put correlation result is not right: %t", success)
+		}
+	}
+	tests := []struct {
+		name string
+		// extract from context
+		extractor propagation.Extractor
+		// extract correlation context
+		extracted map[string]string
+		// put correlation
+		customCase func(ctx context.Context, t *testing.T)
+		// after exported correaltion context
+		want map[string]string
+	}{
+		{
+			name: "no context",
+			extractor: func(headerKey string) (string, error) {
+				return "", nil
+			},
+			extracted: make(map[string]string),
+			customCase: func(ctx context.Context, t *testing.T) {
+				verifyPutResult(ctx, correlationTestKey, correlationTestValue, true, t)
+			},
+			want: func() map[string]string {
+				m := make(map[string]string)
+				m[correlationTestKey] = correlationTestValue
+				return m
+			}(),
+		},
+		{
+			name: "existing context with correlation",
+			extractor: func(headerKey string) (string, error) {
+				if headerKey == propagation.HeaderCorrelation {
+					// test1 = t1
+					return "dGVzdDE=:dDE=", nil
+				}
+				if headerKey == propagation.Header {
+					return "1-MWYyZDRiZjQ3YmY3MTFlYWI3OTRhY2RlNDgwMDExMjI=-MWU3YzIwNGE3YmY3MTFlYWI4NThhY2RlNDgwMDExMjI=" +
+						"-0-c2VydmljZQ==-aW5zdGFuY2U=-cHJvcGFnYXRpb24=-cHJvcGFnYXRpb246NTU2Ng==", nil
+				}
+				return "", nil
+			},
+			extracted: func() map[string]string {
+				m := make(map[string]string)
+				m["test1"] = "t1"
+				return m
+			}(),
+			customCase: func(ctx context.Context, t *testing.T) {
+				verifyPutResult(ctx, correlationTestKey, correlationTestValue, true, t)
+			},
+			want: func() map[string]string {
+				m := make(map[string]string)
+				m[correlationTestKey] = correlationTestValue
+				m["test1"] = "t1"
+				return m
+			}(),
+		},
+		{
+			name: "empty context with put bound judge",
+			extractor: func(headerKey string) (string, error) {
+				return "", nil
+			},
+			customCase: func(ctx context.Context, t *testing.T) {
+				// empty key
+				verifyPutResult(ctx, "", "123", false, t)
+
+				// remove key
+				verifyPutResult(ctx, correlationTestKey, correlationTestValue, true, t)
+				verifyPutResult(ctx, correlationTestKey, "", true, t)
+				if go2sky.GetCorrelation(ctx, correlationTestKey) != "" {
+					t.Errorf("correlation test key should be null")
+				}
+
+				// out of max value size
+				verifyPutResult(ctx, "test-key", "1234567890123456", false, t)
+
+				// out of key count
+				verifyPutResult(ctx, "test-key1", "123", true, t)
+				verifyPutResult(ctx, "test-key2", "123", true, t)
+				verifyPutResult(ctx, "test-key3", "123", true, t)
+				verifyPutResult(ctx, "test-key4", "123", false, t)
+
+				// exists key
+				verifyPutResult(ctx, "test-key1", "123456", true, t)
+			},
+			want: func() map[string]string {
+				m := make(map[string]string)
+				m["test-key1"] = "123456"
+				m["test-key2"] = "123"
+				m["test-key3"] = "123"
+				return m
+			}(),
+		},
+	}
+
+	r, err := reporter.NewLogReporter()
+	if err != nil {
+		log.Fatalf("new reporter error %v \n", err)
+	}
+	defer r.Close()
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := context.Background()
+			tracer, _ := go2sky.NewTracer("correlationTest", go2sky.WithReporter(r), go2sky.WithSampler(1), go2sky.WithCorrelation(3, 10))
+
+			// create entry span from extractor
+			span, ctx, _ := tracer.CreateEntrySpan(ctx, "test-entry", tt.extractor)
+			defer span.End()
+
+			// verify extracted context is same
+			if tt.extracted != nil {
+				for key, value := range tt.extracted {
+					if go2sky.GetCorrelation(ctx, key) != value {
+						t.Errorf("error get previous correlation value, current is: %s", go2sky.GetCorrelation(ctx, key))
+					}
+				}
+			}
+
+			// custom case
+			tt.customCase(ctx, t)
+
+			// put sample local span
+			span, ctx, _ = tracer.CreateLocalSpan(ctx)
+			defer span.End()
+
+			// validate correlation context
+			// verify extracted context is same
+			for key, value := range tt.want {
+				if go2sky.GetCorrelation(ctx, key) != value {
+					t.Errorf("error validate correlation value, current is: %s", go2sky.GetCorrelation(ctx, key))
+				}
+			}
+
+			// export context
+			scx := propagation.SpanContext{}
+			_, err := tracer.CreateExitSpan(ctx, "test-exit", "127.0.0.1:8080", func(headerKey, headerValue string) error {
+				if headerKey == propagation.HeaderCorrelation {
+					err = scx.DecodeSW8Correlation(headerValue)
+					if err != nil {
+						t.Fail()
+					}
+				}
+				return nil
+			})
+			if err != nil {
+				t.Fail()
+			}
+			reflect.DeepEqual(scx, tt.want)
+		})
+	}
+}
+
+func TestGetCorrelation_WithEmptyContext(t *testing.T) {
+	emptyValue := go2sky.GetCorrelation(context.Background(), "empty-key")
+	if emptyValue != "" {
+		t.Errorf("should be empty value")
+	}
+
+	success := go2sky.PutCorrelation(context.Background(), "empty-key", "empty-value")
+	if success {
+		t.Errorf("put correlation key should be failed")
+	}
+
+	emptyValue = go2sky.GetCorrelation(context.Background(), "empty-key")
+	if emptyValue != "" {
+		t.Errorf("should be empty value")
+	}
+}

noop_test.go

@@ -36,15 +36,15 @@ func TestCreateNoopSpan(t *testing.T) {
 		{
 			"Entry",
 			func() (Span, context.Context, error) {
-				return tracer.CreateEntrySpan(context.Background(), "entry", func() (s string, e error) {
+				return tracer.CreateEntrySpan(context.Background(), "entry", func(key string) (s string, e error) {
 					return "", nil
 				})
 			},
 		},
 		{
 			"Exit",
 			func() (s Span, c context.Context, err error) {
-				s, err = tracer.CreateExitSpan(context.Background(), "exit", "localhost:8080", func(header string) error {
+				s, err = tracer.CreateExitSpan(context.Background(), "exit", "localhost:8080", func(key, value string) error {
 					return nil
 				})
 				return
@@ -72,13 +72,13 @@ func TestCreateNoopSpan(t *testing.T) {
 
 func TestNoopSpanFromBegin(t *testing.T) {
 	tracer, _ := NewTracer("service")
-	span, ctx, _ := tracer.CreateEntrySpan(context.Background(), "entry", func() (s string, e error) {
+	span, ctx, _ := tracer.CreateEntrySpan(context.Background(), "entry", func(key string) (s string, e error) {
 		return "", nil
 	})
 	if _, ok := span.(*NoopSpan); !ok {
 		t.Error("Should create noop span")
 	}
-	exitSpan, _ := tracer.CreateExitSpan(ctx, "exit", "localhost:8080", func(header string) error {
+	exitSpan, _ := tracer.CreateExitSpan(ctx, "exit", "localhost:8080", func(key, value string) error {
 		return nil
 	})
 	if _, ok := exitSpan.(*NoopSpan); !ok {