Merge pull request #274 from cschleiden/unify-worker

Unify worker code

Author: cschleiden

backend/redis/expire_test.go

@@ -25,7 +25,7 @@ func Test_AutoExpiration(t *testing.T) {
 	b := setup()
 
 	c := client.New(b)
-	w := worker.New(b, &worker.DefaultWorkerOptions)
+	w := worker.New(b, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
 
@@ -66,7 +66,7 @@ func Test_AutoExpiration_SubWorkflow(t *testing.T) {
 	b := setup()
 
 	c := client.New(b)
-	w := worker.New(b, &worker.DefaultWorkerOptions)
+	w := worker.New(b, nil)
 
 	ctx, cancel := context.WithCancel(context.Background())
 

backend/test/e2e.go

@@ -704,7 +704,7 @@ func EndToEndBackendTest(t *testing.T, setup func(options ...backend.BackendOpti
 		}
 	}
 
-	options := worker.DefaultWorkerOptions
+	options := worker.DefaultOptions
 
 	// Run with cache
 	run("", options)

bench/main.go

@@ -46,7 +46,7 @@ func main() {
 	)
 	defer ba.Close()
 
-	wo := worker.DefaultWorkerOptions
+	wo := worker.DefaultOptions
 	wo.WorkflowExecutorCacheSize = *cacheSize
 	w := worker.New(ba, &wo)
 

internal/worker/activity.go

@@ -2,9 +2,7 @@ package worker
 
 import (
 	"context"
-	"errors"
 	"log/slog"
-	"sync"
 	"time"
 
 	"github.com/benbjohnson/clock"
@@ -14,197 +12,82 @@ import (
 	"github.com/cschleiden/go-workflows/backend/payload"
 	"github.com/cschleiden/go-workflows/internal/activity"
 	"github.com/cschleiden/go-workflows/internal/metrickeys"
-	mi "github.com/cschleiden/go-workflows/internal/metrics"
+	im "github.com/cschleiden/go-workflows/internal/metrics"
 	"github.com/cschleiden/go-workflows/internal/workflow"
 	"github.com/cschleiden/go-workflows/internal/workflowerrors"
 )
 
-type ActivityWorker struct {
-	backend backend.Backend
+func NewActivityWorker(b backend.Backend, registry *workflow.Registry, clock clock.Clock, options WorkerOptions) *Worker[backend.ActivityTask, history.Event] {
+	ae := activity.NewExecutor(b.Logger(), b.Tracer(), b.Converter(), b.ContextPropagators(), registry)
 
-	options *Options
-
-	activityTaskQueue    chan *backend.ActivityTask
-	activityTaskExecutor *activity.Executor
-
-	wg        sync.WaitGroup
-	pollersWg sync.WaitGroup
-
-	clock  clock.Clock
-	logger *slog.Logger
-}
-
-func NewActivityWorker(b backend.Backend, registry *workflow.Registry, clock clock.Clock, options *Options) *ActivityWorker {
-	return &ActivityWorker{
-		backend: b,
-
-		options: options,
-
-		activityTaskQueue:    make(chan *backend.ActivityTask),
-		activityTaskExecutor: activity.NewExecutor(b.Logger(), b.Tracer(), b.Converter(), b.ContextPropagators(), registry),
-
-		clock:  clock,
-		logger: b.Logger(),
-	}
-}
-
-func (aw *ActivityWorker) Start(ctx context.Context) error {
-	aw.pollersWg.Add(aw.options.ActivityPollers)
-
-	for i := 0; i < aw.options.ActivityPollers; i++ {
-		go aw.runPoll(ctx)
+	tw := &ActivityTaskWorker{
+		backend:              b,
+		activityTaskExecutor: ae,
+		clock:                clock,
+		logger:               b.Logger(),
 	}
 
-	go aw.runDispatcher()
-
-	return nil
+	return NewWorker[backend.ActivityTask, history.Event](b, tw, &options)
 }
 
-func (aw *ActivityWorker) WaitForCompletion() error {
-	// Wait for task pollers to finish
-	aw.pollersWg.Wait()
-
-	// Wait for tasks to finish
-	aw.wg.Wait()
-	close(aw.activityTaskQueue)
-
-	return nil
-}
-
-func (aw *ActivityWorker) runPoll(ctx context.Context) {
-	defer aw.pollersWg.Done()
-
-	ticker := time.NewTicker(aw.options.ActivityPollingInterval)
-	defer ticker.Stop()
-	for {
-		task, err := aw.poll(ctx, 30*time.Second)
-		if err != nil {
-			aw.logger.ErrorContext(ctx, "error while polling for activity task", "error", err)
-		}
-		if task != nil {
-			aw.wg.Add(1)
-			aw.activityTaskQueue <- task
-			continue // check for new tasks right away
-		}
-
-		select {
-		case <-ctx.Done():
-			return
-		case <-ticker.C:
-		}
-	}
+type ActivityTaskWorker struct {
+	backend              backend.Backend
+	activityTaskExecutor *activity.Executor
+	clock                clock.Clock
+	logger               *slog.Logger
 }
 
-func (aw *ActivityWorker) runDispatcher() {
-	var sem chan struct{}
-	if aw.options.MaxParallelActivityTasks > 0 {
-		sem = make(chan struct{}, aw.options.MaxParallelActivityTasks)
+func (atw *ActivityTaskWorker) Complete(ctx context.Context, event *history.Event, task *backend.ActivityTask) error {
+	if err := atw.backend.CompleteActivityTask(ctx, task.WorkflowInstance, task.ID, event); err != nil {
+		atw.backend.Logger().Error("completing activity task", "error", err)
 	}
 
-	for task := range aw.activityTaskQueue {
-		if sem != nil {
-			sem <- struct{}{}
-		}
-
-		task := task
-
-		go func() {
-			defer aw.wg.Done()
-
-			// Create new context to allow activities to complete when root context is canceled
-			taskCtx := context.Background()
-			aw.handleTask(taskCtx, task)
-
-			if sem != nil {
-				<-sem
-			}
-		}()
-	}
+	return nil
 }
 
-func (aw *ActivityWorker) handleTask(ctx context.Context, task *backend.ActivityTask) {
+func (atw *ActivityTaskWorker) Execute(ctx context.Context, task *backend.ActivityTask) (*history.Event, error) {
 	a := task.Event.Attributes.(*history.ActivityScheduledAttributes)
-	ametrics := aw.backend.Metrics().WithTags(metrics.Tags{metrickeys.ActivityName: a.Name})
+	ametrics := atw.backend.Metrics().WithTags(metrics.Tags{metrickeys.ActivityName: a.Name})
 
 	// Record how long this task was in the queue
 	scheduledAt := task.Event.Timestamp
 	timeInQueue := time.Since(scheduledAt)
 	ametrics.Distribution(metrickeys.ActivityTaskDelay, metrics.Tags{}, float64(timeInQueue/time.Millisecond))
 
-	// Start heartbeat while activity is running
-	if aw.options.ActivityHeartbeatInterval > 0 {
-		heartbeatCtx, cancelHeartbeat := context.WithCancel(ctx)
-		defer cancelHeartbeat()
-
-		go func(ctx context.Context) {
-			t := time.NewTicker(aw.options.ActivityHeartbeatInterval)
-			defer t.Stop()
-
-			for {
-				select {
-				case <-ctx.Done():
-					return
-				case <-t.C:
-					if err := aw.backend.ExtendActivityTask(ctx, task.ID); err != nil {
-						if !errors.Is(err, context.Canceled) {
-							aw.backend.Logger().Error("extending activity task", "error", err)
-							panic("extending activity task")
-						}
-					}
-				}
-			}
-		}(heartbeatCtx)
-	}
-
-	timer := mi.NewTimer(ametrics, metrickeys.ActivityTaskProcessed, metrics.Tags{})
+	timer := im.NewTimer(ametrics, metrickeys.ActivityTaskProcessed, metrics.Tags{})
 	defer timer.Stop()
 
-	result, err := aw.activityTaskExecutor.ExecuteActivity(ctx, task)
-	event := aw.resultToEvent(task.Event.ScheduleEventID, result, err)
+	result, err := atw.activityTaskExecutor.ExecuteActivity(ctx, task)
+	event := atw.resultToEvent(task.Event.ScheduleEventID, result, err)
 
-	if err := aw.backend.CompleteActivityTask(ctx, task.WorkflowInstance, task.ID, event); err != nil {
-		aw.backend.Logger().Error("completing activity task", "error", err)
-		panic("completing activity task")
-	}
+	return event, nil
+}
+
+func (atw *ActivityTaskWorker) Extend(ctx context.Context, task *backend.ActivityTask) error {
+	return atw.backend.ExtendActivityTask(ctx, task.ID)
 }
 
-func (aw *ActivityWorker) resultToEvent(ScheduleEventID int64, result payload.Payload, err error) *history.Event {
+func (atw *ActivityTaskWorker) Get(ctx context.Context) (*backend.ActivityTask, error) {
+	return atw.backend.GetActivityTask(ctx)
+}
+
+func (atw *ActivityTaskWorker) resultToEvent(scheduleEventID int64, result payload.Payload, err error) *history.Event {
 	if err != nil {
 		return history.NewPendingEvent(
-			aw.clock.Now(),
+			atw.clock.Now(),
 			history.EventType_ActivityFailed,
 			&history.ActivityFailedAttributes{
 				Error: workflowerrors.FromError(err),
 			},
-			history.ScheduleEventID(ScheduleEventID),
+			history.ScheduleEventID(scheduleEventID),
 		)
 	}
 
 	return history.NewPendingEvent(
-		aw.clock.Now(),
+		atw.clock.Now(),
 		history.EventType_ActivityCompleted,
 		&history.ActivityCompletedAttributes{
 			Result: result,
 		},
-		history.ScheduleEventID(ScheduleEventID))
-}
-
-func (aw *ActivityWorker) poll(ctx context.Context, timeout time.Duration) (*backend.ActivityTask, error) {
-	if timeout == 0 {
-		timeout = 30 * time.Second
-	}
-
-	ctx, cancel := context.WithTimeout(ctx, timeout)
-	defer cancel()
-
-	task, err := aw.backend.GetActivityTask(ctx)
-	if err != nil {
-		if errors.Is(err, context.DeadlineExceeded) {
-			return nil, nil
-		}
-
-		return nil, err
-	}
-
-	return task, nil
+		history.ScheduleEventID(scheduleEventID))
 }