AddlambdaToContext for a new background context while the original on…

…e may terminate.

Co-authored-by: zomborsz <balazs.zomborszki@n***.com>
Co-authored-by: Som-Som-CC <84196538+Som-Som-CC@users.noreply.github.com>

ctx.go

@@ -29,3 +29,9 @@ func NewRequestCtx(w http.ResponseWriter, r *http.Request) context.Context {
 func L(ctx context.Context) *lambda.Lambda {
 	return lambda.L(ctx)
 }
+
+// AddLambdaToContext returns the context extended with value of the Lambda.
+// Useful if a new goroutine is created with new context, but it is related to the original Context.
+func AddLambdaToContext(ctx context.Context, l *lambda.Lambda) context.Context {
+	return lambda.AddLambdaToContext(ctx, l)
+}

lambda/ctx.go

@@ -16,6 +16,7 @@ import (
 
 	"github.com/gorilla/mux"
 	"github.com/nokia/restful/trace/tracedata"
+	"github.com/nokia/restful/trace/traceotel"
 	"github.com/nokia/restful/trace/tracer"
 )
 
@@ -142,3 +143,12 @@ func (l *Lambda) ResponseHeaderAddAs(header, value string) {
 func (l *Lambda) TraceID() string {
 	return l.Trace.TraceID()
 }
+
+// AddLambdaToContext will return the context with value of Lambda
+func AddLambdaToContext(parentCtx context.Context, l *Lambda) context.Context {
+	ctx := context.WithValue(parentCtx, ctxName, l)
+	if tracer.GetOTel() {
+		ctx, _ = traceotel.TraceHeadersToContext(ctx, l.r)
+	}
+	return ctx
+}

lambda_test.go

@@ -175,6 +175,8 @@ func TestContext(t *testing.T) {
 		assert.Equal("42", L(ctx).ResponseHeader().Get("Request-Path-Id"))
 		assert.Equal("ss", resp.S)
 		assert.Equal(4, resp.I)
+		otherCtx := AddLambdaToContext(context.Background(), L(ctx))
+		assert.Equal(L(ctx), L(otherCtx))
 	}
 }
 

trace/traceotel/traceotel.go

@@ -32,8 +32,15 @@ type TraceOTel struct {
 //
 // Warning: Does not return trace from request context.
 func NewFromRequest(r *http.Request) *TraceOTel {
-	ctx, spanCtx := traceHeadersToContext(r)
-	if spanCtx.IsValid() {
+	return NewFromRequestWithContext(r.Context(), r)
+}
+
+// NewFromRequestWithContext creates new TraceOTel object derived from parentCtx. Returns nil if tracer not found in request.
+//
+// Warning: Does not return trace from request context.
+func NewFromRequestWithContext(parentCtx context.Context, r *http.Request) *TraceOTel {
+	ctx, ok := TraceHeadersToContext(parentCtx, r)
+	if ok {
 		return &TraceOTel{ctx: ctx}
 	}
 	return nil
@@ -56,13 +63,17 @@ func NewRandom() *TraceOTel {
 	return &TraceOTel{ctx: ctx}
 }
 
-// traceHeadersToContext maps trace headers in request to context.
-// Currently spanContext fails to include debug flag, but propagator sets that to ctx.
-func traceHeadersToContext(r *http.Request) (context.Context, trace.SpanContext) {
+// TraceHeadersToContext maps trace headers in request to context.
+// If there were no tracing headers to be propagated, the original context is returned.
+// The returned bool indicates if the original and the new contexts are the same.
+func TraceHeadersToContext(parentCtx context.Context, r *http.Request) (context.Context, bool) {
 	prop := b3.New()
-	ctx := prop.Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+	ctx := prop.Extract(parentCtx, propagation.HeaderCarrier(r.Header))
 	spanCtx := trace.SpanContextFromContext(ctx)
-	return ctx, spanCtx
+	if spanCtx.IsValid() {
+		return ctx, true
+	}
+	return parentCtx, false
 }
 
 // Span spans the existing trace data and puts that into the request.
@@ -83,15 +94,15 @@ func (t *TraceOTel) Span(r *http.Request) (*http.Request, string) {
 		r = r.WithContext(ctx)
 	} else {
 		// Check if req has tracing headers
-		var newCtx context.Context
-		newCtx, spanCtx = traceHeadersToContext(r)
-		if !spanCtx.IsValid() {
+		newCtx, ok := TraceHeadersToContext(r.Context(), r)
+		if !ok {
 			newCtx = NewRandom().ctx
 		}
+		spanCtx = trace.SpanContextFromContext(newCtx)
 		r = r.WithContext(newCtx)
 	}
 
-	return r, spanCtx.TraceID().String() // Client transport overrides spanID created here, hence not to be logged.
+	return r, spanCtx.TraceID().String()
 }
 
 // SetHeader sets request headers according to the trace data.

trace/tracer/tracer.go

@@ -96,6 +96,25 @@ func NewFromRequest(r *http.Request) *Tracer {
 	return &t
 }
 
+// NewFromRequestWithContext creates new tracer object from request derived from parentCtx if otel. Returns nil if not found.
+func NewFromRequestWithContext(parentCtx context.Context, r *http.Request) *Tracer {
+	var traceData tracedata.TraceData
+	if otelEnabled {
+		traceData = traceotel.NewFromRequestWithContext(parentCtx, r)
+	} else {
+		traceData = traceb3.NewFromRequest(r)
+		if reflect.ValueOf(traceData).IsNil() {
+			traceData = traceparent.NewFromRequest(r)
+		}
+	}
+
+	if traceData == nil || reflect.ValueOf(traceData).IsNil() {
+		return nil
+	}
+	t := Tracer{traceData: traceData, received: true}
+	return &t
+}
+
 // NewFromRequestOrRandom creates new tracer object. If no trace data, then create random. Never returns nil.
 //
 // Warning: Does not return trace from request context.