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

@@ -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 {

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
 }

internal/extension/extension_test.go

@@ -11,13 +11,15 @@ package extension
 import (
 	"bytes"
 	"context"
+	"encoding/base64"
 	"fmt"
 	"net/http"
 	"os"
 	"testing"
 
 	"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"), "")
+}

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)
 		}
 	}