Lambda instrumentation enhancements

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

@@ -47,7 +47,7 @@ import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go
 
 // wrap lambda handler function
 func main() {
-	lambda.Start(otellambda.InstrumentHandler(HandleRequest))
+    lambda.Start(otellambda.InstrumentHandler(HandleRequest))
 }
 ```
 
@@ -66,14 +66,14 @@ func main() {
 var someHeaderKey = "Key" // used by propagator and EventToCarrier function to identify trace header
 
 type mockHTTPRequest struct {
-	Headers map[string][]string
-	Body string
+    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]}}
+    var request mockHTTPRequest
+    _ = json.unmarshal(eventJSON, &request)
+    return propogation.HeaderCarrier{someHeaderKey: []string{request.Headers[someHeaderKey]}}
 }
 
 type mockPropagator struct{}
@@ -82,28 +82,28 @@ type mockPropagator struct{}
 // Fields
 
 func HandleRequest(ctx context.Context, request mockHTTPRequest) error {
-	return fmt.Sprintf("Hello %s!", request.Body ), nil
+    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{})))
+    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/>
 - For more about OpenTelemetry Go: <https://github.com/open-telemetry/opentelemetry-go>
-- For help or feedback on this project, join us in [GitHub Discussions][discussions-url] 
+- For help or feedback on this project, join us in [GitHub Discussions][discussions-url]
 
 ## License
 
@@ -114,4 +114,4 @@ Apache 2.0 - See [LICENSE][license-url] for more information.
 [goref-image]: https://pkg.go.dev/badge/go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda.svg
 [goref-url]: https://pkg.go.dev/go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda
 [discussions-url]: https://github.com/open-telemetry/opentelemetry-go/discussions
-[lambda-detector-url]: https://github.com/open-telemetry/opentelemetry-go-contrib/blob/main/detectors/aws/lambda
+[lambda-detector-url]: https://github.com/open-telemetry/opentelemetry-go-contrib/blob/main/detectors/aws/lambda

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

@@ -34,7 +34,7 @@ func HandleRequest(ctx context.Context, name MyEvent) (string, error) {
 }
 
 func main() {
-	lambda.Start(otellambda.WrapHandlerFunction(HandleRequest))
+	lambda.Start(otellambda.InstrumentHandler(HandleRequest))
 }
 ```
 
@@ -46,7 +46,7 @@ import "go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go
 
 // add options to WrapHandlerFunction call
 func main() {
-	lambda.Start(otellambda.WrapHandlerFunction(HandleRequest, xrayconfig.AllRecommendedOptions()...))
+	lambda.Start(otellambda.InstrumentHandler(HandleRequest, xrayconfig.AllRecommendedOptions()...))
 }
 ```
 ## Recommended AWS Lambda Instrumentation Options
@@ -56,7 +56,7 @@ func main() {
 | `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
@@ -74,4 +74,4 @@ Apache 2.0 - See [LICENSE][license-url] for more information.
 [license-image]: https://img.shields.io/badge/license-Apache_2.0-green.svg?style=flat
 [goref-image]: https://pkg.go.dev/badge/go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig.svg
 [goref-url]: https://pkg.go.dev/go.opentelemetry.io/contrib/instrumentation/github.com/aws/aws-lambda-go/otellambda/xrayconfig
-[discussions-url]: https://github.com/open-telemetry/opentelemetry-go/discussions
+[discussions-url]: https://github.com/open-telemetry/opentelemetry-go/discussions

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)