Processor.Stop will wait until all batches are finished

Processor.Stop() should block until all batches are ended (and resources stored). This is needed to implement graceful shutdown in applications.

Author: elgopher

_example/main.go

@@ -23,6 +23,7 @@ func main() {
 			SaveResource: db.SaveTrain,
 		},
 	)
+	defer processor.Stop()
 
 	trainService := train.Service{
 		BatchProcessor: processor,

batch.go

@@ -6,6 +6,7 @@ package batch
 import (
 	"context"
 	"runtime"
+	"sync"
 	"time"
 )
 
@@ -21,31 +22,41 @@ type Options[Resource any] struct {
 func StartProcessor[Resource any](options Options[Resource]) *Processor[Resource] {
 	options = options.withDefaults()
 
-	operations := make([]chan operation[Resource], options.GoRoutines)
+	workerChannels := make([]chan operation[Resource], options.GoRoutines)
+
+	var workersFinished sync.WaitGroup
+	workersFinished.Add(options.GoRoutines)
+
 	for i := 0; i < options.GoRoutines; i++ {
-		operations[i] = make(chan operation[Resource])
-		w := worker[Resource]{
-			goRoutineNumber: i,
-			operations:      operations[i],
-			loadResource:    options.LoadResource,
-			saveResource:    options.SaveResource,
-			minDuration:     options.MinDuration,
-			maxDuration:     options.MaxDuration,
+		workerChannels[i] = make(chan operation[Resource])
+		_worker := worker[Resource]{
+			goRoutineNumber:    i,
+			incomingOperations: workerChannels[i],
+			loadResource:       options.LoadResource,
+			saveResource:       options.SaveResource,
+			minDuration:        options.MinDuration,
+			maxDuration:        options.MaxDuration,
 		}
-		go w.run()
+
+		go func() {
+			_worker.run()
+			workersFinished.Done()
+		}()
 	}
 
 	return &Processor[Resource]{
-		options:    options,
-		stopped:    make(chan struct{}),
-		operations: operations,
+		options:         options,
+		stopped:         make(chan struct{}),
+		workerChannels:  workerChannels,
+		workersFinished: &workersFinished,
 	}
 }
 
 type Processor[Resource any] struct {
-	options    Options[Resource]
-	stopped    chan struct{}
-	operations []chan operation[Resource]
+	options         Options[Resource]
+	stopped         chan struct{}
+	workerChannels  []chan operation[Resource]
+	workersFinished *sync.WaitGroup
 }
 
 func (s Options[Resource]) withDefaults() Options[Resource] {
@@ -97,7 +108,7 @@ func (p *Processor[Resource]) Run(key string, op func(Resource)) error {
 
 	goRoutineNumber := p.options.GoRoutineNumberForKey(key, p.options.GoRoutines)
 
-	p.operations[goRoutineNumber] <- operation[Resource]{
+	p.workerChannels[goRoutineNumber] <- operation[Resource]{
 		resourceKey: key,
 		run:         op,
 		result:      result,
@@ -107,9 +118,13 @@ func (p *Processor[Resource]) Run(key string, op func(Resource)) error {
 }
 
 // Stop ends all running batches. No new operations will be accepted.
+// Stop blocks until all pending batches are ended and resources saved.
 func (p *Processor[Resource]) Stop() {
 	close(p.stopped)
-	for _, op := range p.operations {
-		close(op)
+
+	for _, channel := range p.workerChannels {
+		close(channel)
 	}
+
+	p.workersFinished.Wait()
 }

batch_test.go

@@ -23,6 +23,7 @@ func TestProcessor_Run(t *testing.T) {
 		futureValue := FutureValue[*resource]()
 
 		processor := batch.StartProcessor(batch.Options[*resource]{})
+		defer processor.Stop()
 		// when
 		err := processor.Run("key", func(c *resource) {
 			futureValue.Set(c)
@@ -44,6 +45,7 @@ func TestProcessor_Run(t *testing.T) {
 					return res, nil
 				},
 			})
+		defer processor.Stop()
 		// when
 		err := processor.Run(key, func(r *resource) {
 			futureValue.Set(r)
@@ -64,6 +66,7 @@ func TestProcessor_Run(t *testing.T) {
 				LoadResource: db.Load,
 				SaveResource: db.Save,
 			})
+		defer processor.Stop()
 		// when
 		err := processor.Run(key, func(r *resource) {
 			r.value = 2
@@ -83,6 +86,7 @@ func TestProcessor_Run(t *testing.T) {
 				MinDuration: 100 * time.Millisecond,
 			},
 		)
+		defer processor.Stop()
 
 		started := time.Now()
 		// when
@@ -95,6 +99,7 @@ func TestProcessor_Run(t *testing.T) {
 
 	t.Run("should run batch with default min duration", func(t *testing.T) {
 		processor := batch.StartProcessor(batch.Options[empty]{})
+		defer processor.Stop()
 
 		started := time.Now()
 		// when
@@ -117,6 +122,7 @@ func TestProcessor_Run(t *testing.T) {
 					return nil
 				},
 			})
+		defer processor.Stop()
 
 		key := ""
 
@@ -146,6 +152,7 @@ func TestProcessor_Run(t *testing.T) {
 					SaveResource: db.Save,
 				},
 			)
+			defer processor.Stop()
 			// when
 			err := processor.Run(key, func(*resource) {})
 			// then
@@ -172,6 +179,7 @@ func TestProcessor_Run(t *testing.T) {
 					SaveResource: db.Save,
 				},
 			)
+			defer processor.Stop()
 			// when
 			err := processor.Run(key, func(*resource) {})
 			// then
@@ -196,6 +204,7 @@ func TestProcessor_Run(t *testing.T) {
 					},
 				},
 			)
+			defer processor.Stop()
 			// when
 			err := processor.Run(key, func(empty) {})
 			// then
@@ -211,6 +220,7 @@ func TestProcessor_Run(t *testing.T) {
 				},
 			},
 		)
+		defer processor.Stop()
 
 		const iterations = 1000
 
@@ -273,6 +283,34 @@ func TestProcessor_Stop(t *testing.T) {
 		elapsed := time.Now().Sub(started)
 		assert.True(t, elapsed < minDuration, "stopped batch should take less time than batch min duration")
 	})
+
+	t.Run("Stop should wait until all batches are finished", func(t *testing.T) {
+		var operationExecuted sync.WaitGroup
+		operationExecuted.Add(1)
+
+		batchFinished := false
+		processor := batch.StartProcessor(
+			batch.Options[empty]{
+				MinDuration: time.Second,
+				MaxDuration: time.Second,
+				SaveResource: func(ctx context.Context, key string, _ empty) error {
+					<-ctx.Done()
+					batchFinished = true
+					return nil
+				},
+			},
+		)
+		go func() {
+			_ = processor.Run("key", func(empty) {
+				operationExecuted.Done()
+			})
+		}()
+		operationExecuted.Wait()
+		// when
+		processor.Stop()
+		// then
+		assert.True(t, batchFinished)
+	})
 }
 
 type resource struct{ value int }

worker.go

@@ -15,12 +15,12 @@ type operation[Resource any] struct {
 }
 
 type worker[Resource any] struct {
-	goRoutineNumber int
-	operations      chan operation[Resource]
-	loadResource    func(context.Context, string) (Resource, error)
-	saveResource    func(context.Context, string, Resource) error
-	minDuration     time.Duration
-	maxDuration     time.Duration
+	goRoutineNumber    int
+	incomingOperations <-chan operation[Resource]
+	loadResource       func(context.Context, string) (Resource, error)
+	saveResource       func(context.Context, string, Resource) error
+	minDuration        time.Duration
+	maxDuration        time.Duration
 }
 
 func (w worker[Resource]) run() {
@@ -34,50 +34,50 @@ func (w worker[Resource]) run() {
 		select {
 		case <-ticker.C:
 			now := time.Now()
-			for _, b := range batchByDeadline {
-				if b.deadline.Before(now) {
-					err := w.saveResource(b.ctx, b.key, b.resource)
-					b.publishResult(err)
-					delete(batchByResourceKey, b.key)
+			for _, _batch := range batchByDeadline {
+				if _batch.deadline.Before(now) {
+					err := w.saveResource(_batch.ctx, _batch.key, _batch.resource)
+					_batch.publishResult(err)
+					delete(batchByResourceKey, _batch.key)
 					batchByDeadline = batchByDeadline[1:]
 					continue
 				}
 			}
-		case op, ok := <-w.operations:
+		case _operation, ok := <-w.incomingOperations:
 			if !ok {
-				for key, b := range batchByResourceKey {
-					err := w.saveResource(b.ctx, key, b.resource)
-					b.publishResult(err)
+				for key, _batch := range batchByResourceKey {
+					err := w.saveResource(_batch.ctx, key, _batch.resource)
+					_batch.publishResult(err)
 					continue
 				}
 				return
 			}
 
-			b, found := batchByResourceKey[op.resourceKey]
+			_batch, found := batchByResourceKey[_operation.resourceKey]
 			if !found {
 				ctx, cancel := context.WithTimeout(context.Background(), w.maxDuration)
 				defer cancel()
 
 				now := time.Now()
 
-				resource, err := w.loadResource(ctx, op.resourceKey)
+				resource, err := w.loadResource(ctx, _operation.resourceKey)
 				if err != nil {
-					op.result <- err
+					_operation.result <- err
 					continue
 				}
-				b = &batch[Resource]{
+				_batch = &batch[Resource]{
 					ctx:      ctx,
-					key:      op.resourceKey,
+					key:      _operation.resourceKey,
 					resource: resource,
 					deadline: now.Add(w.minDuration),
 				}
-				batchByResourceKey[op.resourceKey] = b
-				batchByDeadline = append(batchByDeadline, b)
+				batchByResourceKey[_operation.resourceKey] = _batch
+				batchByDeadline = append(batchByDeadline, _batch)
 			}
 
-			b.results = append(b.results, op.result)
+			_batch.results = append(_batch.results, _operation.result)
 
-			op.run(b.resource)
+			_operation.run(_batch.resource)
 		}
 	}
 }