/
trace.go
44 lines (38 loc) · 1.38 KB
/
trace.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
// 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
}