add WithContext

Author: shijiesheng

internal/context.go

@@ -494,6 +494,6 @@ func GetSpanContext(ctx Context) opentracing.SpanContext {
 	return nil
 }
 
-func contextWithSpan(ctx Context, spanContext opentracing.SpanContext) Context {
+func WithSpanContext(ctx Context, spanContext opentracing.SpanContext) Context {
 	return WithValue(ctx, activeSpanContextKey, spanContext)
 }

internal/tracer.go

@@ -119,5 +119,5 @@ func (t *tracingContextPropagator) ExtractToWorkflow(
 		// did not find a tracing span, just return the current context
 		return ctx, nil
 	}
-	return contextWithSpan(ctx, spanContext), nil
+	return WithSpanContext(ctx, spanContext), nil
 }

internal/tracer_test.go

@@ -82,7 +82,7 @@ func TestTracingContextPropagatorWorkflowContext(t *testing.T) {
 
 	span := tracer.StartSpan("test-operation")
 	assert.NotNil(t, span.Context())
-	ctx := contextWithSpan(Background(), span.Context())
+	ctx := WithSpanContext(Background(), span.Context())
 	header := &shared.Header{
 		Fields: map[string][]byte{},
 	}
@@ -130,7 +130,7 @@ func TestConsistentInjectionExtraction(t *testing.T) {
 	var baggageVal = "e30="
 	span.SetBaggageItem("request-tenancy", baggageVal)
 	assert.NotNil(t, span.Context())
-	ctx := contextWithSpan(Background(), span.Context())
+	ctx := WithSpanContext(Background(), span.Context())
 	header := &shared.Header{
 		Fields: map[string][]byte{},
 	}

workflow/context.go

@@ -93,3 +93,26 @@ func NewDisconnectedContext(parent Context) (ctx Context, cancel CancelFunc) {
 func GetSpanContext(ctx Context) opentracing.SpanContext {
 	return internal.GetSpanContext(ctx)
 }
+
+// WithSpanContext returns [Context] with override [opentracing.SpanContext].
+// This is useful to modify baggage items of current workflow and pass it to activities and child workflows.
+//
+// Example Usage:
+//	func myWorkflow(ctx Context) error {
+// 		// start a new workflow span
+// 		wSpan := opentracing.StartSpan("workflow-operation", opentracing.ChildOf(spanContext))
+// 		// pass the new span context to activity
+// 		aCtx := WithSpanContext(ctx, wSpan.Context())
+// 		var activityFooResult string
+// 		err := ExecuteActivity(aCtx, ActivityFoo).Get(aCtx, &activityFooResult)
+// 		if err != nil {
+// 			wSpan.SetTag("workflow-error", err)
+// 		} else {
+// 			wSpan.SetTag("workflow-result", activityFooResult)
+// 		}
+// 		wSpan.Finish()
+// 		return activityFooResult, err
+// 	}
+func WithSpanContext(ctx Context, spanContext opentracing.SpanContext) Context {
+	return internal.WithSpanContext(ctx, spanContext)
+}