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{}