-
Notifications
You must be signed in to change notification settings - Fork 571
/
worker.go
102 lines (82 loc) · 2.19 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
package directory
import (
"context"
"errors"
"sync"
"github.com/deepfence/ThreatMapper/deepfence_utils/utils"
"github.com/hibiken/asynq"
)
const (
maxSize = 500 * 1024 * 1024 // 500 MB
)
var ErrExhaustedResources = errors.New("exhausted worker resources")
type asynqClients struct {
client *asynq.Client
inspector *asynq.Inspector
}
type WorkEnqueuer struct {
clients asynqClients
}
var workerClientsPool sync.Map
func init() {
workerClientsPool = sync.Map{}
}
func newAsynqClient(endpoints DBConfigs) (*asynqClients, error) {
if endpoints.Redis == nil {
return nil, errors.New("no defined Redis config")
}
redisCfg := asynq.RedisClientOpt{Addr: endpoints.Redis.Endpoint}
return &asynqClients{
client: asynq.NewClient(redisCfg),
inspector: asynq.NewInspector(redisCfg),
}, nil
}
func (ws WorkEnqueuer) Enqueue(taskEnum string, data []byte, opts ...asynq.Option) error {
client := ws.clients.client
inspector := ws.clients.inspector
qs, err := inspector.Queues()
if err != nil {
return err
}
size := 0
for _, q := range qs {
res, err := inspector.GetQueueInfo(q)
if err != nil {
continue
}
size += res.Size
}
if size >= maxSize {
return ErrExhaustedResources
}
_, err = client.Enqueue(asynq.NewTask(taskEnum, data), opts...)
return err
}
// utility func to use when task name is used as task id to ensure unique task is enqueued
func (ws WorkEnqueuer) EnqueueUnique(taskEnum string, data []byte, opts ...asynq.Option) error {
opts = append(opts, asynq.TaskID(taskEnum))
return ws.Enqueue(taskEnum, data, opts...)
}
func (ws WorkEnqueuer) DeleteAllArchivedTasks() (int, []error) {
var deletedTasksCount int
var errs []error
for _, queue := range utils.AsynqQueues {
count, err := ws.clients.inspector.DeleteAllArchivedTasks(queue)
if err != nil {
errs = append(errs, err)
continue
}
deletedTasksCount += count
}
return deletedTasksCount, errs
}
func (ws WorkEnqueuer) Inspector() *asynq.Inspector {
return ws.clients.inspector
}
func Worker(ctx context.Context) (WorkEnqueuer, error) {
client, err := getClient(ctx, &workerClientsPool, newAsynqClient)
if err != nil {
return WorkEnqueuer{}, err
}
return WorkEnqueuer{clients: *client}, err
}