diff --git a/cmd/worker/main.go b/cmd/worker/main.go new file mode 100644 index 0000000..cdce4fe --- /dev/null +++ b/cmd/worker/main.go @@ -0,0 +1,82 @@ +package main + +import ( + "context" + "flag" + "log/slog" + "os" + "os/signal" + "syscall" + "time" + + "github.com/jackc/pgx/v5/pgxpool" + goredis "github.com/redis/go-redis/v9" + + "github.com/smallchungus/disttaskqueue/internal/queue" + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/worker" +) + +func main() { + stage := flag.String("stage", "", "queue stage name (required)") + flag.Parse() + + logger := slog.New(slog.NewJSONHandler(os.Stdout, nil)) + slog.SetDefault(logger) + + if *stage == "" { + slog.Error("--stage is required") + os.Exit(2) + } + + dsn := envOr("DATABASE_URL", "postgres://dtq:dtq@localhost:5432/dtq?sslmode=disable") + redisURL := envOr("REDIS_URL", "redis://localhost:6379/0") + + ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) + defer stop() + + if err := store.Migrate(ctx, dsn); err != nil { + slog.Error("migrate failed", "err", err) + os.Exit(1) + } + + pool, err := pgxpool.New(ctx, dsn) + if err != nil { + slog.Error("pg connect", "err", err) + os.Exit(1) + } + defer pool.Close() + + opts, err := goredis.ParseURL(redisURL) + if err != nil { + slog.Error("redis url", "err", err) + os.Exit(1) + } + redis := goredis.NewClient(opts) + defer func() { _ = redis.Close() }() + + w := worker.New(worker.Config{ + Stage: *stage, + WorkerID: worker.NewWorkerID(), + Store: store.New(pool), + Queue: queue.New(redis), + Handler: worker.NoopHandler{}, + PopTimeout: 30 * time.Second, + HeartbeatTTL: 15 * time.Second, + HeartbeatInterval: 5 * time.Second, + }) + + slog.Info("worker starting", "stage", *stage) + if err := w.Run(ctx); err != nil { + slog.Error("worker run", "err", err) + os.Exit(1) + } + slog.Info("worker stopped") +} + +func envOr(k, def string) string { + if v := os.Getenv(k); v != "" { + return v + } + return def +} diff --git a/internal/queue/queue.go b/internal/queue/queue.go index c24f82a..831c0f8 100644 --- a/internal/queue/queue.go +++ b/internal/queue/queue.go @@ -4,6 +4,7 @@ import ( "context" "errors" "fmt" + "strconv" "time" goredis "github.com/redis/go-redis/v9" @@ -51,3 +52,14 @@ func (q *Queue) Depth(ctx context.Context, stage string) (int64, error) { } return n, nil } + +// Client exposes the underlying Redis client. Test-only. +func (q *Queue) Client() *goredis.Client { return q.cli } + +func (q *Queue) Heartbeat(ctx context.Context, workerID string, ttl time.Duration) error { + now := strconv.FormatInt(time.Now().Unix(), 10) + if err := q.cli.Set(ctx, "heartbeat:"+workerID, now, ttl).Err(); err != nil { + return fmt.Errorf("heartbeat: %w", err) + } + return nil +} diff --git a/internal/queue/queue_integration_test.go b/internal/queue/queue_integration_test.go index e42c49a..a5eab07 100644 --- a/internal/queue/queue_integration_test.go +++ b/internal/queue/queue_integration_test.go @@ -81,3 +81,27 @@ func TestBlockingPop_ReturnsErrEmptyOnTimeout(t *testing.T) { t.Fatalf("got %v, want ErrEmpty", err) } } + +func TestHeartbeat_SetsKeyWithTTL(t *testing.T) { + q := newQueue(t) + cli := q.Client() + + if err := q.Heartbeat(context.Background(), "worker-xyz", 5*time.Second); err != nil { + t.Fatalf("heartbeat: %v", err) + } + + val, err := cli.Get(context.Background(), "heartbeat:worker-xyz").Result() + if err != nil { + t.Fatalf("get: %v", err) + } + if val == "" { + t.Fatalf("empty value") + } + ttl, err := cli.TTL(context.Background(), "heartbeat:worker-xyz").Result() + if err != nil { + t.Fatalf("ttl: %v", err) + } + if ttl <= 0 || ttl > 5*time.Second { + t.Fatalf("ttl out of range: %v", ttl) + } +} diff --git a/internal/worker/backoff.go b/internal/worker/backoff.go new file mode 100644 index 0000000..9c27ae8 --- /dev/null +++ b/internal/worker/backoff.go @@ -0,0 +1,19 @@ +package worker + +import ( + "math" + "math/rand/v2" + "time" +) + +func Compute(attempts int) time.Duration { + if attempts < 1 { + attempts = 1 + } + base := math.Pow(2, float64(attempts)) + if base > 600 { + base = 600 + } + jitter := 0.75 + rand.Float64()*0.5 //nolint:gosec // jitter spread, not security-sensitive + return time.Duration(base * jitter * float64(time.Second)) +} diff --git a/internal/worker/backoff_test.go b/internal/worker/backoff_test.go new file mode 100644 index 0000000..77118ec --- /dev/null +++ b/internal/worker/backoff_test.go @@ -0,0 +1,34 @@ +package worker + +import ( + "testing" + "time" +) + +func TestCompute_GrowsExponentially(t *testing.T) { + cases := []struct { + attempts int + minD time.Duration + maxD time.Duration + }{ + {1, 1500 * time.Millisecond, 2500 * time.Millisecond}, + {2, 3 * time.Second, 5 * time.Second}, + {3, 6 * time.Second, 10 * time.Second}, + {4, 12 * time.Second, 20 * time.Second}, + } + for _, c := range cases { + t.Run("", func(t *testing.T) { + d := Compute(c.attempts) + if d < c.minD || d > c.maxD { + t.Fatalf("attempts=%d got %v, want in [%v,%v]", c.attempts, d, c.minD, c.maxD) + } + }) + } +} + +func TestCompute_CapsAt600s(t *testing.T) { + d := Compute(20) + if d < 450*time.Second || d > 750*time.Second { + t.Fatalf("got %v, want ~600s ±25%%", d) + } +} diff --git a/internal/worker/handler.go b/internal/worker/handler.go new file mode 100644 index 0000000..d23cb1b --- /dev/null +++ b/internal/worker/handler.go @@ -0,0 +1,33 @@ +package worker + +import ( + "context" + "encoding/json" + "time" + + "github.com/smallchungus/disttaskqueue/internal/store" +) + +type Handler interface { + Process(ctx context.Context, job store.Job) error +} + +type NoopHandler struct{} + +func (NoopHandler) Process(ctx context.Context, job store.Job) error { + var p struct { + SleepMs int `json:"sleepMs"` + } + if len(job.Payload) > 0 { + _ = json.Unmarshal(job.Payload, &p) + } + if p.SleepMs <= 0 { + return nil + } + select { + case <-ctx.Done(): + return ctx.Err() + case <-time.After(time.Duration(p.SleepMs) * time.Millisecond): + return nil + } +} diff --git a/internal/worker/handler_test.go b/internal/worker/handler_test.go new file mode 100644 index 0000000..2d94fb4 --- /dev/null +++ b/internal/worker/handler_test.go @@ -0,0 +1,40 @@ +package worker + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/google/uuid" + "github.com/smallchungus/disttaskqueue/internal/store" +) + +func TestNoopHandler_SleepsThenReturnsNil(t *testing.T) { + h := NoopHandler{} + job := store.Job{ID: uuid.New(), Payload: json.RawMessage(`{"sleepMs":50}`)} + + start := time.Now() + err := h.Process(context.Background(), job) + elapsed := time.Since(start) + + if err != nil { + t.Fatalf("err: %v", err) + } + if elapsed < 50*time.Millisecond || elapsed > 200*time.Millisecond { + t.Fatalf("elapsed %v, want ~50ms", elapsed) + } +} + +func TestNoopHandler_DefaultsToZeroSleep(t *testing.T) { + h := NoopHandler{} + job := store.Job{ID: uuid.New(), Payload: json.RawMessage(`{}`)} + + start := time.Now() + if err := h.Process(context.Background(), job); err != nil { + t.Fatal(err) + } + if time.Since(start) > 50*time.Millisecond { + t.Fatalf("should be instant") + } +} diff --git a/internal/worker/heartbeat_integration_test.go b/internal/worker/heartbeat_integration_test.go new file mode 100644 index 0000000..4b09fd6 --- /dev/null +++ b/internal/worker/heartbeat_integration_test.go @@ -0,0 +1,39 @@ +//go:build integration + +package worker_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/worker" +) + +func TestProcessOne_WritesHeartbeatDuringHandler(t *testing.T) { + h := setupHarness(t, worker.NoopHandler{}) + ctx := context.Background() + + job, _ := h.store.EnqueueJob(ctx, store.NewJob{Stage: "test", Payload: json.RawMessage(`{"sleepMs":1500}`)}) + _ = h.queue.Push(ctx, "test", job.ID.String()) + + done := make(chan struct{}) + go func() { + _, _ = h.w.ProcessOne(ctx) + close(done) + }() + + time.Sleep(800 * time.Millisecond) + + val, err := h.queue.Client().Get(ctx, "heartbeat:worker-test-1").Result() + if err != nil { + t.Fatalf("heartbeat get: %v", err) + } + if val == "" { + t.Fatal("heartbeat not set during handler execution") + } + + <-done +} diff --git a/internal/worker/id.go b/internal/worker/id.go new file mode 100644 index 0000000..0d7592d --- /dev/null +++ b/internal/worker/id.go @@ -0,0 +1,20 @@ +package worker + +import ( + "crypto/rand" + "encoding/hex" + "fmt" + "os" + "strings" +) + +func NewWorkerID() string { + host, err := os.Hostname() + if err != nil || host == "" { + host = "unknown" + } + host = strings.ReplaceAll(host, "-", "") + var b [8]byte + _, _ = rand.Read(b[:]) + return fmt.Sprintf("%s-%s", host, hex.EncodeToString(b[:])) +} diff --git a/internal/worker/id_test.go b/internal/worker/id_test.go new file mode 100644 index 0000000..c864580 --- /dev/null +++ b/internal/worker/id_test.go @@ -0,0 +1,28 @@ +package worker + +import ( + "strings" + "testing" +) + +func TestNewWorkerID_HasHostnamePrefixAndRandomSuffix(t *testing.T) { + id := NewWorkerID() + if !strings.Contains(id, "-") { + t.Fatalf("missing hyphen: %q", id) + } + parts := strings.SplitN(id, "-", 2) + if len(parts[0]) == 0 { + t.Fatalf("empty hostname: %q", id) + } + if len(parts[1]) != 16 { + t.Fatalf("suffix not 16 hex chars: %q", id) + } +} + +func TestNewWorkerID_IsUniqueAcrossCalls(t *testing.T) { + a := NewWorkerID() + b := NewWorkerID() + if a == b { + t.Fatalf("ids collided: %q", a) + } +} diff --git a/internal/worker/worker.go b/internal/worker/worker.go new file mode 100644 index 0000000..a50acfb --- /dev/null +++ b/internal/worker/worker.go @@ -0,0 +1,114 @@ +package worker + +import ( + "context" + "errors" + "fmt" + "log/slog" + "time" + + "github.com/google/uuid" + "github.com/smallchungus/disttaskqueue/internal/queue" + "github.com/smallchungus/disttaskqueue/internal/store" +) + +type Config struct { + Stage string + WorkerID string + Store *store.Store + Queue *queue.Queue + Handler Handler + PopTimeout time.Duration + HeartbeatTTL time.Duration + HeartbeatInterval time.Duration +} + +type Worker struct { + cfg Config +} + +func New(cfg Config) *Worker { + return &Worker{cfg: cfg} +} + +func (w *Worker) Run(ctx context.Context) error { + for { + select { + case <-ctx.Done(): + return nil + default: + } + if _, err := w.ProcessOne(ctx); err != nil { + slog.Error("worker process failed", "err", err, "stage", w.cfg.Stage, "worker_id", w.cfg.WorkerID) + } + } +} + +func (w *Worker) ProcessOne(ctx context.Context) (didWork bool, err error) { + rawID, err := w.cfg.Queue.BlockingPop(ctx, w.cfg.Stage, w.cfg.PopTimeout) + if errors.Is(err, queue.ErrEmpty) { + return false, nil + } + if err != nil { + return false, fmt.Errorf("pop: %w", err) + } + + jobID, err := uuid.Parse(rawID) + if err != nil { + return true, fmt.Errorf("parse job id %q: %w", rawID, err) + } + + if err := w.cfg.Store.ClaimJob(ctx, jobID, w.cfg.WorkerID); err != nil { + if errors.Is(err, store.ErrJobNotClaimable) { + slog.Info("claim race lost", "job_id", jobID, "worker_id", w.cfg.WorkerID) + return true, nil + } + return true, fmt.Errorf("claim: %w", err) + } + + job, err := w.cfg.Store.GetJob(ctx, jobID) + if err != nil { + return true, fmt.Errorf("get: %w", err) + } + + hbCtx, hbCancel := context.WithCancel(ctx) + go w.heartbeatLoop(hbCtx) + + handlerErr := w.cfg.Handler.Process(ctx, job) + + hbCancel() + + if handlerErr != nil { + nextRun := time.Now().Add(Compute(job.Attempts + 1)) + if err := w.cfg.Store.MarkFailed(ctx, jobID, handlerErr.Error(), nextRun); err != nil { + return true, fmt.Errorf("mark failed: %w", err) + } + slog.Info("job failed, will retry", "job_id", jobID, "attempts", job.Attempts+1, "next_run", nextRun) + return true, nil + } + + if err := w.cfg.Store.MarkDone(ctx, jobID); err != nil { + return true, fmt.Errorf("mark done: %w", err) + } + return true, nil +} + +func (w *Worker) heartbeatLoop(ctx context.Context) { + tick := time.NewTicker(w.cfg.HeartbeatInterval) + defer tick.Stop() + + if err := w.cfg.Queue.Heartbeat(ctx, w.cfg.WorkerID, w.cfg.HeartbeatTTL); err != nil { + slog.Warn("heartbeat write failed", "err", err) + } + + for { + select { + case <-ctx.Done(): + return + case <-tick.C: + if err := w.cfg.Queue.Heartbeat(ctx, w.cfg.WorkerID, w.cfg.HeartbeatTTL); err != nil { + slog.Warn("heartbeat write failed", "err", err) + } + } + } +} diff --git a/internal/worker/worker_integration_test.go b/internal/worker/worker_integration_test.go new file mode 100644 index 0000000..999ad8d --- /dev/null +++ b/internal/worker/worker_integration_test.go @@ -0,0 +1,116 @@ +//go:build integration + +package worker_test + +import ( + "context" + "encoding/json" + "testing" + "time" + + "github.com/smallchungus/disttaskqueue/internal/queue" + "github.com/smallchungus/disttaskqueue/internal/store" + "github.com/smallchungus/disttaskqueue/internal/testutil" + "github.com/smallchungus/disttaskqueue/internal/worker" +) + +type harness struct { + store *store.Store + queue *queue.Queue + w *worker.Worker +} + +func setupHarness(t *testing.T, h worker.Handler) *harness { + t.Helper() + pool := testutil.StartPostgres(t) + if err := store.Migrate(context.Background(), pool.Config().ConnString()); err != nil { + t.Fatal(err) + } + s := store.New(pool) + q := queue.New(testutil.StartRedis(t)) + w := worker.New(worker.Config{ + Stage: "test", + WorkerID: "worker-test-1", + Store: s, + Queue: q, + Handler: h, + PopTimeout: 500 * time.Millisecond, + HeartbeatTTL: 5 * time.Second, + HeartbeatInterval: 1 * time.Second, + }) + return &harness{store: s, queue: q, w: w} +} + +func TestProcessOne_RunsHandlerAndMarksDone(t *testing.T) { + h := setupHarness(t, worker.NoopHandler{}) + ctx := context.Background() + + job, err := h.store.EnqueueJob(ctx, store.NewJob{Stage: "test", Payload: json.RawMessage(`{"sleepMs":20}`)}) + if err != nil { + t.Fatal(err) + } + if err := h.queue.Push(ctx, "test", job.ID.String()); err != nil { + t.Fatal(err) + } + + didWork, err := h.w.ProcessOne(ctx) + if err != nil { + t.Fatalf("ProcessOne: %v", err) + } + if !didWork { + t.Fatal("expected didWork=true") + } + + got, _ := h.store.GetJob(ctx, job.ID) + if got.Status != store.StatusDone { + t.Fatalf("status: %s, want done", got.Status) + } +} + +func TestProcessOne_ReturnsNoWorkOnEmptyQueue(t *testing.T) { + h := setupHarness(t, worker.NoopHandler{}) + didWork, err := h.w.ProcessOne(context.Background()) + if err != nil { + t.Fatalf("err: %v", err) + } + if didWork { + t.Fatal("expected didWork=false") + } +} + +type errHandler struct{} + +func (errHandler) Process(ctx context.Context, job store.Job) error { + return &handlerErr{"boom"} +} + +type handlerErr struct{ msg string } + +func (e *handlerErr) Error() string { return e.msg } + +func TestProcessOne_MarksFailedOnHandlerError(t *testing.T) { + h := setupHarness(t, errHandler{}) + ctx := context.Background() + + job, _ := h.store.EnqueueJob(ctx, store.NewJob{Stage: "test"}) + _ = h.queue.Push(ctx, "test", job.ID.String()) + + didWork, err := h.w.ProcessOne(ctx) + if err != nil { + t.Fatalf("ProcessOne: %v", err) + } + if !didWork { + t.Fatal("expected didWork=true") + } + + got, _ := h.store.GetJob(ctx, job.ID) + if got.Status != store.StatusQueued { + t.Fatalf("status: %s, want queued (retry)", got.Status) + } + if got.Attempts != 1 { + t.Fatalf("attempts: %d, want 1", got.Attempts) + } + if got.LastError == nil || *got.LastError != "boom" { + t.Fatalf("last_error: %v", got.LastError) + } +}