forked from vmware/vmware-go-kcl
/
worker.go
385 lines (328 loc) · 12.3 KB
/
worker.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
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
/*
* Copyright (c) 2018 VMware, Inc.
*
* Permission is hereby granted, free of charge, to any person obtaining a copy of this software and
* associated documentation files (the "Software"), to deal in the Software without restriction, including
* without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
* copies of the Software, and to permit persons to whom the Software is furnished to do
* so, subject to the following conditions:
*
* The above copyright notice and this permission notice shall be included in all copies or substantial
* portions of the Software.
*
* THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT
* NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
* IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
* WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
* SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
*/
// The implementation is derived from https://github.com/patrobinson/gokini
//
// Copyright 2018 Patrick robinson
//
// Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
package worker
import (
"math/rand"
"sync"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/aws/session"
"github.com/aws/aws-sdk-go/service/kinesis"
"github.com/aws/aws-sdk-go/service/kinesis/kinesisiface"
chk "github.com/Snap-AV/vmware-go-kcl/clientlibrary/checkpoint"
"github.com/Snap-AV/vmware-go-kcl/clientlibrary/config"
kcl "github.com/Snap-AV/vmware-go-kcl/clientlibrary/interfaces"
"github.com/Snap-AV/vmware-go-kcl/clientlibrary/metrics"
par "github.com/Snap-AV/vmware-go-kcl/clientlibrary/partition"
)
/**
* Worker is the high level class that Kinesis applications use to start processing data. It initializes and oversees
* different components (e.g. syncing shard and lease information, tracking shard assignments, and processing data from
* the shards).
*/
type Worker struct {
streamName string
regionName string
workerID string
processorFactory kcl.IRecordProcessorFactory
kclConfig *config.KinesisClientLibConfiguration
kc kinesisiface.KinesisAPI
checkpointer chk.Checkpointer
mService metrics.MonitoringService
stop *chan struct{}
waitGroup *sync.WaitGroup
done bool
rng *rand.Rand
shardStatus map[string]*par.ShardStatus
}
// NewWorker constructs a Worker instance for processing Kinesis stream data.
func NewWorker(factory kcl.IRecordProcessorFactory, kclConfig *config.KinesisClientLibConfiguration) *Worker {
mService := kclConfig.MonitoringService
if mService == nil {
// Replaces nil with noop monitor service (not emitting any metrics).
mService = metrics.NoopMonitoringService{}
}
// Create a pseudo-random number generator and seed it.
rng := rand.New(rand.NewSource(time.Now().UnixNano()))
return &Worker{
streamName: kclConfig.StreamName,
regionName: kclConfig.RegionName,
workerID: kclConfig.WorkerID,
processorFactory: factory,
kclConfig: kclConfig,
mService: mService,
done: false,
rng: rng,
}
}
// WithKinesis is used to provide Kinesis service for either custom implementation or unit testing.
func (w *Worker) WithKinesis(svc kinesisiface.KinesisAPI) *Worker {
w.kc = svc
return w
}
// WithCheckpointer is used to provide a custom checkpointer service for non-dynamodb implementation
// or unit testing.
func (w *Worker) WithCheckpointer(checker chk.Checkpointer) *Worker {
w.checkpointer = checker
return w
}
// Run starts consuming data from the stream, and pass it to the application record processors.
func (w *Worker) Start() error {
log := w.kclConfig.Logger
if err := w.initialize(); err != nil {
log.Errorf("Failed to initialize Worker: %+v", err)
return err
}
// Start monitoring service
log.Infof("Starting monitoring service.")
if err := w.mService.Start(); err != nil {
log.Errorf("Failed to start monitoring service: %+v", err)
return err
}
log.Infof("Starting worker event loop.")
w.waitGroup.Add(1)
go func() {
defer w.waitGroup.Done()
// entering event loop
w.eventLoop()
}()
return nil
}
// Shutdown signals worker to shutdown. Worker will try initiating shutdown of all record processors.
func (w *Worker) Shutdown() {
log := w.kclConfig.Logger
log.Infof("Worker shutdown in requested.")
if w.done || w.stop == nil {
return
}
close(*w.stop)
w.done = true
w.waitGroup.Wait()
w.mService.Shutdown()
log.Infof("Worker loop is complete. Exiting from worker.")
}
// initialize
func (w *Worker) initialize() error {
log := w.kclConfig.Logger
log.Infof("Worker initialization in progress...")
// Create default Kinesis session
if w.kc == nil {
// create session for Kinesis
log.Infof("Creating Kinesis session")
s, err := session.NewSession(&aws.Config{
Region: aws.String(w.regionName),
Endpoint: &w.kclConfig.KinesisEndpoint,
Credentials: w.kclConfig.KinesisCredentials,
})
if err != nil {
// no need to move forward
log.Fatalf("Failed in getting Kinesis session for creating Worker: %+v", err)
}
w.kc = kinesis.New(s)
} else {
log.Infof("Use custom Kinesis service.")
}
// Create default dynamodb based checkpointer implementation
if w.checkpointer == nil {
log.Infof("Creating DynamoDB based checkpointer")
w.checkpointer = chk.NewDynamoCheckpoint(w.kclConfig)
} else {
log.Infof("Use custom checkpointer implementation.")
}
err := w.mService.Init(w.kclConfig.ApplicationName, w.streamName, w.workerID)
if err != nil {
log.Errorf("Failed to start monitoring service: %+v", err)
}
log.Infof("Initializing Checkpointer")
if err := w.checkpointer.Init(); err != nil {
log.Errorf("Failed to start Checkpointer: %+v", err)
return err
}
w.shardStatus = make(map[string]*par.ShardStatus)
stopChan := make(chan struct{})
w.stop = &stopChan
w.waitGroup = &sync.WaitGroup{}
log.Infof("Initialization complete.")
return nil
}
// newShardConsumer to create a shard consumer instance
func (w *Worker) newShardConsumer(shard *par.ShardStatus) *ShardConsumer {
return &ShardConsumer{
streamName: w.streamName,
shard: shard,
kc: w.kc,
checkpointer: w.checkpointer,
recordProcessor: w.processorFactory.CreateProcessor(),
kclConfig: w.kclConfig,
consumerID: w.workerID,
stop: w.stop,
mService: w.mService,
state: WaitingOnParentShards,
}
}
// eventLoop
func (w *Worker) eventLoop() {
log := w.kclConfig.Logger
var foundShards int
for {
// Add [-50%, +50%] random jitter to ShardSyncIntervalMillis. When multiple workers
// starts at the same time, this decreases the probability of them calling
// kinesis.DescribeStream at the same time, and hit the hard-limit on aws API calls.
// On average the period remains the same so that doesn't affect behavior.
shardSyncSleep := w.kclConfig.ShardSyncIntervalMillis/2 + w.rng.Intn(int(w.kclConfig.ShardSyncIntervalMillis))
err := w.syncShard()
if err != nil {
log.Errorf("Error syncing shards: %+v, Retrying in %d ms...", err, shardSyncSleep)
time.Sleep(time.Duration(shardSyncSleep) * time.Millisecond)
continue
}
if foundShards == 0 || foundShards != len(w.shardStatus) {
foundShards = len(w.shardStatus)
log.Infof("Found %d shards", foundShards)
}
// Count the number of leases hold by this worker excluding the processed shard
counter := 0
for _, shard := range w.shardStatus {
if shard.GetLeaseOwner() == w.workerID && shard.Checkpoint != chk.ShardEnd {
counter++
}
}
// max number of lease has not been reached yet
if counter < w.kclConfig.MaxLeasesForWorker {
for _, shard := range w.shardStatus {
// already owner of the shard
if shard.GetLeaseOwner() == w.workerID {
continue
}
err := w.checkpointer.FetchCheckpoint(shard)
if err != nil {
// checkpoint may not existed yet is not an error condition.
if err != chk.ErrSequenceIDNotFound {
log.Errorf(" Error: %+v", err)
// move on to next shard
continue
}
}
// The shard is closed and we have processed all records
if shard.Checkpoint == chk.ShardEnd {
continue
}
err = w.checkpointer.GetLease(shard, w.workerID)
if err != nil {
// cannot get lease on the shard
if err.Error() != chk.ErrLeaseNotAcquired {
log.Errorf("Cannot get lease: %+v", err)
}
continue
}
// log metrics on got lease
w.mService.LeaseGained(shard.ID)
log.Infof("Start Shard Consumer for shard: %v", shard.ID)
sc := w.newShardConsumer(shard)
w.waitGroup.Add(1)
go func() {
defer w.waitGroup.Done()
if err := sc.getRecords(shard); err != nil {
log.Errorf("Error in getRecords: %+v", err)
}
}()
// exit from for loop and not to grab more shard for now.
break
}
}
select {
case <-*w.stop:
log.Infof("Shutting down...")
return
case <-time.After(time.Duration(shardSyncSleep) * time.Millisecond):
log.Debugf("Waited %d ms to sync shards...", shardSyncSleep)
}
}
}
// List all shards and store them into shardStatus table
// If shard has been removed, need to exclude it from cached shard status.
func (w *Worker) getShardIDs(nextToken string, shardInfo map[string]bool) error {
log := w.kclConfig.Logger
args := &kinesis.ListShardsInput{}
// When you have a nextToken, you can't set the streamName
if nextToken != "" {
args.NextToken = aws.String(nextToken)
} else {
args.StreamName = aws.String(w.streamName)
}
listShards, err := w.kc.ListShards(args)
if err != nil {
log.Errorf("Error in ListShards: %s Error: %+v Request: %s", w.streamName, err, args)
return err
}
for _, s := range listShards.Shards {
// record avail shardId from fresh reading from Kinesis
shardInfo[*s.ShardId] = true
// found new shard
if _, ok := w.shardStatus[*s.ShardId]; !ok {
log.Infof("Found new shard with id %s", *s.ShardId)
w.shardStatus[*s.ShardId] = &par.ShardStatus{
ID: *s.ShardId,
ParentShardId: aws.StringValue(s.ParentShardId),
Mux: &sync.Mutex{},
StartingSequenceNumber: aws.StringValue(s.SequenceNumberRange.StartingSequenceNumber),
EndingSequenceNumber: aws.StringValue(s.SequenceNumberRange.EndingSequenceNumber),
}
}
}
if listShards.NextToken != nil {
err := w.getShardIDs(aws.StringValue(listShards.NextToken), shardInfo)
if err != nil {
log.Errorf("Error in ListShards: %s Error: %+v Request: %s", w.streamName, err, args)
return err
}
}
return nil
}
// syncShard to sync the cached shard info with actual shard info from Kinesis
func (w *Worker) syncShard() error {
log := w.kclConfig.Logger
shardInfo := make(map[string]bool)
err := w.getShardIDs("", shardInfo)
if err != nil {
return err
}
for _, shard := range w.shardStatus {
// The cached shard no longer existed, remove it.
if _, ok := shardInfo[shard.ID]; !ok {
// remove the shard from local status cache
delete(w.shardStatus, shard.ID)
// remove the shard entry in dynamoDB as well
// Note: syncShard runs periodically. we don't need to do anything in case of error here.
if err := w.checkpointer.RemoveLeaseInfo(shard.ID); err != nil {
log.Errorf("Failed to remove shard lease info: %s Error: %+v", shard.ID, err)
}
}
}
return nil
}