-
Notifications
You must be signed in to change notification settings - Fork 175
/
consumer.go
380 lines (306 loc) · 9.68 KB
/
consumer.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
package jobqueue
import (
"errors"
"fmt"
"sync"
"github.com/rs/zerolog"
"go.uber.org/atomic"
"github.com/onflow/flow-go/module"
"github.com/onflow/flow-go/storage"
)
type Worker interface {
// returned error must be unexpected fatal error
Run(job module.Job) error
}
type Consumer struct {
mu sync.Mutex
log zerolog.Logger
// Storage
jobs module.Jobs // storage to read jobs from
progress storage.ConsumerProgress // to resume from first unprocessed job after restarting
// dependency
worker Worker // to process job and notify consumer when finish processing a job
// Config
maxProcessing uint64 // max number of jobs to be processed concurrently
maxSearchAhead uint64 // max number of jobs beyond processedIndex to process. 0 means no limit
// State Variables
running bool // a signal to control whether to start processing more jobs. Useful for waiting
// until the workers are ready
isChecking *atomic.Bool // allow only one process checking job processable
// are ready, and stop when shutting down.
runningJobs sync.WaitGroup // to wait for all existing jobs to finish for graceful shutdown
processedIndex uint64
processings map[uint64]*jobStatus // keep track of the status of each on going job
processingsIndex map[module.JobID]uint64 // lookup the index of the job, useful when fast forwarding the
// `processed` variable
}
func NewConsumer(
log zerolog.Logger,
jobs module.Jobs,
progress storage.ConsumerProgress,
worker Worker,
maxProcessing uint64,
maxSearchAhead uint64,
) *Consumer {
return &Consumer{
log: log.With().Str("sub_module", "job_queue").Logger(),
// store dependency
jobs: jobs,
progress: progress,
worker: worker,
// update config
maxProcessing: maxProcessing,
maxSearchAhead: maxSearchAhead,
// init state variables
running: false,
isChecking: atomic.NewBool(false),
processedIndex: 0,
processings: make(map[uint64]*jobStatus),
processingsIndex: make(map[module.JobID]uint64),
}
}
// Start starts consuming the jobs from the job queue.
func (c *Consumer) Start(defaultIndex uint64) error {
c.mu.Lock()
defer c.mu.Unlock()
if c.running {
return nil
}
c.running = true
// on startup, sync with storage for the processed index
// to ensure the consistency
processedIndex, err := c.progress.ProcessedIndex()
if errors.Is(err, storage.ErrNotFound) {
err := c.progress.InitProcessedIndex(defaultIndex)
if errors.Is(err, storage.ErrAlreadyExists) {
return fmt.Errorf("processed index has already been inited, no effect for the second time. default index: %v",
defaultIndex)
}
if err != nil {
return fmt.Errorf("could not init processed index: %w", err)
}
processedIndex = defaultIndex
c.log.Warn().Uint64("processed index", processedIndex).
Msg("processed index not found, initialized.")
} else if err != nil {
return fmt.Errorf("could not read processed index: %w", err)
}
c.processedIndex = processedIndex
c.checkProcessable()
c.log.Info().
Uint64("processed", processedIndex).
Msg("consumer started")
return nil
}
// Stop stops consuming jobs from the job queue.
// It blocks until the existing worker finish processing the job
// Note, it won't stop the existing worker from finishing their job
func (c *Consumer) Stop() {
c.mu.Lock()
c.running = false
// not to use `defer`, otherwise runningJobs.Wait will hold the lock and cause deadlock
c.mu.Unlock()
c.log.Info().Msg("stopping consumer")
c.runningJobs.Wait()
c.log.Info().Msg("consumer stopped")
}
// Size returns number of in-memory jobs that consumer is processing.
func (c *Consumer) Size() uint {
c.mu.Lock()
defer c.mu.Unlock()
return uint(len(c.processings))
}
// LastProcessedIndex returns the last processed job index
func (c *Consumer) LastProcessedIndex() uint64 {
c.mu.Lock()
defer c.mu.Unlock()
return c.processedIndex
}
// NotifyJobIsDone let the consumer know a job has been finished, so that consumer will take
// the next job from the job queue if there are workers available. It returns the last processed job index.
func (c *Consumer) NotifyJobIsDone(jobID module.JobID) uint64 {
c.mu.Lock()
defer c.mu.Unlock()
c.log.Debug().Str("job_id", string(jobID)).Msg("finishing job")
if c.doneJob(jobID) {
c.checkProcessable()
}
return c.processedIndex
}
// Check allows the job publisher to notify the consumer that a new job has been added, so that
// the consumer can check if the job is processable
// since multiple checks at the same time are unnecessary, we could only keep one check by checking.
// an atomic isChecking value.
func (c *Consumer) Check() {
if !c.isChecking.CompareAndSwap(false, true) {
// other process is checking, we could exit and rely on that process to check
// processable jobs
return
}
// still need to lock here, since checkProcessable might update the state vars.
c.mu.Lock()
defer c.mu.Unlock()
c.checkProcessable()
c.isChecking.Store(false)
}
// checkProcessable is a wrap of the `run` function with logging
func (c *Consumer) checkProcessable() {
c.log.Debug().Msg("checking processable jobs")
processingCount, err := c.run()
if err != nil {
c.log.Error().Err(err).Msg("failed to check processables")
return
}
if processingCount > 0 {
c.log.Info().Int64("processing", processingCount).Msg("processing jobs")
} else {
c.log.Debug().Bool("running", c.running).Msg("no job found")
}
}
// run checks if there are processable jobs and process them by giving
// them to the callback functions.
// this function is passive, it won't trigger itself, but can only be
// triggered by either Start or NotifyJobIsDone
func (c *Consumer) run() (int64, error) {
processedFrom := c.processedIndex
processables, processedTo, err := c.processableJobs()
if err != nil {
return 0, fmt.Errorf("could not query processable jobs: %w", err)
}
c.log.Debug().
Uint64("processed_from", processedFrom).
Uint64("processed_to", processedTo).
Int("processables", len(processables)).
Bool("running", c.running).
Msg("running")
for _, indexedJob := range processables {
jobID := indexedJob.job.ID()
c.processingsIndex[jobID] = indexedJob.index
c.processings[indexedJob.index] = &jobStatus{
jobID: jobID,
done: false,
}
c.runningJobs.Add(1)
go func(j *jobAtIndex) {
err := c.worker.Run(j.job)
if err != nil {
c.log.Fatal().Err(err).Msg("could not run the job")
}
c.runningJobs.Done()
}(indexedJob)
}
err = c.progress.SetProcessedIndex(processedTo)
if err != nil {
return 0, fmt.Errorf("could not set processed index %v, %w", processedTo, err)
}
for index := c.processedIndex + 1; index <= processedTo; index++ {
jobStatus, ok := c.processings[index]
if !ok {
continue
}
delete(c.processings, index)
delete(c.processingsIndex, jobStatus.jobID)
}
c.processedIndex = processedTo
return int64(len(processables)), nil
}
func (c *Consumer) processableJobs() ([]*jobAtIndex, uint64, error) {
processables, processedTo, err := processableJobs(
c.jobs,
c.processings,
c.maxProcessing,
c.maxSearchAhead,
c.processedIndex,
)
if err != nil {
return nil, 0, err
}
// if the consumer has been stopped, we allow the existing worker to update the progressed index
// but won't return any new job for processing
if !c.running {
return nil, processedTo, nil
}
return processables, processedTo, nil
}
// processableJobs check the worker's capacity and if sufficient, read
// jobs from the storage, return the processable jobs, and the processed
// index
func processableJobs(jobs module.Jobs, processings map[uint64]*jobStatus, maxProcessing uint64, maxSearchAhead uint64, processedIndex uint64) ([]*jobAtIndex, uint64,
error) {
processables := make([]*jobAtIndex, 0)
// count how many jobs are still processing,
// in order to decide whether to process a new job
processing := uint64(0)
// determine if the consumer should pause processing new jobs because it's too far ahead of
// the lowest in progress index
shouldPause := func(index uint64) bool {
if maxSearchAhead == 0 {
return false
}
return index-processedIndex > maxSearchAhead
}
// if still have processing capacity, find the next processable job
for i := processedIndex + 1; processing < maxProcessing && !shouldPause(i); i++ {
status, ok := processings[i]
// if no worker is processing the next job, try to read it and process
if !ok {
// take one job
job, err := jobs.AtIndex(i)
// if there is no more job at this index, we could stop
if errors.Is(err, storage.ErrNotFound) {
break
}
// exception
if err != nil {
return nil, 0, fmt.Errorf("could not read job at index %v, %w", i, err)
}
processing++
processables = append(processables, &jobAtIndex{
job: job,
index: i,
})
continue
}
// only increment the processing variable when
// the job is not done, meaning still processing
if !status.done {
processing++
continue
}
if i == processedIndex+1 {
processedIndex++
}
}
return processables, processedIndex, nil
}
// doneJob updates the internal state to mark the job has been processed
// return true if the job is changed from processing to finished
// return false if the job is already finished, or removed
func (c *Consumer) doneJob(jobID module.JobID) bool {
// lock
index, ok := c.processingsIndex[jobID]
if !ok {
// job must has been processed
return false
}
status, ok := c.processings[index]
if !ok {
// must be a bug, if went here
c.log.Fatal().Msgf("bug, job (%v) can not be found by index (%v)", jobID, index)
return false
}
if status.done {
// job has been done already
return false
}
status.done = true
return true
}
type jobAtIndex struct {
job module.Job
index uint64
}
type jobStatus struct {
jobID module.JobID
done bool
}