-
Notifications
You must be signed in to change notification settings - Fork 0
/
sequnteialPerIdWorkerQueue.go
120 lines (104 loc) · 2.9 KB
/
sequnteialPerIdWorkerQueue.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
package worker
import (
"context"
"sync"
)
type SequentialPerIdJob interface {
Job
Id() string
}
// SequentialPerIdWorkerQueue is a worker queue, that guarantees a single threaded execution for all jobs with the
// same comparable result.
type SequentialPerIdWorkerQueue interface {
// Do register a DistinctiveJob for being processed by a worker when the worker is ready, the
// job will be scheduled after other jobs with same id and will never run in parallel with them.
Do(job SequentialPerIdJob)
// Close the jobs channel so no new jobs can be queued. It is not necessary to close the channel for cleanup.
Close()
// GetQueueSize return the size of all waiting jobs in the queue
GetQueueSize() int
}
type sequentialPerIdWorkerQueue struct {
sync.RWMutex
sequentialQueue map[string]chan SequentialPerIdJob
sequentialQueueBufferSize int
queue chan SequentialPerIdJob
cancel context.CancelFunc
}
func NewSequentialPerIdWorkerQueue(ctx context.Context, workerCount int, bufferSize int, sequentialQueueBufferSize int) SequentialPerIdWorkerQueue {
instance := &sequentialPerIdWorkerQueue{
sequentialQueue: make(map[string]chan SequentialPerIdJob),
queue: make(chan SequentialPerIdJob, bufferSize),
sequentialQueueBufferSize: sequentialQueueBufferSize,
}
var innerCtx context.Context
innerCtx, instance.cancel = context.WithCancel(ctx)
for i := 0; i < workerCount; i++ {
_ctx := context.WithValue(innerCtx, workerId, i)
go instance.worker(_ctx)
}
return instance
}
func (q *sequentialPerIdWorkerQueue) Do(job SequentialPerIdJob) {
q.queue <- job
}
func (q *sequentialPerIdWorkerQueue) registerJob(ctx context.Context, originJob SequentialPerIdJob) {
// check if there are jobs in the sequentialQueue to process first
q.Lock()
sq, ok := q.sequentialQueue[originJob.Id()]
if !ok {
sq = make(chan SequentialPerIdJob, q.sequentialQueueBufferSize)
q.sequentialQueue[originJob.Id()] = sq
q.Unlock()
select {
case sq <- originJob:
break
case <-ctx.Done():
return
}
} else {
// add the new job to the sequential queue and let it being processed by the worker that created the cahnnel
defer q.Unlock()
select {
case sq <- originJob:
return
case <-ctx.Done():
return
}
}
// work the sequential queue
for job := range sq {
if ctx.Err() != nil {
return
}
job.Run(ctx)
if len(sq) == 0 {
q.Lock()
if len(sq) > 0 {
q.Unlock()
continue
}
delete(q.sequentialQueue, job.Id())
q.Unlock()
break
}
}
}
func (q *sequentialPerIdWorkerQueue) worker(ctx context.Context) {
for job := range q.queue {
q.registerJob(ctx, job)
}
}
func (q *sequentialPerIdWorkerQueue) Close() {
close(q.queue)
q.cancel()
}
func (q *sequentialPerIdWorkerQueue) GetQueueSize() int {
l := len(q.queue)
q.RLock()
defer q.RUnlock()
for _, sq := range q.sequentialQueue {
l += len(sq)
}
return l
}