Moves files

internal/worker/doc.go

@@ -0,0 +1,2 @@
+// Package worker implements a worker pool.
+package worker

internal/worker/pool.go

@@ -0,0 +1,101 @@
+package worker
+
+import "context"
+
+// Job is one unit of work. It is given to a Func, and the returned result and
+// error are stored in Result and Error.
+type Job struct {
+	Data   interface{}
+	Result interface{}
+	Error  error
+}
+
+// Func does the actual work within a Pool.
+type Func func(ctx context.Context, job Job) (result interface{}, err error)
+
+// Pool implements a worker pool.
+type Pool struct {
+	f     Func
+	jobCh <-chan Job
+	resCh chan<- Job
+
+	numWorkers     int
+	workersExit    chan struct{}
+	allWorkersDone chan struct{}
+}
+
+// New returns a new worker pool with n goroutines, each running the function
+// f. The workers are started immediately.
+func New(ctx context.Context, n int, f Func, jobChan <-chan Job, resultChan chan<- Job) *Pool {
+	p := &Pool{
+		f:              f,
+		workersExit:    make(chan struct{}),
+		allWorkersDone: make(chan struct{}),
+		numWorkers:     n,
+		jobCh:          jobChan,
+		resCh:          resultChan,
+	}
+
+	for i := 0; i < n; i++ {
+		go p.runWorker(ctx, i)
+	}
+
+	go p.waitForExit()
+
+	return p
+}
+
+// waitForExit receives from p.workersExit until all worker functions have
+// exited, then closes the result channel.
+func (p *Pool) waitForExit() {
+	n := p.numWorkers
+	for n > 0 {
+		<-p.workersExit
+		n--
+	}
+	close(p.allWorkersDone)
+	close(p.resCh)
+}
+
+// runWorker runs a worker function.
+func (p *Pool) runWorker(ctx context.Context, numWorker int) {
+	defer func() {
+		p.workersExit <- struct{}{}
+	}()
+
+	var (
+		// enable the input channel when starting up a new goroutine
+		inCh = p.jobCh
+		// but do not enable the output channel until we have a result
+		outCh chan<- Job
+
+		job Job
+		ok  bool
+	)
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+
+		case job, ok = <-inCh:
+			if !ok {
+				return
+			}
+
+			job.Result, job.Error = p.f(ctx, job)
+			inCh = nil
+			outCh = p.resCh
+
+		case outCh <- job:
+			outCh = nil
+			inCh = p.jobCh
+		}
+	}
+}
+
+// Wait waits for all worker goroutines to terminate, afterwards the output
+// channel is closed.
+func (p *Pool) Wait() {
+	<-p.allWorkersDone
+}

internal/worker/pool_test.go

@@ -0,0 +1,92 @@
+package worker_test
+
+import (
+	"context"
+	"testing"
+
+	"restic/errors"
+
+	"restic/worker"
+)
+
+const concurrency = 10
+
+var errTooLarge = errors.New("too large")
+
+func square(ctx context.Context, job worker.Job) (interface{}, error) {
+	n := job.Data.(int)
+	if n > 2000 {
+		return nil, errTooLarge
+	}
+	return n * n, nil
+}
+
+func newBufferedPool(ctx context.Context, bufsize int, n int, f worker.Func) (chan worker.Job, chan worker.Job, *worker.Pool) {
+	inCh := make(chan worker.Job, bufsize)
+	outCh := make(chan worker.Job, bufsize)
+
+	return inCh, outCh, worker.New(ctx, n, f, inCh, outCh)
+}
+
+func TestPool(t *testing.T) {
+	inCh, outCh, p := newBufferedPool(context.TODO(), 200, concurrency, square)
+
+	for i := 0; i < 150; i++ {
+		inCh <- worker.Job{Data: i}
+	}
+
+	close(inCh)
+	p.Wait()
+
+	for res := range outCh {
+		if res.Error != nil {
+			t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
+			continue
+		}
+
+		n := res.Data.(int)
+		m := res.Result.(int)
+
+		if m != n*n {
+			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
+		}
+	}
+}
+
+func TestPoolErrors(t *testing.T) {
+	inCh, outCh, p := newBufferedPool(context.TODO(), 200, concurrency, square)
+
+	for i := 0; i < 150; i++ {
+		inCh <- worker.Job{Data: i + 1900}
+	}
+
+	close(inCh)
+	p.Wait()
+
+	for res := range outCh {
+		n := res.Data.(int)
+
+		if n > 2000 {
+			if res.Error == nil {
+				t.Errorf("expected error not found, result is %v", res)
+				continue
+			}
+
+			if res.Error != errTooLarge {
+				t.Errorf("unexpected error found, result is %v", res)
+			}
+
+			continue
+		} else {
+			if res.Error != nil {
+				t.Errorf("unexpected error for job %v received: %v", res.Data, res.Error)
+				continue
+			}
+		}
+
+		m := res.Result.(int)
+		if m != n*n {
+			t.Errorf("wrong value for job %d returned: want %d, got %d", n, n*n, m)
+		}
+	}
+}