feat(worker): Phase 2C — Worker scaffold (claim, heartbeat, retry)

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
+}

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
+}

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)
+	}
+}

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))
+}

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)
+	}
+}

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
+	}
+}

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")
+	}
+}

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
+}

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[:]))
+}

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)
+	}
+}