Introduce HandlerTrace

Author: [email protected]

lambda/handler.go

@@ -70,6 +70,42 @@ func validateReturns(handler reflect.Type) error {
 	return nil
 }
 
+// HandlerTrace allows handlers which wrap the return value of 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{}
+
+// WithHandlerTrace adds callbacks to the provided context which allows handlers
+// which wrap the return value of NewHandler to access to the request and
+// response events.
+func WithHandlerTrace(ctx context.Context, trace HandlerTrace) context.Context {
+	existing := contextHandlerTrace(ctx)
+	return context.WithValue(ctx, handlerTraceKey{}, HandlerTrace{
+		RequestEvent:  callbackCompose(existing.RequestEvent, trace.RequestEvent),
+		ResponseEvent: callbackCompose(existing.ResponseEvent, trace.ResponseEvent),
+	})
+}
+
+func contextHandlerTrace(ctx context.Context) HandlerTrace {
+	trace, _ := ctx.Value(handlerTraceKey{}).(HandlerTrace)
+	return trace
+}
+
 // NewHandler creates a base lambda handler from the given handler function. The
 // returned Handler performs JSON deserialization and deserialization, and
 // delegates to the input handler function.  The handler function parameter must
@@ -95,6 +131,9 @@ func NewHandler(handlerFunc interface{}) Handler {
 	}
 
 	return lambdaHandler(func(ctx context.Context, payload []byte) (interface{}, error) {
+
+		trace := contextHandlerTrace(ctx)
+
 		// construct arguments
 		var args []reflect.Value
 		if takesContext {
@@ -107,7 +146,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 +164,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

@@ -206,3 +206,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 = WithHandlerTrace(ctx, HandlerTrace{}) // empty HandlerTrace
+	ctx = WithHandlerTrace(ctx, HandlerTrace{   // with RequestEvent
+		RequestEvent: func(c context.Context, e interface{}) {
+			requestHistory += "A"
+			checkInt(e, 123)
+		},
+	})
+	ctx = WithHandlerTrace(ctx, HandlerTrace{ // with ResponseEvent
+		ResponseEvent: func(c context.Context, e interface{}) {
+			responseHistory += "X"
+			checkInt(e, 456)
+		},
+	})
+	ctx = WithHandlerTrace(ctx, 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 = WithHandlerTrace(ctx, 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)
+	}
+}