Lambda instrumentation enhancements

instrumentation/github.com/aws/aws-lambda-go/otellambda/README.md

@@ -16,38 +16,65 @@ See [./example](https://github.com/open-telemetry/opentelemetry-go-contrib/tree/
 
 ## Usage
 
-Create a sample Lambda Go application such as below.
+Create a sample Lambda Go application instrumented by the `otellambda` package such as below.
 
 ```go
 package main
 
 import (
 	"context"
 	"fmt"
+
 	"github.com/aws/aws-lambda-go/lambda"
+	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
 )
 
-type MyEvent struct {
-	Name string `json:"name"`
-}
-
-func HandleRequest(ctx context.Context, name MyEvent) (string, error) {
-	return fmt.Sprintf("Hello %s!", name.Name ), nil
+func HandleRequest(ctx context.Context) (error) {
+	fmt.Println("Hello World!" )
+	return nil
 }
 
 func main() {
-	lambda.Start(HandleRequest)
+	ctx := context.Background()
+	lambda.Start(otellambda.InstrumentHandler(HandleRequest(ctx)))
 }
 ```
 
-Now use the provided wrapper to instrument your basic Lambda function:
+Now configure the instrumentation with the provided options to export traces to AWS X-Ray via [the OpenTelemetry Collector](https://github.com/open-telemetry/opentelemetry-collector) running as a Lambda Extension. Instructions for running the OTel Collector as a Lambda Extension can be found in the [AWS OpenTelemetry Documentation](https://aws-otel.github.io/docs/getting-started/lambda).
+
 ```go
-// Add import
-import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
+// Add imports
+import (
+    "context"
+    "fmt"
+
+    "github.com/aws/aws-lambda-go/lambda"
+    "go.opentelemetry.io/contrib/propagators/aws/xray"
+    "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
+    "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig"
+    "go.opentelemetry.io/otel"
+)
 
-// wrap lambda handler function
+// add options to WrapHandlerFunction call
 func main() {
-	lambda.Start(otellambda.InstrumentHandler(HandleRequest))
+    ctx := context.Background()
+
+    tp, err := xrayconfig.PrepareTracerProvider(ctx)
+    if err != nil {
+        fmt.Printf("error creating tracer provider: %v", err)
+    }
+
+    defer func(ctx context.Context) {
+        err := tp.Shutdown(ctx)
+        if err != nil {
+            fmt.Printf("error shutting down tracer provider: %v", err)
+        }
+    }(ctx)
+
+    otel.SetTracerProvider(tp)
+    otel.SetTextMapPropagator(xray.Propagator{})
+
+    lambda.Start(otellambda.InstrumentHandler(HandleRequest(ctx), xrayconfig.AllRecommendedOptions(tp)...))
 }
 ```
 
@@ -60,45 +87,6 @@ func main() {
 | `WithEventToCarrier` | `func(eventJSON []byte) propagation.TextMapCarrier{}` | Function for providing custom logic to support retrieving trace header from different event types that are handled by AWS Lambda (e.g., SQS, CloudWatch, Kinesis, API Gateway) and returning them in a `propagation.TextMapCarrier` which a Propagator can use to extract the trace header into the context. | Function which returns an empty `TextMapCarrier` - new spans will be part of a new Trace and have no parent past Lambda instrumentation span
 | `WithPropagator` | `propagation.Propagator` | The `Propagator` the instrumentation will use to extract trace information into the context. | `otel.GetTextMapPropagator()` |
 
-### Usage With Options Example
-
-```go
-var someHeaderKey = "Key" // used by propagator and EventToCarrier function to identify trace header
-
-type mockHTTPRequest struct {
-	Headers map[string][]string
-	Body string
-}
-
-func mockEventToCarrier(eventJSON []byte) propagation.TextMapCarrier{
-	var request mockHTTPRequest
-	_ = json.unmarshal(eventJSON, &request)
-	return propogation.HeaderCarrier{someHeaderKey: []string{request.Headers[someHeaderKey]}}
-}
-
-type mockPropagator struct{}
-// Extract - read from `someHeaderKey`
-// Inject
-// Fields
-
-func HandleRequest(ctx context.Context, request mockHTTPRequest) error {
-	return fmt.Sprintf("Hello %s!", request.Body ), nil
-}
-
-func main() {
-	exp, _ := stdouttrace.New()
-
-	tp := sdktrace.NewTracerProvider(
-		    sdktrace.WithBatcher(exp))
-
-	lambda.Start(otellambda.InstrumentHandler(HandleRequest,
-		                                    otellambda.WithTracerProvider(tp),
-		                                    otellambda.WithFlusher(tp),
-		                                    otellambda.WithEventToCarrier(mockEventToCarrier),
-		                                    otellambda.WithPropagator(mockPropagator{})))
-}
-```
-
 ## Useful links
 
 - For more information on OpenTelemetry, visit: <https://opentelemetry.io/>

instrumentation/github.com/aws/aws-lambda-go/otellambda/config.go

@@ -16,6 +16,7 @@ package otellambda
 
 import (
 	"context"
+	"runtime"
 
 	"go.opentelemetry.io/otel/propagation"
 	"go.opentelemetry.io/otel/trace"
@@ -92,6 +93,20 @@ type config struct {
 	Propagator propagation.TextMapPropagator
 }
 
+type asyncSafeFlusher struct {
+	flusher Flusher
+}
+
+func (f asyncSafeFlusher) ForceFlush(ctx context.Context) error {
+	// yield processor to attempt to ensure all spans have
+	// been consumed and are ready to be flushed
+	// - see https://github.com/open-telemetry/opentelemetry-go/issues/2080
+	// to be removed upon resolution of above issue
+	runtime.Gosched()
+
+	return f.flusher.ForceFlush(ctx)
+}
+
 func WithTracerProvider(tracerProvider trace.TracerProvider) Option {
 	return optionFunc(func(c *config) {
 		c.TracerProvider = tracerProvider
@@ -115,3 +130,9 @@ func WithPropagator(propagator propagation.TextMapPropagator) Option {
 		c.Propagator = propagator
 	})
 }
+
+func WithAsyncSafeFlusher(f Flusher) Option {
+	return optionFunc(func(c *config) {
+		c.Flusher = asyncSafeFlusher{f}
+	})
+}

instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig/README.md

@@ -11,52 +11,14 @@ This module provides recommended configuration options for [`AWS Lambda Instrume
 go get -u go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig
 ```
 
-## Usage
-
-Create a sample Lambda Go application instrumented by the `otellambda` package such as below.
-
-```go
-package main
-
-import (
-	"context"
-	"fmt"
-	"github.com/aws/aws-lambda-go/lambda"
-	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
-)
-
-type MyEvent struct {
-	Name string `json:"name"`
-}
-
-func HandleRequest(ctx context.Context, name MyEvent) (string, error) {
-	return fmt.Sprintf("Hello %s!", name.Name ), nil
-}
-
-func main() {
-	lambda.Start(otellambda.WrapHandlerFunction(HandleRequest))
-}
-```
-
-Now configure the instrumentation with the provided options to export traces to AWS X-Ray via [the OpenTelemetry Collector](https://github.com/open-telemetry/opentelemetry-collector) running as a Lambda Extension. Instructions for running the OTel Collector as a Lambda Extension can be found in the [AWS OpenTelemetry Documentation](https://aws-otel.github.io/docs/getting-started/lambda).
-
-```go
-// Add import
-import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig"
-
-// add options to WrapHandlerFunction call
-func main() {
-	lambda.Start(otellambda.WrapHandlerFunction(HandleRequest, xrayconfig.AllRecommendedOptions()...))
-}
-```
 ## Recommended AWS Lambda Instrumentation Options
 
 | Instrumentation Option | Recommended Value | Exported As |
 | --- | --- | --- |
 | `WithTracerProvider` | An `sdktrace.TracerProvider` configured to export in batches to an OTel Collector running locally in Lambda | Not individually exported. Can only be used via `AllRecommendedOptions()`
 | `WithFlusher` | An `otellambda.Flusher` which yields before calling ForceFlush on the configured `sdktrace.TracerProvider`. Yielding mitigates data delays caused by asynchronous nature of batching TracerProvider when in Lambda | Not individually exported. Can only be used via `AllRecommendedOptions()`
 | `WithEventToCarrier` | Function which reads X-Ray TraceID from Lambda environment and inserts it into a `propagtation.TextMapCarrier` | Individually exported as `EventToCarrier()`, also included in `AllRecommendedOptions()`
-| `WithPropagator` | An `xray.propagator` | Individually exported as `EventToCarrier()`, also included in `AllRecommendedOptions()`
+| `WithPropagator` | An `xray.propagator` | Individually exported as `Propagator()`, also included in `AllRecommendedOptions()`
 
 
 ## Useful links

instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig/xrayconfig.go

@@ -18,7 +18,6 @@ import (
 	"context"
 	"log"
 	"os"
-	"runtime"
 
 	lambdadetector "go.opentelemetry.io/contrib/detectors/aws/lambda"
 	"go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda"
@@ -35,48 +34,37 @@ func xrayEventToCarrier([]byte) propagation.TextMapCarrier {
 	return propagation.HeaderCarrier{"X-Amzn-Trace-Id": []string{xrayTraceID}}
 }
 
-type asyncSafeFlusher struct {
-	tp *sdktrace.TracerProvider
-}
-
-func (f asyncSafeFlusher) ForceFlush(ctx context.Context) error {
-	// yield processor to attempt to ensure all spans have
-	// been consumed and are ready to be flushed
-	// - see https://github.com/open-telemetry/opentelemetry-go/issues/2080
-	// to be removed upon resolution of above issue
-	runtime.Gosched()
-
-	return f.tp.ForceFlush(ctx)
-}
-
-// tracerProviderAndFlusher returns a list of otellambda.Option(s) to
-// enable using a TracerProvider configured for AWS XRay via a collector
-// and an otellambda.Flusher to flush this TracerProvider.
-// tracerProviderAndFlusher is not exported because it should not be used
-// without the provided EventToCarrier function and XRay Propagator
-func tracerProviderAndFlusher() ([]otellambda.Option, error) {
-	ctx := context.Background()
-
-	// Do not need transport security in Lambda because collector
-	// runs locally in Lambda execution environment
+// PrepareTracerProvider returns a TracerProvider configured with exporter,
+// id generator and lambda resource detector to send trace data to AWS X-Ray via Collector
+func PrepareTracerProvider(ctx context.Context) (*sdktrace.TracerProvider, error) {
+	log.Println("creating trace exporter")
 	exp, err := otlptracegrpc.New(ctx, otlptracegrpc.WithInsecure())
 	if err != nil {
-		return []otellambda.Option{}, err
+		errorLogger.Println("failed to create exporter: ", err)
+		return nil, err
 	}
 
 	detector := lambdadetector.NewResourceDetector()
-	res, err := detector.Detect(ctx)
+	resource, err := detector.Detect(ctx)
 	if err != nil {
-		return []otellambda.Option{}, err
+		errorLogger.Println("failed to detect lambda resources: ", err)
+		return nil, err
 	}
 
-	tp := sdktrace.NewTracerProvider(
+	return sdktrace.NewTracerProvider(
 		sdktrace.WithBatcher(exp),
 		sdktrace.WithIDGenerator(xray.NewIDGenerator()),
-		sdktrace.WithResource(res),
-	)
+		sdktrace.WithResource(resource),
+	), nil
+}
 
-	return []otellambda.Option{otellambda.WithTracerProvider(tp), otellambda.WithFlusher(asyncSafeFlusher{tp: tp})}, nil
+// tracerProviderAndFlusher returns a list of otellambda.Option(s) to
+// enable using a TracerProvider configured for AWS XRay via a collector
+// and an otellambda.Flusher to flush this TracerProvider.
+// tracerProviderAndFlusher is not exported because it should not be used
+// without the provided EventToCarrier function and XRay Propagator
+func tracerProviderAndFlusher(tp *sdktrace.TracerProvider) ([]otellambda.Option, error) {
+	return []otellambda.Option{otellambda.WithTracerProvider(tp), otellambda.WithAsyncSafeFlusher(tp)}, nil
 }
 
 // EventToCarrier returns an otellambda.Option to enable
@@ -89,14 +77,13 @@ func EventToCarrier() otellambda.Option {
 
 // Propagator returns an otellambda.Option to enable the xray.Propagator
 func Propagator() otellambda.Option {
-
 	return otellambda.WithPropagator(xray.Propagator{})
 }
 
 // AllRecommendedOptions returns a list of all otellambda.Option(s)
 // recommended for the otellambda package when using AWS XRay
-func AllRecommendedOptions() []otellambda.Option {
-	options, err := tracerProviderAndFlusher()
+func AllRecommendedOptions(tp *sdktrace.TracerProvider) []otellambda.Option {
+	options, err := tracerProviderAndFlusher(tp)
 	if err != nil {
 		// should we fail to create the TracerProvider, do not alter otellambda's default configuration
 		errorLogger.Println("failed to create recommended configuration: ", err)

instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig/xrayconfig_test.go

@@ -154,6 +154,10 @@ func assertSpanEqualsIgnoreTimeAndSpanID(t *testing.T, expected *v1trace.Resourc
 func TestWrapEndToEnd(t *testing.T) {
 	setEnvVars()
 
+	ctx := context.Background()
+
+	tp, _ := PrepareTracerProvider(ctx)
+
 	customerHandler := func() (string, error) {
 		return "hello world", nil
 	}
@@ -163,7 +167,7 @@ func TestWrapEndToEnd(t *testing.T) {
 	}()
 	<-time.After(5 * time.Millisecond)
 
-	wrapped := otellambda.InstrumentHandler(customerHandler, AllRecommendedOptions()...)
+	wrapped := otellambda.InstrumentHandler(customerHandler, AllRecommendedOptions(tp)...)
 	wrappedCallable := reflect.ValueOf(wrapped)
 	resp := wrappedCallable.Call([]reflect.Value{reflect.ValueOf(mockContext)})
 	assert.Len(t, resp, 2)