GODT-1779: Remove go-imap

internal/pool/pool.go

@@ -0,0 +1,177 @@
+package pool
+
+import (
+	"context"
+	"errors"
+	"sync"
+
+	"github.com/ProtonMail/gluon/queue"
+)
+
+// ErrJobCancelled indicates the job was cancelled.
+var ErrJobCancelled = errors.New("Job cancelled by surrounding context")
+
+// Pool is a worker pool that handles input of type In and returns results of type Out.
+type Pool[In comparable, Out any] struct {
+	queue *queue.QueuedChannel[*Job[In, Out]]
+	size  int
+}
+
+// DoneFunc must be called to free up pool resources.
+type DoneFunc func()
+
+// New returns a new pool.
+func New[In comparable, Out any](size int, work func(context.Context, In) (Out, error)) *Pool[In, Out] {
+	queue := queue.NewQueuedChannel[*Job[In, Out]](0, 0)
+
+	for i := 0; i < size; i++ {
+		go func() {
+			for job := range queue.GetChannel() {
+				select {
+				case <-job.ctx.Done():
+					job.postFailure(ErrJobCancelled)
+
+				default:
+					res, err := work(job.ctx, job.req)
+					if err != nil {
+						job.postFailure(err)
+					} else {
+						job.postSuccess(res)
+					}
+
+					job.waitDone()
+				}
+			}
+		}()
+	}
+
+	return &Pool[In, Out]{
+		queue: queue,
+		size:  size,
+	}
+}
+
+// NewJob submits a job to the pool. It returns a job handle and a DoneFunc.
+// The job handle allows the job result to be obtained. The DoneFunc is used to mark the job as done,
+// which frees up the worker in the pool for reuse.
+func (pool *Pool[In, Out]) NewJob(ctx context.Context, req In) (*Job[In, Out], DoneFunc) {
+	job := newJob[In, Out](ctx, req)
+
+	pool.queue.Enqueue(job)
+
+	return job, func() { close(job.done) }
+}
+
+// Process submits jobs to the pool. The callback provides access to the result, or an error if one occurred.
+func (pool *Pool[In, Out]) Process(ctx context.Context, reqs []In, fn func(In, Out, error) error) error {
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	var (
+		wg      sync.WaitGroup
+		errList []error
+		lock    sync.Mutex
+	)
+
+	for _, req := range reqs {
+		req := req
+
+		wg.Add(1)
+
+		go func() {
+			defer wg.Done()
+
+			job, done := pool.NewJob(ctx, req)
+			defer done()
+
+			res, err := job.Result()
+
+			if err := fn(req, res, err); err != nil {
+				lock.Lock()
+				defer lock.Unlock()
+
+				// Cancel ongoing jobs.
+				cancel()
+
+				// Collect the error.
+				errList = append(errList, err)
+			}
+		}()
+	}
+
+	wg.Wait()
+
+	// TODO: Join the errors somehow?
+	if len(errList) > 0 {
+		return errList[0]
+	}
+
+	return nil
+}
+
+// ProcessAll submits jobs to the pool. All results are returned once available.
+func (pool *Pool[In, Out]) ProcessAll(ctx context.Context, reqs []In) (map[In]Out, error) {
+	var (
+		data = make(map[In]Out)
+		lock = sync.Mutex{}
+	)
+
+	if err := pool.Process(ctx, reqs, func(req In, res Out, err error) error {
+		if err != nil {
+			return err
+		}
+
+		lock.Lock()
+		defer lock.Unlock()
+
+		data[req] = res
+
+		return nil
+	}); err != nil {
+		return nil, err
+	}
+
+	return data, nil
+}
+
+func (pool *Pool[In, Out]) Done() {
+	pool.queue.Close()
+}
+
+type Job[In, Out any] struct {
+	ctx context.Context
+	req In
+
+	res chan Out
+	err chan error
+
+	done chan struct{}
+}
+
+func newJob[In, Out any](ctx context.Context, req In) *Job[In, Out] {
+	return &Job[In, Out]{
+		ctx:  ctx,
+		req:  req,
+		res:  make(chan Out),
+		err:  make(chan error),
+		done: make(chan struct{}),
+	}
+}
+
+func (job *Job[In, Out]) Result() (Out, error) {
+	return <-job.res, <-job.err
+}
+
+func (job *Job[In, Out]) postSuccess(res Out) {
+	close(job.err)
+	job.res <- res
+}
+
+func (job *Job[In, Out]) postFailure(err error) {
+	close(job.res)
+	job.err <- err
+}
+
+func (job *Job[In, Out]) waitDone() {
+	<-job.done
+}