feat(worker): Add retryTask method and integrate with task loop

Author: KengoWada

tests/worker_test.go

@@ -194,4 +194,76 @@ func TestWorkerProcesses(t *testing.T) {
 
 		assert.Contains(t, logOutput.String(), "Task channel closed, worker exiting...")
 	})
+
+	t.Run("should handle publish failure", func(t *testing.T) {
+		mockBroker := NewMockBroker(1)
+		wg := &sync.WaitGroup{}
+		cfg := taskqueue.WorkerConfig{ID: 1, Broker: mockBroker, Backoff: nil, WG: wg}
+		worker := taskqueue.DefaultWorkerFactory(cfg)
+		worker.Register(failTaskName, failTaskHandler)
+
+		task := taskqueue.Task{
+			Name:     failTaskName,
+			Args:     taskqueue.TaskArgs{},
+			MaxRetry: 3,
+		}
+		err := mockBroker.Publish(task)
+		assert.Nil(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		var logOutput bytes.Buffer
+		log.SetOutput(&logOutput)
+		defer log.SetOutput(nil)
+
+		wg.Add(1)
+		go worker.Start(ctx)
+		mockBroker.badPublish = true
+		time.Sleep(100 * time.Millisecond)
+
+		cancel()
+		wg.Wait()
+
+		expected := fmt.Sprintf("Worker %d: failed to re-queue task %s", cfg.ID, task.Name)
+		assert.Contains(t, logOutput.String(), expected)
+	})
+
+	t.Run("should handle publish failure with backoff", func(t *testing.T) {
+		backoff := &taskqueue.BackoffPolicy{
+			BaseDelay: 1 * time.Second, // small delay for fast test
+			MaxDelay:  20 * time.Second,
+		}
+		mockBroker := NewMockBroker(1)
+		wg := &sync.WaitGroup{}
+		cfg := taskqueue.WorkerConfig{ID: 1, Broker: mockBroker, Backoff: backoff, WG: wg}
+		worker := taskqueue.DefaultWorkerFactory(cfg)
+		worker.Register(failTaskName, failTaskHandler)
+
+		task := taskqueue.Task{
+			Name:     failTaskName,
+			Args:     taskqueue.TaskArgs{},
+			MaxRetry: 3,
+		}
+		err := mockBroker.Publish(task)
+		assert.Nil(t, err)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		var logOutput bytes.Buffer
+		log.SetOutput(&logOutput)
+		defer log.SetOutput(nil)
+
+		wg.Add(1)
+		go worker.Start(ctx)
+		mockBroker.badPublish = true
+		time.Sleep(100 * time.Millisecond)
+
+		cancel()
+		wg.Wait()
+
+		expected := fmt.Sprintf("Worker %d: failed to re-queue task %s", cfg.ID, task.Name)
+		assert.Contains(t, logOutput.String(), expected)
+	})
 }

worker.go

@@ -2,6 +2,7 @@ package taskqueue
 
 import (
 	"context"
+	"fmt"
 	"log"
 	"sync"
 	"time"
@@ -190,16 +191,36 @@ func (w *DefaultWorker) Start(ctx context.Context) {
 				log.Printf("Worker %d: task %s exceeded retries\n", w.id, task.Name)
 				continue
 			}
+			w.retryTask(ctx, task)
+		}
+	}
+}
 
-			task.Retry++
-			if w.backoff != nil {
-				backoffDelay := w.backoff.Calculate(task.Retry)
-				log.Printf("Task %s  failed, backing off for %v (retry %d of %d)", task.Name, backoffDelay, task.Retry, task.MaxRetry)
-				time.Sleep(backoffDelay)
-			}
+func (w *DefaultWorker) retryTask(ctx context.Context, task Task) error {
+	failedToRepublish := fmt.Sprintf("Worker %d: failed to re-queue task %s", w.id, task.Name)
+
+	task.Retry++
+	if w.backoff != nil {
+		backoffDelay := w.backoff.Calculate(task.Retry)
+		log.Printf("Task %s  failed, backing off for %v (retry %d of %d)", task.Name, backoffDelay, task.Retry, task.MaxRetry)
+
+		select {
+		case <-time.After(backoffDelay):
 
-			log.Printf("Task %s failed, retrying now (retry %d of %d)", task.Name, task.Retry, task.MaxRetry)
-			w.broker.Publish(task)
+		case <-ctx.Done():
+			log.Printf("Worker %d: context cancelled during backoff, re-queuing task %s", w.id, task.Name)
+			err := w.broker.Publish(task)
+			if err != nil {
+				log.Println(failedToRepublish)
+			}
+			return err
 		}
 	}
+
+	log.Printf("Task %s failed, retrying now (retry %d of %d)", task.Name, task.Retry, task.MaxRetry)
+	err := w.broker.Publish(task)
+	if err != nil {
+		log.Println(failedToRepublish)
+	}
+	return err
 }