Introduce HandlerTrace callbacks (#141)

* Introduce HandlerTrace

* Move handler trace logic into handlertrace package

* Address review feedback: Improve handlertrace names

Author: willnewrelic

lambda/handler.go

@@ -7,6 +7,8 @@ import (
 	"encoding/json"
 	"fmt"
 	"reflect"
+
+	"github.com/aws/aws-lambda-go/lambda/handlertrace"
 )
 
 type Handler interface {
@@ -95,6 +97,9 @@ func NewHandler(handlerFunc interface{}) Handler {
 	}
 
 	return lambdaHandler(func(ctx context.Context, payload []byte) (interface{}, error) {
+
+		trace := handlertrace.FromContext(ctx)
+
 		// construct arguments
 		var args []reflect.Value
 		if takesContext {
@@ -107,7 +112,9 @@ func NewHandler(handlerFunc interface{}) Handler {
 			if err := json.Unmarshal(payload, event.Interface()); err != nil {
 				return nil, err
 			}
-
+			if nil != trace.RequestEvent {
+				trace.RequestEvent(ctx, event.Elem().Interface())
+			}
 			args = append(args, event.Elem())
 		}
 
@@ -123,6 +130,10 @@ func NewHandler(handlerFunc interface{}) Handler {
 		var val interface{}
 		if len(response) > 1 {
 			val = response[0].Interface()
+
+			if nil != trace.ResponseEvent {
+				trace.ResponseEvent(ctx, val)
+			}
 		}
 
 		return val, err

lambda/handler_test.go

@@ -8,6 +8,7 @@ import (
 	"fmt"
 	"testing"
 
+	"github.com/aws/aws-lambda-go/lambda/handlertrace"
 	"github.com/stretchr/testify/assert"
 )
 
@@ -206,3 +207,64 @@ func TestInvalidJsonInput(t *testing.T) {
 	assert.Equal(t, "unexpected end of JSON input", err.Error())
 
 }
+
+func TestHandlerTrace(t *testing.T) {
+	handler := NewHandler(func(ctx context.Context, x int) (int, error) {
+		if x != 123 {
+			t.Error(x)
+		}
+		return 456, nil
+	})
+	requestHistory := ""
+	responseHistory := ""
+	checkInt := func(e interface{}, expected int) {
+		nt, ok := e.(int)
+		if !ok {
+			t.Error("not int as expected", e)
+			return
+		}
+		if nt != expected {
+			t.Error("unexpected value", nt, expected)
+		}
+	}
+	ctx := context.Background()
+	ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{}) // empty HandlerTrace
+	ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{   // with RequestEvent
+		RequestEvent: func(c context.Context, e interface{}) {
+			requestHistory += "A"
+			checkInt(e, 123)
+		},
+	})
+	ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{ // with ResponseEvent
+		ResponseEvent: func(c context.Context, e interface{}) {
+			responseHistory += "X"
+			checkInt(e, 456)
+		},
+	})
+	ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{ // with RequestEvent and ResponseEvent
+		RequestEvent: func(c context.Context, e interface{}) {
+			requestHistory += "B"
+			checkInt(e, 123)
+		},
+		ResponseEvent: func(c context.Context, e interface{}) {
+			responseHistory += "Y"
+			checkInt(e, 456)
+		},
+	})
+	ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{}) // empty HandlerTrace
+
+	payload := []byte(`123`)
+	js, err := handler.Invoke(ctx, payload)
+	if err != nil {
+		t.Error("unexpected handler error", err)
+	}
+	if string(js) != "456" {
+		t.Error("unexpected handler output", string(js))
+	}
+	if requestHistory != "AB" {
+		t.Error("request callbacks not called as expected", requestHistory)
+	}
+	if responseHistory != "XY" {
+		t.Error("response callbacks not called as expected", responseHistory)
+	}
+}

lambda/handlertrace/trace.go

@@ -0,0 +1,44 @@
+// Package handlertrace allows middleware authors using lambda.NewHandler to
+// instrument request and response events.
+package handlertrace
+
+import (
+	"context"
+)
+
+// HandlerTrace allows handlers which wrap the return value of lambda.NewHandler
+// to access to the request and response events.
+type HandlerTrace struct {
+	RequestEvent  func(context.Context, interface{})
+	ResponseEvent func(context.Context, interface{})
+}
+
+func callbackCompose(f1, f2 func(context.Context, interface{})) func(context.Context, interface{}) {
+	return func(ctx context.Context, event interface{}) {
+		if nil != f1 {
+			f1(ctx, event)
+		}
+		if nil != f2 {
+			f2(ctx, event)
+		}
+	}
+}
+
+type handlerTraceKey struct{}
+
+// NewContext adds callbacks to the provided context which allows handlers which
+// wrap the return value of lambda.NewHandler to access to the request and
+// response events.
+func NewContext(ctx context.Context, trace HandlerTrace) context.Context {
+	existing := FromContext(ctx)
+	return context.WithValue(ctx, handlerTraceKey{}, HandlerTrace{
+		RequestEvent:  callbackCompose(existing.RequestEvent, trace.RequestEvent),
+		ResponseEvent: callbackCompose(existing.ResponseEvent, trace.ResponseEvent),
+	})
+}
+
+// FromContext returns the HandlerTrace associated with the provided context.
+func FromContext(ctx context.Context) HandlerTrace {
+	trace, _ := ctx.Value(handlerTraceKey{}).(HandlerTrace)
+	return trace
+}