refactor worker

Author: elgopher

worker.go

@@ -21,67 +21,85 @@ type worker[Resource any] struct {
 	saveResource       func(context.Context, string, Resource) error
 	minDuration        time.Duration
 	maxDuration        time.Duration
+
+	batchByResourceKey map[string]*batch[Resource]
+	batchByDeadline    []*batch[Resource]
 }
 
-func (w worker[Resource]) run() {
-	batchByResourceKey := map[string]*batch[Resource]{}
-	var batchByDeadline []*batch[Resource]
+func (w *worker[Resource]) run() {
+	w.batchByResourceKey = map[string]*batch[Resource]{}
 
 	ticker := time.NewTicker(time.Millisecond)
 	defer ticker.Stop()
 
 	for {
 		select {
 		case <-ticker.C:
-			now := time.Now()
-			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
-				}
-			}
+			w.endBatchesAfterDeadline()
 		case _operation, ok := <-w.incomingOperations:
 			if !ok {
-				for key, _batch := range batchByResourceKey {
-					err := w.saveResource(_batch.ctx, key, _batch.resource)
-					_batch.publishResult(err)
-					continue
-				}
+				w.endAllBatches()
 				return
 			}
 
-			_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, _operation.resourceKey)
-				if err != nil {
-					_operation.result <- err
-					continue
-				}
-				_batch = &batch[Resource]{
-					ctx:      ctx,
-					key:      _operation.resourceKey,
-					resource: resource,
-					deadline: now.Add(w.minDuration),
-				}
-				batchByResourceKey[_operation.resourceKey] = _batch
-				batchByDeadline = append(batchByDeadline, _batch)
-			}
+			w.runOperation(_operation)
+		}
+	}
+}
 
-			_batch.results = append(_batch.results, _operation.result)
+func (w *worker[Resource]) endBatchesAfterDeadline() {
+	now := time.Now()
 
-			_operation.run(_batch.resource)
+	for _, _batch := range w.batchByDeadline {
+		if _batch.deadline.After(now) {
+			return
 		}
+
+		err := w.saveResource(_batch.ctx, _batch.key, _batch.resource)
+		_batch.publishResult(err)
+		delete(w.batchByResourceKey, _batch.key)
+		w.batchByDeadline = w.batchByDeadline[1:]
 	}
 }
 
+func (w *worker[Resource]) endAllBatches() {
+	for key, _batch := range w.batchByResourceKey {
+		err := w.saveResource(_batch.ctx, key, _batch.resource)
+		_batch.publishResult(err)
+	}
+
+	w.batchByResourceKey = map[string]*batch[Resource]{}
+	w.batchByDeadline = nil
+}
+
+func (w *worker[Resource]) runOperation(_operation operation[Resource]) {
+	_batch, found := w.batchByResourceKey[_operation.resourceKey]
+	if !found {
+		ctx, _ := context.WithTimeout(context.Background(), w.maxDuration)
+
+		now := time.Now()
+
+		resource, err := w.loadResource(ctx, _operation.resourceKey)
+		if err != nil {
+			_operation.result <- err
+			return
+		}
+
+		_batch = &batch[Resource]{
+			ctx:      ctx,
+			key:      _operation.resourceKey,
+			resource: resource,
+			deadline: now.Add(w.minDuration),
+		}
+		w.batchByResourceKey[_operation.resourceKey] = _batch
+		w.batchByDeadline = append(w.batchByDeadline, _batch)
+	}
+
+	_batch.results = append(_batch.results, _operation.result)
+
+	_operation.run(_batch.resource)
+}
+
 type batch[Resource any] struct {
 	ctx      context.Context
 	key      string