github.com/supabase/cli@v1.168.1/internal/utils/queue.go (about)

     1  package utils
     2  
     3  import (
     4  	"sync"
     5  
     6  	"github.com/go-errors/errors"
     7  )
     8  
     9  // JobQueue implements a background job processor using a single channel and wait group.
    10  //
    11  // The channel is initialised with a maximum number of concurrent workers. Adding a new job
    12  // consumes a worker from the channel. After finishing a job, the worker is added back to
    13  // the channel. When all workers are consumed, adding new job will block.
    14  //
    15  // Example usage:
    16  //
    17  //	jq := NewJobQueue(5)
    18  //	err := jq.Put(func() error {
    19  //		return nil
    20  //	})
    21  //	errors.Join(err, jq.Collect())
    22  type JobQueue struct {
    23  	wg    sync.WaitGroup
    24  	errCh chan error
    25  }
    26  
    27  func NewJobQueue(maxConcurrency uint) *JobQueue {
    28  	q := &JobQueue{
    29  		errCh: make(chan error, maxConcurrency),
    30  	}
    31  	for i := 0; i < cap(q.errCh); i++ {
    32  		q.errCh <- nil
    33  	}
    34  	return q
    35  }
    36  
    37  // Put runs a job in the background, returning any error from the previous job.
    38  func (q *JobQueue) Put(job func() error) error {
    39  	err := <-q.errCh
    40  	q.wg.Add(1)
    41  	go func() {
    42  		defer q.wg.Done()
    43  		q.errCh <- job()
    44  	}()
    45  	return err
    46  }
    47  
    48  // Collect waits for all jobs to finish, returning any errors.
    49  func (q *JobQueue) Collect() error {
    50  	q.wg.Wait()
    51  	var err []error
    52  	for i := 0; i < cap(q.errCh); i++ {
    53  		err = append(err, <-q.errCh)
    54  		q.errCh <- nil
    55  	}
    56  	return errors.Join(err...)
    57  }