From 13e948d4abbeeee9466cbd1a801791d100963ab6 Mon Sep 17 00:00:00 2001
From: Rogee <rogeecn@users.noreply.github.com>
Date: Fri, 14 Aug 2026 22:43:01 +0800
Subject: [PATCH] H-141: harden worker shutdown persistence (#24)

Co-authored-by: Rogee <rogee@ipao.vip>
---
 .../dispatch/dispatcher_worker_test.go        |  4 +--
 backend/internal/worker/worker.go             |  5 +++
 backend/internal/worker/worker_test.go        | 34 +++++++++++++++++--
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/backend/internal/dispatch/dispatcher_worker_test.go b/backend/internal/dispatch/dispatcher_worker_test.go
index 8e426176..a2e16bbf 100644
--- a/backend/internal/dispatch/dispatcher_worker_test.go
+++ b/backend/internal/dispatch/dispatcher_worker_test.go
@@ -2,6 +2,7 @@ package dispatch
 
 import (
 	"context"
+	"path/filepath"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -28,7 +29,7 @@ func (l *dispatchWorkerListener) OnEvent(ctx context.Context, event *channel.Cha
 
 func newDispatchWorkerDB(t *testing.T) *gorm.DB {
 	t.Helper()
-	db, err := gorm.Open(sqlite.Open("file:dispatch-worker?mode=memory&cache=shared"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "dispatch-worker.db")), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
 	if err != nil {
 		t.Fatalf("open sqlite: %v", err)
 	}
@@ -41,7 +42,6 @@ func newDispatchWorkerDB(t *testing.T) *gorm.DB {
 		t.Fatalf("migrate background jobs: %v", err)
 	}
 	t.Cleanup(func() {
-		db.Exec("DELETE FROM background_jobs")
 		sqlDB.Close()
 	})
 	return db
diff --git a/backend/internal/worker/worker.go b/backend/internal/worker/worker.go
index f2666fec..ab3219cd 100644
--- a/backend/internal/worker/worker.go
+++ b/backend/internal/worker/worker.go
@@ -387,6 +387,9 @@ func (wp *WorkerPool) ProcessOne(ctx context.Context) (bool, error) {
 	if wp.db == nil {
 		return false, ErrWorkerDatabaseRequired
 	}
+	if err := ctx.Err(); err != nil {
+		return false, err
+	}
 	job, err := wp.claimNext(ctx)
 	if err != nil {
 		if errors.Is(err, gorm.ErrRecordNotFound) {
@@ -695,6 +698,7 @@ func (wp *WorkerPool) perform(ctx context.Context, job *model.BackgroundJob) err
 	if err := handler(ctx, job); err != nil {
 		return wp.fail(ctx, job, err)
 	}
+	ctx = context.WithoutCancel(ctx)
 	finishedAt := wp.now()
 	updates := map[string]any{
 		"status":      model.BackgroundJobStatusCompleted,
@@ -707,6 +711,7 @@ func (wp *WorkerPool) perform(ctx context.Context, job *model.BackgroundJob) err
 }
 
 func (wp *WorkerPool) fail(ctx context.Context, job *model.BackgroundJob, err error) error {
+	ctx = context.WithoutCancel(ctx)
 	now := wp.now()
 	updates := map[string]any{
 		"locked_at":  nil,
diff --git a/backend/internal/worker/worker_test.go b/backend/internal/worker/worker_test.go
index 5765e1af..aae550df 100644
--- a/backend/internal/worker/worker_test.go
+++ b/backend/internal/worker/worker_test.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"path/filepath"
 	"sync/atomic"
 	"testing"
 	"time"
@@ -18,7 +19,7 @@ import (
 
 func newWorkerTestDB(t *testing.T) *gorm.DB {
 	t.Helper()
-	db, err := gorm.Open(sqlite.Open("file:worker-test?mode=memory&cache=shared"), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
+	db, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "worker.db")), &gorm.Config{Logger: logger.Default.LogMode(logger.Silent)})
 	if err != nil {
 		t.Fatalf("open sqlite: %v", err)
 	}
@@ -31,7 +32,6 @@ func newWorkerTestDB(t *testing.T) *gorm.DB {
 		t.Fatalf("migrate background jobs: %v", err)
 	}
 	t.Cleanup(func() {
-		db.Exec("DELETE FROM background_jobs")
 		sqlDB.Close()
 	})
 	return db
@@ -242,6 +242,36 @@ func TestWorkerPoolStartAndStopProcessJobs(t *testing.T) {
 	}
 }
 
+func TestWorkerPoolStopPersistsCancelledJobRetry(t *testing.T) {
+	db := newWorkerTestDB(t)
+	wp := NewWorkerPoolWithOptions(db, WithPollInterval(5*time.Millisecond), WithBackoff(func(int) time.Duration { return 0 }))
+	started := make(chan struct{})
+	wp.Register("cancelled", func(ctx context.Context, job *model.BackgroundJob) error {
+		close(started)
+		<-ctx.Done()
+		return ctx.Err()
+	})
+	job, err := wp.Enqueue(context.Background(), "cancelled", nil)
+	if err != nil {
+		t.Fatalf("enqueue: %v", err)
+	}
+	if err := wp.Start(); err != nil {
+		t.Fatalf("start worker: %v", err)
+	}
+	select {
+	case <-started:
+	case <-time.After(time.Second):
+		t.Fatal("worker did not start job")
+	}
+	if err := wp.Stop(); err != nil {
+		t.Fatalf("stop worker: %v", err)
+	}
+	reloaded := loadJob(t, db, job.ID)
+	if reloaded.Status != model.BackgroundJobStatusRetrying || reloaded.LockedAt != nil || reloaded.LockedBy != "" || reloaded.LastError != context.Canceled.Error() {
+		t.Fatalf("expected cancelled job to be retryable and unlocked: %+v", reloaded)
+	}
+}
+
 // --- Redis Stream path tests ---
 
 func newMiniRedis(t *testing.T) (*miniredis.Miniredis, *redis.Client) {