github.com/cozy/cozy-stack@v0.0.0-20240603063001-31110fa4cae1/model/job/mem_broker.go (about) 1 package job 2 3 import ( 4 "container/list" 5 "context" 6 "errors" 7 "fmt" 8 "sync" 9 "sync/atomic" 10 11 "github.com/cozy/cozy-stack/pkg/config/config" 12 "github.com/cozy/cozy-stack/pkg/limits" 13 "github.com/cozy/cozy-stack/pkg/logger" 14 "github.com/cozy/cozy-stack/pkg/prefixer" 15 multierror "github.com/hashicorp/go-multierror" 16 ) 17 18 type ( 19 // memQueue is a queue in-memory implementation of the Queue interface. 20 memQueue struct { 21 MaxCapacity int 22 Jobs chan *Job 23 closed chan struct{} 24 25 list *list.List 26 run bool 27 jmu sync.RWMutex 28 } 29 30 // memBroker is an in-memory broker implementation of the Broker interface. 31 memBroker struct { 32 queues map[string]*memQueue 33 workers []*Worker 34 workersTypes []string 35 running uint32 36 } 37 ) 38 39 // newMemQueue creates and a new in-memory queue. 40 func newMemQueue(workerType string) *memQueue { 41 return &memQueue{ 42 list: list.New(), 43 Jobs: make(chan *Job), 44 closed: make(chan struct{}), 45 } 46 } 47 48 // Enqueue into the queue 49 func (q *memQueue) Enqueue(job *Job) error { 50 q.jmu.Lock() 51 defer q.jmu.Unlock() 52 q.list.PushBack(job.Clone()) 53 if !q.run { 54 q.run = true 55 go q.send() 56 } 57 return nil 58 } 59 60 func (q *memQueue) send() { 61 for { 62 q.jmu.Lock() 63 e := q.list.Front() 64 if e == nil || !q.run { 65 q.run = false 66 q.jmu.Unlock() 67 return 68 } 69 q.list.Remove(e) 70 q.jmu.Unlock() 71 select { 72 case <-q.closed: 73 return 74 case q.Jobs <- e.Value.(*Job): 75 } 76 } 77 } 78 79 func (q *memQueue) close() { 80 q.jmu.Lock() 81 defer q.jmu.Unlock() 82 if !q.run { 83 return 84 } 85 q.run = false 86 go func() { q.closed <- struct{}{} }() 87 } 88 89 // Len returns the length of the queue 90 func (q *memQueue) Len() int { 91 q.jmu.RLock() 92 defer q.jmu.RUnlock() 93 return q.list.Len() 94 } 95 96 // NewMemBroker creates a new in-memory broker system. 97 // 98 // The in-memory implementation of the job system has the specifity that 99 // workers are actually launched by the broker at its creation. 100 func NewMemBroker() Broker { 101 return &memBroker{ 102 queues: make(map[string]*memQueue), 103 } 104 } 105 106 func (b *memBroker) StartWorkers(ws WorkersList) error { 107 if !atomic.CompareAndSwapUint32(&b.running, 0, 1) { 108 return ErrClosed 109 } 110 111 for _, conf := range ws { 112 b.workersTypes = append(b.workersTypes, conf.WorkerType) 113 if conf.Concurrency <= 0 { 114 continue 115 } 116 q := newMemQueue(conf.WorkerType) 117 w := NewWorker(conf) 118 b.queues[conf.WorkerType] = q 119 b.workers = append(b.workers, w) 120 if err := w.Start(q.Jobs); err != nil { 121 return err 122 } 123 } 124 125 if len(b.workers) > 0 { 126 joblog.Infof("Started in-memory broker for %d workers type", len(b.workers)) 127 } 128 129 // XXX for retro-compat 130 if slots := config.GetConfig().Jobs.NbWorkers; len(b.workers) > 0 && slots > 0 { 131 joblog.Warnf("Limiting the number of total concurrent workers to %d", slots) 132 joblog.Warnf("Please update your configuration file to avoid a hard limit") 133 setNbSlots(slots) 134 } 135 136 return nil 137 } 138 139 func (b *memBroker) ShutdownWorkers(ctx context.Context) error { 140 if !atomic.CompareAndSwapUint32(&b.running, 1, 0) { 141 return ErrClosed 142 } 143 if len(b.workers) == 0 { 144 return nil 145 } 146 147 fmt.Print(" shutting down in-memory broker...") 148 149 for _, q := range b.queues { 150 q.close() 151 } 152 153 errs := make(chan error) 154 for _, w := range b.workers { 155 go func(w *Worker) { errs <- w.Shutdown(ctx) }(w) 156 } 157 var errm error 158 for i := 0; i < len(b.workers); i++ { 159 if err := <-errs; err != nil { 160 errm = multierror.Append(errm, err) 161 } 162 } 163 164 if errm != nil { 165 fmt.Println("failed:", errm) 166 } else { 167 fmt.Println("ok.") 168 } 169 return errm 170 } 171 172 // PushJob will produce a new Job with the given options and enqueue the job in 173 // the proper queue. 174 func (b *memBroker) PushJob(db prefixer.Prefixer, req *JobRequest) (*Job, error) { 175 if atomic.LoadUint32(&b.running) == 0 { 176 return nil, ErrClosed 177 } 178 179 workerType := req.WorkerType 180 var worker *Worker 181 for _, w := range b.workers { 182 if w.Type == workerType { 183 worker = w 184 break 185 } 186 } 187 if worker == nil && workerType != "client" { 188 return nil, ErrUnknownWorker 189 } 190 191 // Check for limits 192 ct, err := GetCounterTypeFromWorkerType(req.WorkerType) 193 if err == nil { 194 err := config.GetRateLimiter().CheckRateLimit(db, ct) 195 if errors.Is(err, limits.ErrRateLimitReached) { 196 joblog.WithFields(logger.Fields{ 197 "worker_type": req.WorkerType, 198 "instance": db.DomainName(), 199 }).Warn(err.Error()) 200 return nil, err 201 } 202 if limits.IsLimitReachedOrExceeded(err) { 203 return nil, err 204 } 205 } 206 207 job := NewJob(db, req) 208 if worker != nil && worker.Conf.BeforeHook != nil { 209 ok, err := worker.Conf.BeforeHook(job) 210 if err != nil { 211 return nil, err 212 } 213 if !ok { 214 return job, nil 215 } 216 } 217 218 if err := job.Create(); err != nil { 219 return nil, err 220 } 221 222 // For client jobs, we don't need to enqueue the job. 223 if workerType == "client" { 224 return job, nil 225 } 226 227 q := b.queues[workerType] 228 if err := q.Enqueue(job); err != nil { 229 return nil, err 230 } 231 return job, nil 232 } 233 234 // WorkerQueueLen returns the size of the number of elements in queue of the 235 // specified worker type. 236 func (b *memBroker) WorkerQueueLen(workerType string) (int, error) { 237 q, ok := b.queues[workerType] 238 if !ok { 239 return 0, ErrUnknownWorker 240 } 241 return q.Len(), nil 242 } 243 244 func (b *memBroker) WorkerIsReserved(workerType string) (bool, error) { 245 for _, w := range b.workers { 246 if w.Type == workerType { 247 return w.Conf.Reserved, nil 248 } 249 } 250 return false, ErrUnknownWorker 251 } 252 253 func (b *memBroker) WorkersTypes() []string { 254 return b.workersTypes 255 } 256 257 var _ Broker = &memBroker{}