forked from newrelic/go-agent
-
Notifications
You must be signed in to change notification settings - Fork 0
/
handler.go
162 lines (138 loc) · 4.77 KB
/
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
// Copyright 2020 New Relic Corporation. All rights reserved.
// SPDX-License-Identifier: Apache-2.0
// Package nrlambda adds support for AWS Lambda.
//
// Use this package to instrument your AWS Lambda handler function. Data is
// sent to CloudWatch when the Lambda is invoked. CloudWatch collects Lambda
// log data and sends it to a New Relic log-ingestion Lambda. The log-ingestion
// Lambda sends that data to us.
//
// Monitoring AWS Lambda requires several steps shown here:
// https://docs.newrelic.com/docs/serverless-function-monitoring/aws-lambda-monitoring/get-started/enable-new-relic-monitoring-aws-lambda
//
// Example: https://github.com/newrelic/go-agent/tree/master/_integrations/nrlambda/example/main.go
package nrlambda
import (
"context"
"io"
"net/http"
"os"
"sync"
"github.com/aws/aws-lambda-go/lambda"
"github.com/aws/aws-lambda-go/lambda/handlertrace"
"github.com/aws/aws-lambda-go/lambdacontext"
newrelic "github.com/newrelic/go-agent"
"github.com/newrelic/go-agent/internal"
"github.com/newrelic/go-agent/internal/integrationsupport"
)
type response struct {
header http.Header
code int
}
var _ http.ResponseWriter = &response{}
func (r *response) Header() http.Header { return r.header }
func (r *response) Write([]byte) (int, error) { return 0, nil }
func (r *response) WriteHeader(int) {}
func requestEvent(ctx context.Context, event interface{}) {
txn := newrelic.FromContext(ctx)
if nil == txn {
return
}
if sourceARN := getEventSourceARN(event); "" != sourceARN {
integrationsupport.AddAgentAttribute(txn, internal.AttributeAWSLambdaEventSourceARN, sourceARN, nil)
}
if request := eventWebRequest(event); nil != request {
txn.SetWebRequest(request)
}
}
func responseEvent(ctx context.Context, event interface{}) {
txn := newrelic.FromContext(ctx)
if nil == txn {
return
}
if rw := eventResponse(event); nil != rw && 0 != rw.code {
txn.SetWebResponse(rw)
txn.WriteHeader(rw.code)
}
}
func (h *wrappedHandler) Invoke(ctx context.Context, payload []byte) ([]byte, error) {
var arn, requestID string
if lctx, ok := lambdacontext.FromContext(ctx); ok {
arn = lctx.InvokedFunctionArn
requestID = lctx.AwsRequestID
}
defer internal.ServerlessWrite(h.app, arn, h.writer)
txn := h.app.StartTransaction(h.functionName, nil, nil)
defer txn.End()
integrationsupport.AddAgentAttribute(txn, internal.AttributeAWSRequestID, requestID, nil)
integrationsupport.AddAgentAttribute(txn, internal.AttributeAWSLambdaARN, arn, nil)
h.firstTransaction.Do(func() {
integrationsupport.AddAgentAttribute(txn, internal.AttributeAWSLambdaColdStart, "", true)
})
ctx = newrelic.NewContext(ctx, txn)
ctx = handlertrace.NewContext(ctx, handlertrace.HandlerTrace{
RequestEvent: requestEvent,
ResponseEvent: responseEvent,
})
response, err := h.original.Invoke(ctx, payload)
if nil != err {
txn.NoticeError(err)
}
return response, err
}
type wrappedHandler struct {
original lambda.Handler
app newrelic.Application
// functionName is copied from lambdacontext.FunctionName for
// deterministic tests that don't depend on environment variables.
functionName string
// Although we are told that each Lambda will only handle one request at
// a time, we use a synchronization primitive to determine if this is
// the first transaction for defensiveness in case of future changes.
firstTransaction sync.Once
// writer is used to log the data JSON at the end of each transaction.
// This field exists (rather than hardcoded os.Stdout) for testing.
writer io.Writer
}
// WrapHandler wraps the provided handler and returns a new handler with
// instrumentation. StartHandler should generally be used in place of
// WrapHandler: this function is exposed for consumers who are chaining
// middlewares.
func WrapHandler(handler lambda.Handler, app newrelic.Application) lambda.Handler {
if nil == app {
return handler
}
return &wrappedHandler{
original: handler,
app: app,
functionName: lambdacontext.FunctionName,
writer: os.Stdout,
}
}
// Wrap wraps the provided handler and returns a new handler with
// instrumentation. Start should generally be used in place of Wrap.
func Wrap(handler interface{}, app newrelic.Application) lambda.Handler {
return WrapHandler(lambda.NewHandler(handler), app)
}
// Start should be used in place of lambda.Start. Replace:
//
// lambda.Start(myhandler)
//
// With:
//
// nrlambda.Start(myhandler, app)
//
func Start(handler interface{}, app newrelic.Application) {
lambda.StartHandler(Wrap(handler, app))
}
// StartHandler should be used in place of lambda.StartHandler. Replace:
//
// lambda.StartHandler(myhandler)
//
// With:
//
// nrlambda.StartHandler(myhandler, app)
//
func StartHandler(handler lambda.Handler, app newrelic.Application) {
lambda.StartHandler(WrapHandler(handler, app))
}