-
Notifications
You must be signed in to change notification settings - Fork 2
/
main.go
184 lines (160 loc) · 5.57 KB
/
main.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
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
package main
import (
"context"
"os"
"os/signal"
"path/filepath"
"sync"
"syscall"
"time"
"github.com/rs/zerolog"
"gopkg.in/alecthomas/kingpin.v2"
"github.com/david7482/lambda-extension-log-shipper/extension"
"github.com/david7482/lambda-extension-log-shipper/forwardservice"
"github.com/david7482/lambda-extension-log-shipper/forwardservice/forwarders/newrelic"
"github.com/david7482/lambda-extension-log-shipper/forwardservice/forwarders/stdout"
"github.com/david7482/lambda-extension-log-shipper/logservice"
)
const (
// ListenPort is the port that our log server listens on.
listenPort = 8443
// MaxItems is the maximum number of events to be buffered in memory. (default: 10000, minimum: 1000, maximum: 10000)
maxItems = 10000
// MaxBytes is the maximum size in bytes of the logs to be buffered in memory. (default: 262144, minimum: 262144, maximum: 1048576)
maxBytes = 262144
// TimeoutMS is the maximum time (in milliseconds) for a batch to be buffered. (default: 1000, minimum: 100, maximum: 30000)
timeoutMS = 1000
)
var (
extensionName = filepath.Base(os.Args[0]) // extension name has to match the filename
logTypes = []extension.LogType{extension.Platform, extension.Function}
forwarders = []forwardservice.Forwarder{stdout.New(), newrelic.New()}
)
type generalConfig struct {
AWSLambdaName *string
AWSRegion *string
AWSRuntimeAPI *string
LogLevel *string
LogTimeFormat *string
EnablePlatformReport *bool
}
func setupGeneralConfigs(app *kingpin.Application) generalConfig {
var config generalConfig
// the followings would read from lambda runtime environment variables
config.AWSLambdaName = app.
Flag("lambda-name", "The name of the lambda function").
Envar("AWS_LAMBDA_FUNCTION_NAME").
Required().String()
config.AWSRegion = app.
Flag("region", "The AWS Region where the Lambda function is executed").
Envar("AWS_REGION").
Required().String()
config.AWSRuntimeAPI = app.
Flag("runtime-api", "The endpoint URL of lambda extension runtime API").
Envar("AWS_LAMBDA_RUNTIME_API").
Required().String()
// the followings are general settings
config.LogLevel = app.
Flag("log-level", "The level of the internal logger").
Envar("LS_LOG_LEVEL").
Default("info").Enum("error", "warn", "info", "debug")
config.LogTimeFormat = app.
Flag("log-timeformat", "The time format of the internal logger").
Envar("LS_LOG_TIMEFORMAT").
Default("2006-01-02T15:04:05.000Z07:00").String()
config.EnablePlatformReport = app.
Flag("enable-platform-report", "Send Lambda platform report to all forwarders").
Envar("LS_ENABLE_PLATFORM_REPORT").
Default("true").Bool()
return config
}
func setupForwarderConfigs(app *kingpin.Application) {
// let each forwarder setup its own configurations
for _, f := range forwarders {
f.SetupConfigs(app)
}
}
func main() {
// Setup configurations
app := kingpin.New("lambda-extension-log-shipper", "Lambda Extension Log Shipper")
cfg := setupGeneralConfigs(app)
setupForwarderConfigs(app)
kingpin.MustParse(app.Parse(os.Args[1:]))
// Setup zerolog
lvl, _ := zerolog.ParseLevel(*cfg.LogLevel)
zerolog.SetGlobalLevel(lvl)
zerolog.TimeFieldFormat = *cfg.LogTimeFormat
rootLogger := zerolog.New(os.Stdout).With().Timestamp().Logger()
// Create root context
rootCtx, rootCtxCancelFunc := context.WithCancel(context.Background())
rootCtx = rootLogger.WithContext(rootCtx)
rootLogger.Info().Interface("config", cfg).Msg("lambda-extension-log-shipper start...")
// Register extension as soon as possible
extensionClient := extension.NewClient(*cfg.AWSRuntimeAPI)
_, err := extensionClient.RegisterExtension(rootCtx, extensionName)
if err != nil {
rootLogger.Fatal().Err(err).Msg("fail to register extension")
}
// Create the logs queue
logsQueue := make(chan []logservice.Log, 8)
// Start services
wg := sync.WaitGroup{}
wg.Add(1)
logSrv := logservice.New(logservice.ServiceParams{
LogAPIClient: extensionClient,
LogTypes: logTypes,
LogsQueue: logsQueue,
ListenPort: listenPort,
MaxItems: maxItems,
MaxBytes: maxBytes,
TimeoutMS: timeoutMS,
EnablePlatformReport: *cfg.EnablePlatformReport,
})
logSrv.Run(rootCtx, &wg)
wg.Add(1)
forwardSrv := forwardservice.New(forwardservice.ServiceParams{
Forwarders: forwarders,
LogsQueue: logsQueue,
LambdaName: *cfg.AWSLambdaName,
AWSRegion: *cfg.AWSRegion,
})
forwardSrv.Run(rootCtx, &wg)
// Listen to SIGTEM/SIGINT to close
var gracefulStop = make(chan os.Signal, 1)
signal.Notify(gracefulStop, syscall.SIGTERM, syscall.SIGINT)
// Will block until invoke or shutdown event is received or cancelled via the context.
LOOP:
for {
select {
case s := <-gracefulStop:
rootLogger.Info().Msgf("received signal to terminate: %s", s.String())
break LOOP
default:
// This is a blocking call
res, err := extensionClient.NextEvent(rootCtx)
if err != nil {
rootLogger.Error().Err(err).Msg("fail to invoke NextEvent")
return
}
// Exit if we receive a SHUTDOWN event
if res.EventType == extension.Shutdown {
rootLogger.Info().Msg("received SHUTDOWN event")
break LOOP
}
}
}
// Close root context to terminate everything
rootCtxCancelFunc()
// Wait for all services to close with a specific timeout
var waitUntilDone = make(chan struct{})
go func() {
wg.Wait()
close(waitUntilDone)
}()
select {
case <-waitUntilDone:
rootLogger.Info().Msg("success to close all services")
case <-time.After(1950 * time.Millisecond):
rootLogger.Err(context.DeadlineExceeded).Msg("fail to close all services")
}
}