diff --git a/lambda/handler.go b/lambda/handler.go index 0923a6e1..96eeac2e 100644 --- a/lambda/handler.go +++ b/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 diff --git a/lambda/handler_test.go b/lambda/handler_test.go index 15674dcf..962ac38c 100644 --- a/lambda/handler_test.go +++ b/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) + } +} diff --git a/lambda/handlertrace/trace.go b/lambda/handlertrace/trace.go new file mode 100644 index 00000000..cdd452ff --- /dev/null +++ b/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 +}