From 43eac5d3335278f45bbef7e59620fe6448b2ad46 Mon Sep 17 00:00:00 2001 From: Rey Abolofia Date: Fri, 26 Apr 2024 10:41:50 -0700 Subject: [PATCH] Set error type and stack with universal instrumentation (#171) * fix: set error type and stack with universal instrumentation * Tests for error type/stack headers. --------- Co-authored-by: dali546 <35352237+dali546@users.noreply.github.com> --- ddlambda.go | 2 +- internal/extension/extension.go | 64 ++++++++++++++++++++++++---- internal/extension/extension_test.go | 49 ++++++++++++++++++++- internal/trace/listener.go | 12 +++--- 4 files changed, 110 insertions(+), 17 deletions(-) diff --git a/ddlambda.go b/ddlambda.go index 2cbb3561..c8ee2caf 100644 --- a/ddlambda.go +++ b/ddlambda.go @@ -221,7 +221,7 @@ func (cfg *Config) toTraceConfig() trace.Config { } if tracingEnabled, err := strconv.ParseBool(os.Getenv(DatadogTraceEnabledEnvVar)); err == nil { - traceConfig.DDTraceEnabled = tracingEnabled; + traceConfig.DDTraceEnabled = tracingEnabled // Only read the OTEL env var if DD tracing is disabled if tracingEnabled { if otelTracerEnabled, err := strconv.ParseBool(os.Getenv(OtelTracerEnabled)); err == nil { diff --git a/internal/extension/extension.go b/internal/extension/extension.go index a1b424a0..3f7179ce 100644 --- a/internal/extension/extension.go +++ b/internal/extension/extension.go @@ -11,10 +11,15 @@ package extension import ( "bytes" "context" + "encoding/base64" "encoding/json" "fmt" "net/http" "os" + "reflect" + "runtime" + "strconv" + "strings" "time" "github.com/DataDog/datadog-lambda-go/internal/logger" @@ -25,12 +30,14 @@ import ( type ddTraceContext string const ( - DdTraceId ddTraceContext = "x-datadog-trace-id" - DdParentId ddTraceContext = "x-datadog-parent-id" - DdSpanId ddTraceContext = "x-datadog-span-id" - DdSamplingPriority ddTraceContext = "x-datadog-sampling-priority" - DdInvocationError ddTraceContext = "x-datadog-invocation-error" - DdInvocationErrorMsg ddTraceContext = "x-datadog-invocation-error-msg" + DdTraceId ddTraceContext = "x-datadog-trace-id" + DdParentId ddTraceContext = "x-datadog-parent-id" + DdSpanId ddTraceContext = "x-datadog-span-id" + DdSamplingPriority ddTraceContext = "x-datadog-sampling-priority" + DdInvocationError ddTraceContext = "x-datadog-invocation-error" + DdInvocationErrorMsg ddTraceContext = "x-datadog-invocation-error-msg" + DdInvocationErrorType ddTraceContext = "x-datadog-invocation-error-type" + DdInvocationErrorStack ddTraceContext = "x-datadog-invocation-error-stack" DdSeverlessSpan ddTraceContext = "dd-tracer-serverless-span" DdLambdaResponse ddTraceContext = "dd-response" @@ -125,7 +132,7 @@ func (em *ExtensionManager) SendStartInvocationRequest(ctx context.Context, even return ctx } -func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functionExecutionSpan ddtrace.Span, err error) { +func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functionExecutionSpan ddtrace.Span, cfg ddtrace.FinishConfig) { // Handle Lambda response lambdaResponse := ctx.Value(DdLambdaResponse) content, responseErr := json.Marshal(lambdaResponse) @@ -136,9 +143,11 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi req, _ := http.NewRequest(http.MethodPost, em.endInvocationUrl, body) // Mark the invocation as an error if any - if err != nil { + if cfg.Error != nil { req.Header.Set(string(DdInvocationError), "true") - req.Header.Set(string(DdInvocationErrorMsg), err.Error()) + req.Header.Set(string(DdInvocationErrorMsg), cfg.Error.Error()) + req.Header.Set(string(DdInvocationErrorType), reflect.TypeOf(cfg.Error).String()) + req.Header.Set(string(DdInvocationErrorStack), takeStacktrace(cfg)) } // Extract the DD trace context and pass them to the extension via request headers @@ -165,6 +174,43 @@ func (em *ExtensionManager) SendEndInvocationRequest(ctx context.Context, functi } } +// defaultStackLength specifies the default maximum size of a stack trace. +const defaultStackLength = 32 + +// takeStacktrace takes a stack trace of maximum n entries, skipping the first skip entries. +// If n is 0, up to 20 entries are retrieved. +func takeStacktrace(opts ddtrace.FinishConfig) string { + if opts.StackFrames == 0 { + opts.StackFrames = defaultStackLength + } + var builder strings.Builder + pcs := make([]uintptr, opts.StackFrames) + + // +3 to exclude runtime.Callers, takeStacktrace and SendEndInvocationRequest + numFrames := runtime.Callers(3+int(opts.SkipStackFrames), pcs) + if numFrames == 0 { + return "" + } + frames := runtime.CallersFrames(pcs[:numFrames]) + for i := 0; ; i++ { + frame, more := frames.Next() + if i != 0 { + builder.WriteByte('\n') + } + builder.WriteString(frame.Function) + builder.WriteByte('\n') + builder.WriteByte('\t') + builder.WriteString(frame.File) + builder.WriteByte(':') + builder.WriteString(strconv.Itoa(frame.Line)) + if !more { + break + } + } + + return base64.StdEncoding.EncodeToString([]byte(builder.String())) +} + func (em *ExtensionManager) IsExtensionRunning() bool { return em.isExtensionRunning } diff --git a/internal/extension/extension_test.go b/internal/extension/extension_test.go index d1c45036..b4da119b 100644 --- a/internal/extension/extension_test.go +++ b/internal/extension/extension_test.go @@ -11,6 +11,7 @@ package extension import ( "bytes" "context" + "encoding/base64" "fmt" "net/http" "os" @@ -18,6 +19,7 @@ import ( "github.com/DataDog/datadog-lambda-go/internal/logger" "github.com/stretchr/testify/assert" + "gopkg.in/DataDog/dd-trace-go.v1/ddtrace" "gopkg.in/DataDog/dd-trace-go.v1/ddtrace/tracer" ) @@ -203,7 +205,7 @@ func TestExtensionEndInvocation(t *testing.T) { httpClient: &ClientSuccessEndInvoke{}, } span := tracer.StartSpan("aws.lambda") - logOutput := captureLog(func() { em.SendEndInvocationRequest(context.TODO(), span, nil) }) + logOutput := captureLog(func() { em.SendEndInvocationRequest(context.TODO(), span, ddtrace.FinishConfig{}) }) span.Finish() assert.Equal(t, "", logOutput) @@ -215,8 +217,51 @@ func TestExtensionEndInvocationError(t *testing.T) { httpClient: &ClientErrorMock{}, } span := tracer.StartSpan("aws.lambda") - logOutput := captureLog(func() { em.SendEndInvocationRequest(context.TODO(), span, nil) }) + logOutput := captureLog(func() { em.SendEndInvocationRequest(context.TODO(), span, ddtrace.FinishConfig{}) }) span.Finish() assert.Contains(t, logOutput, "could not send end invocation payload to the extension") } + +type capturingClient struct { + hdr http.Header +} + +func (c capturingClient) Do(req *http.Request) (*http.Response, error) { + for k, v := range req.Header { + c.hdr[k] = v + } + return &http.Response{StatusCode: 200}, nil +} + +func TestExtensionEndInvocationErrorHeaders(t *testing.T) { + hdr := http.Header{} + em := &ExtensionManager{httpClient: capturingClient{hdr: hdr}} + span := tracer.StartSpan("aws.lambda") + cfg := ddtrace.FinishConfig{Error: fmt.Errorf("ooooops")} + + em.SendEndInvocationRequest(context.TODO(), span, cfg) + + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error"), "true") + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error-Msg"), "ooooops") + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error-Type"), "*errors.errorString") + + data, err := base64.StdEncoding.DecodeString(hdr.Get("X-Datadog-Invocation-Error-Stack")) + assert.Nil(t, err) + assert.Contains(t, string(data), "github.com/DataDog/datadog-lambda-go") + assert.Contains(t, string(data), "TestExtensionEndInvocationErrorHeaders") +} + +func TestExtensionEndInvocationErrorHeadersNilError(t *testing.T) { + hdr := http.Header{} + em := &ExtensionManager{httpClient: capturingClient{hdr: hdr}} + span := tracer.StartSpan("aws.lambda") + cfg := ddtrace.FinishConfig{Error: nil} + + em.SendEndInvocationRequest(context.TODO(), span, cfg) + + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error"), "") + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error-Msg"), "") + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error-Type"), "") + assert.Equal(t, hdr.Get("X-Datadog-Invocation-Error-Stack"), "") +} diff --git a/internal/trace/listener.go b/internal/trace/listener.go index 6862fda0..41c390cc 100644 --- a/internal/trace/listener.go +++ b/internal/trace/listener.go @@ -31,7 +31,7 @@ type ( ddTraceEnabled bool mergeXrayTraces bool universalInstrumentation bool - otelTracerEnabled bool + otelTracerEnabled bool extensionManager *extension.ExtensionManager traceContextExtractor ContextExtractor } @@ -41,7 +41,7 @@ type ( DDTraceEnabled bool MergeXrayTraces bool UniversalInstrumentation bool - OtelTracerEnabled bool + OtelTracerEnabled bool TraceContextExtractor ContextExtractor } ) @@ -58,7 +58,7 @@ func MakeListener(config Config, extensionManager *extension.ExtensionManager) L ddTraceEnabled: config.DDTraceEnabled, mergeXrayTraces: config.MergeXrayTraces, universalInstrumentation: config.UniversalInstrumentation, - otelTracerEnabled: config.OtelTracerEnabled, + otelTracerEnabled: config.OtelTracerEnabled, extensionManager: extensionManager, traceContextExtractor: config.TraceContextExtractor, } @@ -85,7 +85,7 @@ func (l *Listener) HandlerStarted(ctx context.Context, msg json.RawMessage) cont opts := []tracer.StartOption{ tracer.WithService(serviceName), tracer.WithLambdaMode(extensionNotRunning), - tracer.WithGlobalTag("_dd.origin","lambda"), + tracer.WithGlobalTag("_dd.origin", "lambda"), tracer.WithSendRetries(2), } if l.otelTracerEnabled { @@ -115,8 +115,10 @@ func (l *Listener) HandlerFinished(ctx context.Context, err error) { if functionExecutionSpan != nil { functionExecutionSpan.Finish(tracer.WithError(err)) + finishConfig := ddtrace.FinishConfig{Error: err} + if l.universalInstrumentation && l.extensionManager.IsExtensionRunning() { - l.extensionManager.SendEndInvocationRequest(ctx, functionExecutionSpan, err) + l.extensionManager.SendEndInvocationRequest(ctx, functionExecutionSpan, finishConfig) } }