-
Notifications
You must be signed in to change notification settings - Fork 39
/
wrap_handler.go
163 lines (139 loc) · 4.71 KB
/
wrap_handler.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
/*
* Unless explicitly stated otherwise all files in this repository are licensed
* under the Apache License Version 2.0.
*
* This product includes software developed at Datadog (https://www.datadoghq.com/).
* Copyright 2019 Datadog, Inc.
*/
package wrapper
import (
"context"
"encoding/json"
"errors"
"fmt"
"reflect"
"github.com/DataDog/datadog-lambda-go/internal/logger"
)
var (
// CurrentContext is the last create lambda context object.
CurrentContext context.Context
)
type (
// HandlerListener is a point where listener logic can be injected into a handler
HandlerListener interface {
HandlerStarted(ctx context.Context, msg json.RawMessage) context.Context
HandlerFinished(ctx context.Context)
}
)
// WrapHandlerWithListeners wraps a lambda handler, and calls listeners before and after every invocation.
func WrapHandlerWithListeners(handler interface{}, listeners ...HandlerListener) interface{} {
err := validateHandler(handler)
if err != nil {
// This wasn't a valid handler function, pass back to AWS SDK to let it handle the error.
logger.Error(fmt.Errorf("handler function was in format ddlambda doesn't recognize: %v", err))
return handler
}
coldStart := true
// Return custom handler, to be called once per invocation
return func(ctx context.Context, msg json.RawMessage) (interface{}, error) {
ctx = context.WithValue(ctx, "cold_start", coldStart)
for _, listener := range listeners {
ctx = listener.HandlerStarted(ctx, msg)
}
CurrentContext = ctx
result, err := callHandler(ctx, msg, handler)
if err != nil {
ctx = context.WithValue(ctx, "error", true)
}
for _, listener := range listeners {
listener.HandlerFinished(ctx)
}
coldStart = false
CurrentContext = nil
return result, err
}
}
func validateHandler(handler interface{}) error {
// Detect the handler follows the right format, based on the GO AWS SDK.
// https://docs.aws.amazon.com/lambda/latest/dg/go-programming-model-handler-types.html
handlerType := reflect.TypeOf(handler)
if handlerType.Kind() != reflect.Func {
return errors.New("handler is not a function")
}
if handlerType.NumIn() == 2 {
contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
firstArgType := handlerType.In(0)
if !firstArgType.Implements(contextType) {
return errors.New("handler should take context as first argument")
}
}
if handlerType.NumIn() > 2 {
return errors.New("handler takes too many arguments")
}
errorType := reflect.TypeOf((*error)(nil)).Elem()
if handlerType.NumOut() > 2 {
return errors.New("handler returns more than two values")
}
if handlerType.NumOut() > 0 {
rt := handlerType.Out(handlerType.NumOut() - 1) // Last returned value
if !rt.Implements(errorType) {
return errors.New("handler doesn't return error as it's last value")
}
}
return nil
}
func callHandler(ctx context.Context, msg json.RawMessage, handler interface{}) (interface{}, error) {
ev, err := unmarshalEventForHandler(msg, handler)
if err != nil {
return nil, err
}
handlerType := reflect.TypeOf(handler)
args := []reflect.Value{}
if handlerType.NumIn() == 1 {
// When there is only one argument, argument is either the event payload, or the context.
contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
firstArgType := handlerType.In(0)
if firstArgType.Implements(contextType) {
args = []reflect.Value{reflect.ValueOf(ctx)}
} else {
args = []reflect.Value{ev.Elem()}
}
} else if handlerType.NumIn() == 2 {
// Or when there are two arguments, context is always first, followed by event payload.
args = []reflect.Value{reflect.ValueOf(ctx), ev.Elem()}
}
handlerValue := reflect.ValueOf(handler)
output := handlerValue.Call(args)
var response interface{}
var errResponse error
if len(output) > 0 {
// If there are any output values, the last should always be an error
val := output[len(output)-1].Interface()
if errVal, ok := val.(error); ok {
errResponse = errVal
}
}
if len(output) > 1 {
// If there is more than one output value, the first should be the response payload.
response = output[0].Interface()
}
return response, errResponse
}
func unmarshalEventForHandler(ev json.RawMessage, handler interface{}) (reflect.Value, error) {
handlerType := reflect.TypeOf(handler)
if handlerType.NumIn() == 0 {
return reflect.ValueOf(nil), nil
}
messageType := handlerType.In(handlerType.NumIn() - 1)
contextType := reflect.TypeOf((*context.Context)(nil)).Elem()
firstArgType := handlerType.In(0)
if handlerType.NumIn() == 1 && firstArgType.Implements(contextType) {
return reflect.ValueOf(nil), nil
}
newMessage := reflect.New(messageType)
err := json.Unmarshal(ev, newMessage.Interface())
if err != nil {
return reflect.ValueOf(nil), err
}
return newMessage, err
}