refactor: change (almost) everything

Author: guilhem

workers.go

@@ -3,247 +3,199 @@ package workers
 import (
 	"context"
 	"errors"
-	"os"
-	"os/signal"
 	"sync"
-	"syscall"
 	"time"
 )
 
-var defaultWatchSignals = []os.Signal{syscall.SIGINT, syscall.SIGTERM, syscall.SIGKILL}
-
-// Worker Contains the work function. Allows an input and output to a channel or another worker for pipeline work.
-// Return nil if you want the Runner to continue otherwise any error will cause the Runner to shutdown and return the
-// error.
+// Worker Contains the work function.
+// Work() get input and could put in outChan for followers.
+// ⚠️ outChan could be closed if follower is stoped before producer.
+// error returned can be process by afterFunc but will be ignored by default.
 type Worker interface {
-	Work(in interface{}, out chan<- interface{}) error
-}
-
-// Runner Handles the running the Worker logic.
-type Runner interface {
-	BeforeFunc(func(ctx context.Context) error) Runner
-	AfterFunc(func(ctx context.Context, err error) error) Runner
-	SetDeadline(t time.Time) Runner
-	SetTimeout(duration time.Duration) Runner
-	SetFollower()
-	Send(in interface{})
-	InFrom(w ...Runner) Runner
-	SetOut(chan interface{})
-	Start() Runner
-	Stop() chan error
-	Wait() error
+	Work(ctx context.Context, in interface{}, out chan<- interface{}) error
 }
 
-type runner struct {
-	ctx        context.Context
-	cancel     context.CancelFunc
-	inChan     chan interface{}
-	outChan    chan interface{}
-	errChan    chan error
-	signalChan chan os.Signal
-	limiter    chan struct{}
+type Runner struct {
+	ctx         context.Context
+	cancel      context.CancelFunc
+	inputCtx    context.Context
+	inputCancel context.CancelFunc
+	inChan      chan interface{}
+	outChan     chan interface{}
+	limiter     chan struct{}
 
 	afterFunc  func(ctx context.Context, err error) error
-	workFunc   func(in interface{}, out chan<- interface{}) error
+	workFunc   func(ctx context.Context, in interface{}, out chan<- interface{}) error
 	beforeFunc func(ctx context.Context) error
 
-	timeout  time.Duration
-	deadline time.Duration
-
-	isLeader   bool
-	stopCalled bool
+	timeout time.Duration
 
 	numWorkers int64
-	lock       *sync.RWMutex
-	wg         *sync.WaitGroup
-	done       *sync.Once
-	once       *sync.Once
+	started    *sync.Once
+	done       chan struct{}
 }
 
 // NewRunner Factory function for a new Runner.  The Runner will handle running the workers logic.
-func NewRunner(ctx context.Context, w Worker, numWorkers int64) Runner {
-	var runnerCtx, runnerCancel = context.WithCancel(ctx)
-	var runner = &runner{
-		ctx:        runnerCtx,
-		cancel:     runnerCancel,
-		inChan:     make(chan interface{}, numWorkers),
-		outChan:    nil,
-		errChan:    make(chan error, 1),
-		signalChan: make(chan os.Signal, 1),
-		limiter:    make(chan struct{}, numWorkers),
-		afterFunc:  func(ctx context.Context, err error) error { return err },
-		workFunc:   w.Work,
-		beforeFunc: func(ctx context.Context) error { return nil },
-		numWorkers: numWorkers,
-		isLeader:   true,
-		lock:       new(sync.RWMutex),
-		wg:         new(sync.WaitGroup),
-		once:       new(sync.Once),
-		done:       new(sync.Once),
+func NewRunner(ctx context.Context, w Worker, numWorkers int64, buffer int64) *Runner {
+	runnerCtx, runnerCancel := context.WithCancel(ctx)
+	inputCtx, inputCancel := context.WithCancel(runnerCtx)
+
+	runner := &Runner{
+		ctx:         runnerCtx,
+		cancel:      runnerCancel,
+		inputCtx:    inputCtx,
+		inputCancel: inputCancel,
+		inChan:      make(chan interface{}, buffer),
+		outChan:     nil,
+		limiter:     make(chan struct{}, numWorkers),
+		afterFunc:   func(ctx context.Context, err error) error { return nil },
+		workFunc:    w.Work,
+		beforeFunc:  func(ctx context.Context) error { return nil },
+		numWorkers:  numWorkers,
+		started:     new(sync.Once),
+		done:        make(chan struct{}),
 	}
-	runner.waitForSignal(defaultWatchSignals...)
 	return runner
 }
 
-// Send Send an object to the worker for processing.
-func (r *runner) Send(in interface{}) {
+var ErrInputClosed = errors.New("input closed")
+
+// Send Send an object to the worker for processing if context is not Done.
+func (r *Runner) Send(in interface{}) error {
 	select {
-	case <-r.ctx.Done():
-		return
+	case <-r.inputCtx.Done():
+		return ErrInputClosed
 	case r.inChan <- in:
 	}
+	return nil
 }
 
 // InFrom Set a worker to accept output from another worker(s).
-func (r *runner) InFrom(w ...Runner) Runner {
-	r.SetFollower()
+func (r *Runner) InFrom(w ...*Runner) *Runner {
 	for _, wr := range w {
-		wr.SetOut(r.inChan)
+		// in := make(chan interface{})
+		// go func(in chan interface{}) {
+		// 	for msg := range in {
+		// 		if err := r.Send(msg); err != nil {
+		// 			return
+		// 		}
+		// 	}
+		// }(in)
+		wr.SetOut(r.inChan) // nolint
 	}
 	return r
 }
 
-// SetFollower Sets the worker as a follower and does not need to close it's in channel.
-func (r *runner) SetFollower() {
-	r.lock.Lock()
-	r.isLeader = false
-	r.lock.Unlock()
-}
-
-// Start Starts the worker on processing.
-func (r *runner) Start() Runner {
-	r.startWork()
-	return r
+// Start execute beforeFunc and launch worker processing.
+func (r *Runner) Start() error {
+	r.started.Do(func() {
+		if err := r.beforeFunc(r.ctx); err == nil {
+			go r.work()
+		}
+	})
+	return nil
 }
 
 // BeforeFunc Function to be run before worker starts processing.
-func (r *runner) BeforeFunc(f func(ctx context.Context) error) Runner {
+func (r *Runner) BeforeFunc(f func(ctx context.Context) error) *Runner {
 	r.beforeFunc = f
 	return r
 }
 
 // AfterFunc Function to be run after worker has stopped.
-func (r *runner) AfterFunc(f func(ctx context.Context, err error) error) Runner {
+// It can be used for logging and error management.
+// input can be retreive with context value:
+//   ctx.Value(workers.InputKey{})
+// ⚠️ If an error is returned it stop Runner execution.
+func (r *Runner) AfterFunc(f func(ctx context.Context, err error) error) *Runner {
 	r.afterFunc = f
 	return r
 }
 
+var ErrOutAlready = errors.New("out already set")
+
 // SetOut Allows the setting of a workers out channel, if not already set.
-func (r *runner) SetOut(c chan interface{}) {
+func (r *Runner) SetOut(c chan interface{}) error {
 	if r.outChan != nil {
-		return
+		return ErrOutAlready
 	}
 	r.outChan = c
+	return nil
 }
 
-// SetDeadline allows a time to be set when the workers should stop.
-// Deadline needs to be handled by the IsDone method.
-func (r *runner) SetDeadline(t time.Time) Runner {
-	r.lock.Lock()
-	defer r.lock.Unlock()
+// SetDeadline allows a time to be set when the Runner should stop.
+// ⚠️ Should only be called before Start
+func (r *Runner) SetDeadline(t time.Time) *Runner {
 	r.ctx, r.cancel = context.WithDeadline(r.ctx, t)
 	return r
 }
 
-// SetTimeout allows a time duration to be set when the workers should stop.
-// Timeout needs to be handled by the IsDone method.
-func (r *runner) SetTimeout(duration time.Duration) Runner {
-	r.lock.Lock()
-	defer r.lock.Unlock()
+// SetWorkerTimeout allows a time duration to be set when the workers should stop.
+// ⚠️ Should only be called before Start
+func (r *Runner) SetWorkerTimeout(duration time.Duration) *Runner {
 	r.timeout = duration
 	return r
 }
 
-// Wait calls stop on workers and waits for the channel to drain.
-// !!Should only be called when certain nothing will send to worker.
-func (r *runner) Wait() error {
-	r.waitForDrain()
-	if err := <-r.Stop(); err != nil && !errors.Is(err, context.Canceled) {
-		return err
+// Wait close the input channel and waits it to drain and process.
+func (r *Runner) Wait() *Runner {
+	if r.inputCtx.Err() == nil {
+		r.inputCancel()
+		close(r.inChan)
 	}
-	return nil
-}
 
-// Stop Stops the processing of a worker and closes it's channel in.
-// Returns a blocking channel with type error.
-// !!Should only be called when certain nothing will send to worker.
-func (r *runner) Stop() chan error {
-	r.done.Do(func() {
-		if r.inChan != nil && r.isLeader {
-			close(r.inChan)
-		}
-	})
-	return r.errChan
+	<-r.done
+
+	return r
 }
 
-// IsDone returns a channel signaling the workers context has been canceled.
-func (r *runner) IsDone() <-chan struct{} {
-	return r.ctx.Done()
+// Stop Stops the processing of a worker and waits for workers to finish.
+func (r *Runner) Stop() *Runner {
+	r.cancel()
+	r.Wait()
+	return r
 }
 
-// waitForSignal make sure we wait for a term signal and shutdown correctly
-func (r *runner) waitForSignal(signals ...os.Signal) {
-	go func() {
-		signal.Notify(r.signalChan, signals...)
-		<-r.signalChan
-		if r.cancel != nil {
-			r.cancel()
-		}
+type InputKey struct{}
+
+// work starts processing input and limits worker instance number.
+func (r *Runner) work() {
+	var wg sync.WaitGroup
+
+	defer func() {
+		wg.Wait()
+		r.cancel()
+		close(r.done)
 	}()
-}
 
-// waitForDrain Waits for the limiter to be zeroed out and the in channel to be empty.
-func (r *runner) waitForDrain() {
-	for len(r.limiter) > 0 || len(r.inChan) > 0 {
-		// Wait for the drain.
-	}
-}
+	for {
+		select {
+		case <-r.ctx.Done():
+			return
+		case input, open := <-r.inChan:
+			if !open {
+				return
+			}
+			wg.Add(1)
 
-// startWork Runs the before function and starts processing until one of three things happen.
-// 1. A term signal is received or cancellation of context.
-// 2. Stop function is called.
-// 3. Worker returns an error.
-func (r *runner) startWork() {
-	var err error
-	if err = r.beforeFunc(r.ctx); err != nil {
-		r.errChan <- err
-		return
-	}
-	if r.timeout > 0 {
-		r.ctx, r.cancel = context.WithTimeout(r.ctx, r.timeout)
-	}
-	r.wg.Add(1)
-	go func() {
-		var workerWG = new(sync.WaitGroup)
-		var closeOnce = new(sync.Once)
-
-		// write out error if not nil on exit.
-		defer func() {
-			workerWG.Wait()
-			r.errChan <- err
-			closeOnce.Do(func() {
-				if r.outChan != nil {
-					close(r.outChan)
-				}
-			})
-			r.wg.Done()
-		}()
-		for in := range r.inChan {
-			input := in
 			r.limiter <- struct{}{}
-			workerWG.Add(1)
+
+			inputCtx := context.WithValue(r.ctx, InputKey{}, input)
+			workCtx, workCancel := context.WithCancel(inputCtx)
+			if r.timeout > 0 {
+				workCtx, workCancel = context.WithTimeout(inputCtx, r.timeout)
+			}
+
 			go func() {
 				defer func() {
 					<-r.limiter
-					workerWG.Done()
+					workCancel()
+					wg.Done()
 				}()
-				if err := r.afterFunc(r.ctx, r.workFunc(input, r.outChan)); err != nil {
-					r.once.Do(func() {
-						r.errChan <- err
-						r.cancel()
-					})
+				if err := r.afterFunc(inputCtx, r.workFunc(workCtx, input, r.outChan)); err != nil {
+					r.cancel()
 				}
 			}()
 		}
-	}()
+	}
 }