feat(worker): Added middleware support for worker registration (#380)

Author: villfa

worker/README.md

@@ -13,6 +13,7 @@
 * [Documentation](#documentation)
   * [Workers](#workers)
   * [WorkerPool](#workerpool)
+  * [Middlewares](#middlewares)
   * [Logging](#logging)
   * [Tracing](#tracing)
   * [Metrics](#metrics)
@@ -156,6 +157,115 @@ func main() {
 }
 ```
 
+### Middlewares
+
+This module provides middleware support for workers, allowing you to add behaviors without modifying the worker's core implementation.
+
+Middlewares wrap a worker's `Run` method and can perform actions before and after the worker execution, or even modify the execution flow.
+
+#### Using Middlewares
+
+Here's an example of a middleware that returns an error if a worker takes too long to execute:
+
+```go
+package main
+
+import (
+	"context"
+	"errors"
+	"time"
+
+	"github.com/ankorstore/yokai/worker"
+)
+
+// SimpleWorker is a basic worker that sleeps for a specified duration
+type SimpleWorker struct {
+	sleepDuration time.Duration
+}
+
+// NewSimpleWorker creates a new SimpleWorker
+func NewSimpleWorker(sleepDuration time.Duration) *SimpleWorker {
+	return &SimpleWorker{
+		sleepDuration: sleepDuration,
+	}
+}
+
+// Name returns the worker name
+func (w *SimpleWorker) Name() string {
+	return "simple-worker"
+}
+
+// Run executes the worker
+func (w *SimpleWorker) Run(ctx context.Context) error {
+	// Simulate work by sleeping
+	time.Sleep(w.sleepDuration)
+	
+	return nil
+}
+
+// TimeoutMiddleware implements the worker.Middleware interface
+type TimeoutMiddleware struct {
+	timeout time.Duration
+}
+
+// NewTimeoutMiddleware creates a new TimeoutMiddleware with the specified timeout
+func NewTimeoutMiddleware(timeout time.Duration) *TimeoutMiddleware {
+	return &TimeoutMiddleware{
+		timeout: timeout,
+	}
+}
+
+// Name returns the middleware name
+func (m *TimeoutMiddleware) Name() string {
+	return "timeout-middleware"
+}
+
+// Handle returns the middleware function
+func (m *TimeoutMiddleware) Handle() worker.MiddlewareFunc {
+	return func(next worker.HandlerFunc) worker.HandlerFunc {
+		return func(ctx context.Context) error {
+			// Create a timeout context
+			timeoutCtx, cancel := context.WithTimeout(ctx, m.timeout)
+			defer cancel()
+
+			// Create a channel to receive the result of the worker execution
+			done := make(chan error)
+
+			// Execute the worker in a goroutine
+			go func() {
+				done <- next(timeoutCtx)
+			}()
+
+			// Wait for either the worker to complete or the timeout to occur
+			select {
+			case err := <-done:
+				return err
+			case <-timeoutCtx.Done():
+				if errors.Is(timeoutCtx.Err(), context.DeadlineExceeded) {
+					return errors.New("worker execution timed out")
+				}
+				return timeoutCtx.Err()
+			}
+		}
+	}
+}
+
+func main() {
+	// Create a worker pool with a worker and the timeout middleware
+	pool, _ := worker.NewDefaultWorkerPoolFactory().Create(
+		worker.WithWorker(
+			NewSimpleWorker(10 * time.Second), // Worker that takes 10 seconds to complete
+			worker.WithMiddlewares(NewTimeoutMiddleware(5 * time.Second)), // Middleware that times out after 5 seconds
+		),
+	)
+
+	// Start the pool
+	pool.Start(context.Background())
+	
+	// The worker will be interrupted after 5 seconds with a timeout error
+}
+```
+
 ### Logging
 
 You can use the [CtxLogger()](context.go) function to retrieve the

worker/execution.go

@@ -15,6 +15,7 @@ type WorkerExecution struct {
 	maxExecutionsAttempts   int
 	deferredStartThreshold  float64
 	events                  []*WorkerExecutionEvent
+	middlewares             []Middleware
 }
 
 // NewWorkerExecution returns a new [WorkerExecution].
@@ -27,6 +28,7 @@ func NewWorkerExecution(id string, name string, options ExecutionOptions) *Worke
 		maxExecutionsAttempts:   options.MaxExecutionsAttempts,
 		deferredStartThreshold:  options.DeferredStartThreshold,
 		events:                  []*WorkerExecutionEvent{},
+		middlewares:             options.Middlewares,
 	}
 }
 
@@ -157,3 +159,10 @@ func (e *WorkerExecution) HasEvent(message string) bool {
 
 	return false
 }
+
+func (e *WorkerExecution) Middlewares() []Middleware {
+	e.mutex.Lock()
+	defer e.mutex.Unlock()
+
+	return e.middlewares
+}

worker/middleware.go

@@ -0,0 +1,15 @@
+package worker
+
+import "context"
+
+// Middleware is the interface to implement to provide worker middlewares.
+type Middleware interface {
+	Name() string
+	Handle() MiddlewareFunc
+}
+
+// MiddlewareFunc wraps handlers in the middleware chain to perform operations before and after execution.
+type MiddlewareFunc func(next HandlerFunc) HandlerFunc
+
+// HandlerFunc executes a worker's Run method or another middleware in the chain with context information.
+type HandlerFunc func(ctx context.Context) error

worker/middleware_test.go

@@ -0,0 +1,15 @@
+package worker_test
+
+import "github.com/ankorstore/yokai/worker"
+
+type TestMiddleware struct {
+	Func worker.MiddlewareFunc
+}
+
+func (m *TestMiddleware) Name() string {
+	return "TestMiddleware"
+}
+
+func (m *TestMiddleware) Handle() worker.MiddlewareFunc {
+	return m.Func
+}

worker/option.go

@@ -71,13 +71,15 @@ func WithWorker(worker Worker, options ...WorkerExecutionOption) WorkerPoolOptio
 type ExecutionOptions struct {
 	DeferredStartThreshold float64
 	MaxExecutionsAttempts  int
+	Middlewares            []Middleware
 }
 
 // DefaultWorkerExecutionOptions are the default options for the [Worker] executions.
 func DefaultWorkerExecutionOptions() ExecutionOptions {
 	return ExecutionOptions{
 		DeferredStartThreshold: DefaultDeferredStartThreshold,
 		MaxExecutionsAttempts:  DefaultMaxExecutionsAttempts,
+		Middlewares:            []Middleware{},
 	}
 }
 
@@ -97,3 +99,13 @@ func WithMaxExecutionsAttempts(l int) WorkerExecutionOption {
 		o.MaxExecutionsAttempts = l
 	}
 }
+
+// WithMiddlewares is used to add middlewares to a worker registration.
+func WithMiddlewares(middlewares ...Middleware) WorkerExecutionOption {
+	return func(o *ExecutionOptions) {
+		if o.Middlewares == nil {
+			o.Middlewares = []Middleware{}
+		}
+		o.Middlewares = append(o.Middlewares, middlewares...)
+	}
+}

worker/option_test.go

@@ -79,3 +79,15 @@ func TestWorkerExecutionOptionsWithMaxExecutionsAttempts(t *testing.T) {
 
 	assert.Equal(t, 2, opt.MaxExecutionsAttempts)
 }
+
+func TestWorkerExecutionOptionsWithMiddlewares(t *testing.T) {
+	t.Parallel()
+
+	opt := worker.DefaultWorkerExecutionOptions()
+
+	assert.Len(t, opt.Middlewares, 0)
+
+	worker.WithMiddlewares(&TestMiddleware{}, &TestMiddleware{})(&opt)
+
+	assert.Len(t, opt.Middlewares, 2)
+}

worker/pool.go

@@ -214,7 +214,14 @@ func (p *WorkerPool) startWorkerRegistration(ctx context.Context, registration *
 
 		p.metrics.IncrementWorkerExecutionStart(registration.Worker().Name())
 
-		if err := registration.Worker().Run(ctx); err != nil {
+		runFunc := registration.Worker().Run
+
+		for i := len(execution.Middlewares()) - 1; i >= 0; i-- {
+			middleware := execution.Middlewares()[i]
+			runFunc = middleware.Handle()(runFunc)
+		}
+
+		if err := runFunc(ctx); err != nil {
 			message = fmt.Sprintf(
 				"stopping execution attempt %d/%d with error: %v",
 				workerExecution.CurrentExecutionAttempt(),

worker/pool_test.go

@@ -544,3 +544,67 @@ func TestExecutionWithRestartingPanicWorker(t *testing.T) {
 	)
 	assert.NoError(t, err)
 }
+
+func TestExecutionWithMiddlewares(t *testing.T) {
+	t.Parallel()
+
+	// Create a slice to track the execution order
+	executionOrder := []string{}
+
+	// Create a worker
+	testWorker := workers.NewClassicWorker()
+
+	// Create the first middleware
+	middleware1 := &TestMiddleware{
+		Func: func(next worker.HandlerFunc) worker.HandlerFunc {
+			return func(ctx context.Context) error {
+				executionOrder = append(executionOrder, "middleware1 before")
+				err := next(ctx)
+				executionOrder = append(executionOrder, "middleware1 after")
+
+				return err
+			}
+		},
+	}
+
+	// Create the second middleware
+	middleware2 := &TestMiddleware{
+		Func: func(next worker.HandlerFunc) worker.HandlerFunc {
+			return func(ctx context.Context) error {
+				executionOrder = append(executionOrder, "middleware2 before")
+				err := next(ctx)
+				executionOrder = append(executionOrder, "middleware2 after")
+
+				return err
+			}
+		},
+	}
+
+	// Create a worker pool with the worker and middlewares
+	pool, err := worker.NewDefaultWorkerPoolFactory().Create(
+		worker.WithWorker(
+			testWorker,
+			worker.WithMiddlewares(middleware1, middleware2),
+		),
+	)
+	assert.NoError(t, err)
+
+	// Start the worker pool
+	err = pool.Start(context.Background())
+	assert.NoError(t, err)
+
+	// Wait for the worker to complete
+	time.Sleep(10 * time.Millisecond)
+
+	// Stop the worker pool
+	err = pool.Stop()
+	assert.NoError(t, err)
+
+	// Verify the execution order
+	assert.Equal(t, []string{
+		"middleware1 before",
+		"middleware2 before",
+		"middleware2 after",
+		"middleware1 after",
+	}, executionOrder)
+}