Support trace logger in workflows & activities (#714)

temporalio · Feb 7, 2022 · b5ba1d3 · b5ba1d3
1 parent 5773685
commit b5ba1d3
Showing 1 changed file with 20 additions and 1 deletion.
diff --git a/interceptor/tracing_interceptor.go b/interceptor/tracing_interceptor.go
@@ -27,6 +27,7 @@ import (
 	"fmt"
 
 	commonpb "go.temporal.io/api/common/v1"
+
 	"go.temporal.io/sdk/activity"
 	"go.temporal.io/sdk/client"
 	"go.temporal.io/sdk/converter"
@@ -290,11 +291,29 @@ func (t *tracingClientOutboundInterceptor) QueryWorkflow(
 	return val, err
 }
 
+type tracingActivityOutboundInterceptor struct {
+	ActivityOutboundInterceptorBase
+	root *tracingInterceptor
+}
+
+func (t *tracingActivityOutboundInterceptor) GetLogger(ctx context.Context) log.Logger {
+	if span := t.root.tracer.SpanFromContext(ctx); span != nil {
+		return t.root.tracer.GetLogger(t.Next.GetLogger(ctx), span)
+	}
+	return t.Next.GetLogger(ctx)
+}
+
 type tracingActivityInboundInterceptor struct {
 	ActivityInboundInterceptorBase
 	root *tracingInterceptor
 }
 
+func (t *tracingActivityInboundInterceptor) Init(outbound ActivityOutboundInterceptor) error {
+	i := &tracingActivityOutboundInterceptor{root: t.root}
+	i.Next = outbound
+	return t.Next.Init(i)
+}
+
 func (t *tracingActivityInboundInterceptor) ExecuteActivity(
 	ctx context.Context,
 	in *ExecuteActivityInput,
@@ -453,7 +472,7 @@ func (t *tracingWorkflowOutboundInterceptor) ExecuteLocalActivity(
 
 func (t *tracingWorkflowOutboundInterceptor) GetLogger(ctx workflow.Context) log.Logger {
 	if span, _ := ctx.Value(t.root.options.SpanContextKey).(TracerSpan); span != nil {
-		t.root.tracer.GetLogger(t.Next.GetLogger(ctx), span)
+		return t.root.tracer.GetLogger(t.Next.GetLogger(ctx), span)
 	}
 	return t.Next.GetLogger(ctx)
 }