-
Notifications
You must be signed in to change notification settings - Fork 133
/
jobdb.go
196 lines (178 loc) · 5.34 KB
/
jobdb.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
package jobdb
import (
"sync"
"github.com/benbjohnson/immutable"
"github.com/google/uuid"
"github.com/pkg/errors"
"golang.org/x/exp/maps"
)
var emptyList = immutable.NewSortedSet[*Job](JobPriorityComparer{})
type JobDb struct {
jobsById map[string]*Job
jobsByRunId map[uuid.UUID]string
jobsByQueue map[string]immutable.SortedSet[*Job]
copyMutex sync.Mutex
writerMutex sync.Mutex
}
func NewJobDb() *JobDb {
return &JobDb{
jobsById: map[string]*Job{},
jobsByRunId: map[uuid.UUID]string{},
jobsByQueue: map[string]immutable.SortedSet[*Job]{},
copyMutex: sync.Mutex{},
}
}
// Upsert will insert the given jobs if they don't already exist or update the if they do
func (jobDb *JobDb) Upsert(txn *Txn, jobs []*Job) error {
if err := jobDb.checkWritableTransaction(txn); err != nil {
return err
}
for _, job := range jobs {
existingJob := txn.jobsById[job.id]
if existingJob != nil {
existingQueue, ok := txn.jobsByQueue[existingJob.queue]
if ok {
txn.jobsByQueue[existingJob.queue] = existingQueue.Delete(existingJob)
}
}
txn.jobsById[job.id] = job
for _, run := range job.runsById {
txn.jobsByRunId[run.id] = job.id
}
if job.Queued() {
newQueue, ok := txn.jobsByQueue[job.queue]
if !ok {
q := emptyList
newQueue = q
}
newQueue = newQueue.Add(job)
txn.jobsByQueue[job.queue] = newQueue
}
}
return nil
}
// GetById returns the job with the given Id or nil if no such job exists
// The Job returned by this function *must not* be subsequently modified
func (jobDb *JobDb) GetById(txn *Txn, id string) *Job {
return txn.jobsById[id]
}
// GetByRunId returns the job with the given run id or nil if no such job exists
// The Job returned by this function *must not* be subsequently modified
func (jobDb *JobDb) GetByRunId(txn *Txn, runId uuid.UUID) *Job {
jobId := txn.jobsByRunId[runId]
return jobDb.GetById(txn, jobId)
}
// HasQueuedJobs returns true if the queue has any jobs in the running state or false otherwise
func (jobDb *JobDb) HasQueuedJobs(txn *Txn, queue string) bool {
queuedJobs, ok := txn.jobsByQueue[queue]
if !ok {
return false
}
return queuedJobs.Len() > 0
}
// QueuedJobs returns true if the queue has any jobs in the running state or false otherwise
func (jobDb *JobDb) QueuedJobs(txn *Txn, queue string) *immutable.SortedSetIterator[*Job] {
jobQueue, ok := txn.jobsByQueue[queue]
if ok {
return jobQueue.Iterator()
} else {
return emptyList.Iterator()
}
}
// GetAll returns all jobs in the database.
// The Jobs returned by this function *must not* be subsequently modified
func (jobDb *JobDb) GetAll(txn *Txn) []*Job {
return maps.Values(txn.jobsById)
}
// BatchDelete removes the jobs with the given ids from the database. Any ids that are not in the database will be
// ignored
func (jobDb *JobDb) BatchDelete(txn *Txn, ids []string) error {
if err := jobDb.checkWritableTransaction(txn); err != nil {
return err
}
for _, id := range ids {
job, present := txn.jobsById[id]
if present {
delete(txn.jobsById, id)
for _, run := range job.runsById {
delete(txn.jobsByRunId, run.id)
}
queue, ok := txn.jobsByQueue[job.queue]
if ok {
newQueue := queue.Delete(job)
txn.jobsByQueue[job.queue] = newQueue
}
}
}
return nil
}
func (jobDb *JobDb) checkWritableTransaction(txn *Txn) error {
if txn.readOnly {
return errors.New("Cannot write using a read only transaction")
}
if !txn.active {
return errors.New("Cannot write using an inactive transaction")
}
return nil
}
// ReadTxn returns a read-only transaction.
// Multiple read-only transactions can access the db concurrently
func (jobDb *JobDb) ReadTxn() *Txn {
jobDb.copyMutex.Lock()
defer jobDb.copyMutex.Unlock()
return &Txn{
readOnly: true,
jobsById: jobDb.jobsById,
jobsByRunId: jobDb.jobsByRunId,
jobsByQueue: jobDb.jobsByQueue,
active: true,
jobDb: jobDb,
}
}
// WriteTxn returns a writeable transaction.
// Only a single write transaction may access the db at any given time so note that this function will block until
// any outstanding write transactions have been committed or aborted
func (jobDb *JobDb) WriteTxn() *Txn {
jobDb.writerMutex.Lock()
jobDb.copyMutex.Lock()
defer jobDb.copyMutex.Unlock()
return &Txn{
readOnly: false,
jobsById: maps.Clone(jobDb.jobsById),
jobsByRunId: maps.Clone(jobDb.jobsByRunId),
jobsByQueue: maps.Clone(jobDb.jobsByQueue),
active: true,
jobDb: jobDb,
}
}
// Txn is a JobDb Transaction. Transactions provide a consistent view of the database, allowing readers to
// perform multiple actions without the database changing from underneath them.
// Write transactions also allow callers to perform write operations that will not be visible to other users
// until the transaction is committed.
type Txn struct {
readOnly bool
jobsById map[string]*Job
jobsByRunId map[uuid.UUID]string
jobsByQueue map[string]immutable.SortedSet[*Job]
jobDb *JobDb
active bool
}
func (txn *Txn) Commit() {
if txn.readOnly || !txn.active {
return
}
txn.jobDb.copyMutex.Lock()
defer txn.jobDb.copyMutex.Unlock()
defer txn.jobDb.writerMutex.Unlock()
txn.jobDb.jobsById = txn.jobsById
txn.jobDb.jobsByRunId = txn.jobsByRunId
txn.jobDb.jobsByQueue = txn.jobsByQueue
txn.active = false
}
func (txn *Txn) Abort() {
if txn.readOnly || !txn.active {
return
}
txn.active = false
txn.jobDb.writerMutex.Unlock()
}