From 985d0a63cc970491edd23081c30b94a20130a357 Mon Sep 17 00:00:00 2001 From: Raphael Simon Date: Tue, 6 Oct 2026 00:36:56 -0700 Subject: [PATCH] Save new worker registration with its first heartbeat --- pool/node.go | 30 +- pool/worker.go | 35 +- pool/worker_registration.go | 155 +++++++ pool/worker_registration_lifecycle_test.go | 513 +++++++++++++++++++++ pool/worker_registration_test.go | 400 ++++++++++++++++ 5 files changed, 1102 insertions(+), 31 deletions(-) create mode 100644 pool/worker_registration.go create mode 100644 pool/worker_registration_lifecycle_test.go create mode 100644 pool/worker_registration_test.go diff --git a/pool/node.go b/pool/node.go index 382ebb3..26079d1 100644 --- a/pool/node.go +++ b/pool/node.go @@ -538,24 +538,42 @@ func AddNode(ctx context.Context, poolName string, rdb *redis.Client, opts ...No // NotificationHandler and MessageHandler interfaces to handle job-scoped // notifications and hash-routed messages. func (node *Node) AddWorker(ctx context.Context, handler JobHandler) (*Worker, error) { + // Reject local closure before storage access. The Redis lifecycle check + // runs without this lock because retirement must acquire its write side. node.lock.RLock() - defer node.lock.RUnlock() if node.closing { + node.lock.RUnlock() return nil, fmt.Errorf("AddWorker: pool %q is closed", node.PoolName) } if node.clientOnly { + node.lock.RUnlock() return nil, fmt.Errorf("AddWorker: pool %q is client-only", node.PoolName) } + node.lock.RUnlock() if err := node.ensureGenerationActive(ctx); err != nil { return nil, fmt.Errorf("AddWorker: %w", err) } + + // Local close waits for construction and publication. Redis registration + // checks the pool again, so retirement after the first check cannot admit work. + node.lock.RLock() + if node.closing { + node.lock.RUnlock() + return nil, fmt.Errorf("AddWorker: pool %q is closed", node.PoolName) + } w, err := newWorker(ctx, node, handler) - if err != nil { - return nil, err + if err == nil { + node.localWorkers.Store(w.ID, w) + node.workerStreams.Store(w.ID, w.stream) + } + node.lock.RUnlock() + + // An ended pool or node is retired after releasing the read lock. Keep + // the original failure even when the final lifecycle check also fails. + if errors.Is(err, ErrPoolGenerationLost) { + return nil, errors.Join(err, node.ensureGenerationActive(context.WithoutCancel(ctx))) } - node.localWorkers.Store(w.ID, w) - node.workerStreams.Store(w.ID, w.stream) - return w, nil + return w, err } // RemoveWorker joins worker intake and waits for Stop to release every locally diff --git a/pool/worker.go b/pool/worker.go index 5d51d47..be1be97 100644 --- a/pool/worker.go +++ b/pool/worker.go @@ -10,7 +10,6 @@ import ( "errors" "fmt" "sort" - "strconv" "sync" "time" @@ -135,9 +134,6 @@ var ( // newWorker creates a new worker. func newWorker(ctx context.Context, node *Node, h JobHandler) (*Worker, error) { - if err := node.ensureGenerationActive(ctx); err != nil { - return nil, err - } wid := ulid.Make().String() createdAt, err := node.rdb.Time(ctx).Result() if err != nil { @@ -162,30 +158,19 @@ func newWorker(ctx context.Context, node *Node, h JobHandler) (*Worker, error) { destroyErr, ) } - if err := node.setPoolMapAndWait( - ctx, - node.workerMap, - node.resources.workers, - wid, - strconv.FormatInt(createdAt.UnixNano(), 10), - ); err != nil { - reader.Close() - destroyErr := stream.Destroy(context.WithoutCancel(ctx)) - return nil, errors.Join( - fmt.Errorf("failed to add worker %q to pool %q: %w", wid, node.PoolName, err), - destroyErr, - ) - } - now, err := node.updateWorkerHeartbeat(ctx, wid) - if err == nil { - err = waitPoolMapValue(ctx, node.workerKeepAliveMap, wid, now) - } - if err != nil { - removeErr := node.deletePoolMap(context.WithoutCancel(ctx), node.resources.workers, wid) + // Save discovery and the first heartbeat together so cleanup cannot + // observe a newly published worker without a heartbeat. + if err := node.registerWorker(ctx, wid, createdAt, stream); err != nil { reader.Close() destroyErr := stream.Destroy(context.WithoutCancel(ctx)) + var removeErr error + // Retire the stream before deleting discovery. If retirement fails, + // keep the entry so another node can finish the worker's cleanup. + if destroyErr == nil { + removeErr = node.removeWorkerRegistration(context.WithoutCancel(ctx), wid, createdAt, stream) + } return nil, errors.Join( - fmt.Errorf("failed to update worker keep-alive: %w", err), + fmt.Errorf("failed to register worker %q in pool %q: %w", wid, node.PoolName, err), removeErr, destroyErr, ) diff --git a/pool/worker_registration.go b/pool/worker_registration.go new file mode 100644 index 0000000..9901446 --- /dev/null +++ b/pool/worker_registration.go @@ -0,0 +1,155 @@ +// Worker registration saves the worker's identity and first heartbeat together. +// Pool cleanup therefore observes either no worker or a worker with a current +// heartbeat. Repeated registration cannot revive a destroyed worker stream. +package pool + +import ( + "context" + "fmt" + "strconv" + "time" + + redis "github.com/redis/go-redis/v9" + + "goa.design/pulse/streaming" +) + +var ( + registerWorkerScript = redis.NewScript(` +if redis.call("HGET", KEYS[1], "state") ~= ARGV[1] +or redis.call("HGET", KEYS[1], "generation") ~= ARGV[2] then + return redis.error_reply("POOLGENERATIONLOST") +end +if redis.call("HGET", KEYS[7], ARGV[5]) then + return redis.error_reply("NODECLEANUPLOST") +end +if redis.call("HGET", KEYS[8], "state") ~= ARGV[1] +or redis.call("HGET", KEYS[8], "generation") ~= ARGV[6] then + return redis.error_reply("STREAMDESTROYED") +end +local registration = redis.call("HGET", KEYS[2], ARGV[3]) +if (registration and registration ~= ARGV[4]) +or redis.call("HGET", KEYS[6], ARGV[3]) then + return redis.error_reply("WORKERCLEANUPLOST") +end +-- Repeating this creation returns the saved heartbeat without replacing the +-- timestamp that the constructor is waiting to observe in its local map. +if registration then + local heartbeat = redis.call("HGET", KEYS[4], ARGV[3]) + if not heartbeat or not tonumber(heartbeat) then + return redis.error_reply("WORKERHEARTBEATINVALID") + end + return heartbeat +end +local clock = redis.call("TIME") +local timestamp = clock[1] .. string.format("%06d", clock[2]) .. "000" +local function save(content, channel, value) + redis.call("HSET", content, ARGV[3], value) + local rev = tostring(redis.call("HINCRBY", content, "=rev", 1)) + redis.call("HSET", content, "=kind", "set") + local message = struct.pack( + "ic0ic0ic0", + string.len(ARGV[3]), ARGV[3], + string.len(value), value, + string.len(rev), rev + ) + redis.call("PUBLISH", channel, "set:" .. message) +end +save(KEYS[4], KEYS[5], timestamp) +save(KEYS[2], KEYS[3], ARGV[4]) +return timestamp +`) + + removeWorkerRegistrationScript = redis.NewScript(` +if redis.call("HGET", KEYS[1], "state") ~= ARGV[1] +or redis.call("HGET", KEYS[1], "generation") ~= ARGV[2] then + return redis.error_reply("POOLGENERATIONLOST") +end +if redis.call("HGET", KEYS[6], "state") ~= "destroyed" +or redis.call("HGET", KEYS[6], "generation") ~= ARGV[5] then + return redis.error_reply("STREAMDESTROYED") +end +local registration = redis.call("HGET", KEYS[2], ARGV[3]) +if (registration and registration ~= ARGV[4]) +or redis.call("HGET", KEYS[7], ARGV[3]) then + return redis.error_reply("WORKERCLEANUPLOST") +end +local function remove(content, channel) + if redis.call("HDEL", content, ARGV[3]) == 0 then + return + end + local rev = tostring(redis.call("HINCRBY", content, "=rev", 1)) + redis.call("HSET", content, "=kind", "del") + local message = struct.pack( + "ic0ic0", + string.len(ARGV[3]), ARGV[3], + string.len(rev), rev + ) + redis.call("PUBLISH", channel, "del:" .. message) +end +remove(KEYS[2], KEYS[3]) +remove(KEYS[4], KEYS[5]) +return 1 +`) +) + +// registerWorker saves a worker and its first Redis-time heartbeat while the +// pool, node, and original worker stream remain active. It waits for both local +// maps to observe the saved values before the constructor starts the worker. +// Repeated Redis commands reuse the saved heartbeat rather than replace the +// timestamp that an earlier successful response returned. +func (node *Node) registerWorker(ctx context.Context, workerID string, createdAt time.Time, stream *streaming.Stream) error { + creation := strconv.FormatInt(createdAt.UnixNano(), 10) + heartbeat, err := registerWorkerScript.Run( + ctx, + node.rdb, + []string{ + fmt.Sprintf("pulse:stream:%s:lifecycle", node.poolStream.Name), + rmapContentKey(node.resources.workers), + rmapUpdateChannel(node.resources.workers), + rmapContentKey(node.resources.workerKeepAlive), + rmapUpdateChannel(node.resources.workerKeepAlive), + rmapContentKey(node.resources.workerCleanup), + rmapContentKey(node.resources.nodeKeepAlive), + fmt.Sprintf("pulse:stream:%s:lifecycle", stream.Name), + }, + "active", + node.resources.generation, + workerID, + creation, + nodeCleanupField(node.ID), + stream.Generation(), + ).Text() + if err != nil { + return poolBoundaryError(err) + } + if err := waitPoolMapValue(ctx, node.workerMap, workerID, creation); err != nil { + return err + } + return waitPoolMapValue(ctx, node.workerKeepAliveMap, workerID, heartbeat) +} + +// removeWorkerRegistration removes both saved entries for a failed constructor +// only after its original stream is destroyed. A changed pool or worker stream +// generation, or another cleanup owner's claim, is rejected without changing +// the records that owner needs to finish cleanup. +func (node *Node) removeWorkerRegistration(ctx context.Context, workerID string, createdAt time.Time, stream *streaming.Stream) error { + return poolBoundaryError(removeWorkerRegistrationScript.Run( + ctx, + node.rdb, + []string{ + fmt.Sprintf("pulse:stream:%s:lifecycle", node.poolStream.Name), + rmapContentKey(node.resources.workers), + rmapUpdateChannel(node.resources.workers), + rmapContentKey(node.resources.workerKeepAlive), + rmapUpdateChannel(node.resources.workerKeepAlive), + fmt.Sprintf("pulse:stream:%s:lifecycle", stream.Name), + rmapContentKey(node.resources.workerCleanup), + }, + "active", + node.resources.generation, + workerID, + strconv.FormatInt(createdAt.UnixNano(), 10), + stream.Generation(), + ).Err()) +} diff --git a/pool/worker_registration_lifecycle_test.go b/pool/worker_registration_lifecycle_test.go new file mode 100644 index 0000000..76b7bc4 --- /dev/null +++ b/pool/worker_registration_lifecycle_test.go @@ -0,0 +1,513 @@ +// These tests interrupt AddWorker at its Redis registration command. Failed +// replies must leave unfinished cleanup discoverable. Each saved stream version, +// called its generation, must reject an old constructor after that version ends. +package pool + +import ( + "context" + "errors" + "fmt" + "reflect" + "strconv" + "sync" + "sync/atomic" + "testing" + "time" + + redis "github.com/redis/go-redis/v9" + + "goa.design/pulse/pulse" + "goa.design/pulse/streaming" + ptesting "goa.design/pulse/testing" +) + +type ( + workerLifecycleContextKey struct{} + + workerLifecycleRegistration struct { + workerID string + createdAt string + generation string + lifecycleKey string + } + + workerLifecycleHook struct { + workersKey string + afterSave bool + replyError error + destroyError error + observed chan workerLifecycleRegistration + resume chan struct{} + releaseOnce sync.Once + intercepted atomic.Bool + destroyAttempts atomic.Int32 + registration workerLifecycleRegistration + } + + workerLifecycleResult struct { + worker *Worker + err error + } + + workerLifecycleStreamSnapshot struct { + lifecycle map[string]string + events []redis.XMessage + } +) + +func TestWorkerRegistrationFailedDestroyPreservesDiscovery(t *testing.T) { + ctx, rdb := newWorkerLifecycleRedis(t) + node := newWorkerLifecycleNode(t, ctx, rdb) + replyError := errors.New("synthetic registration reply failure") + destroyError := errors.New("synthetic exact stream destruction failure") + hook := newWorkerLifecycleHook(t, rdb, node) + hook.afterSave = true + hook.replyError = replyError + hook.destroyError = destroyError + result := startWorkerLifecycleConstructor(ctx, node) + registration := awaitWorkerLifecycleRegistration(t, ctx, hook) + original := openWorkerLifecycleStream(t, ctx, rdb, registration) + before := workerLifecycleMaps(t, ctx, rdb, node) + streamBefore := readWorkerLifecycleStream(t, ctx, rdb, registration.lifecycleKey) + if before[0][registration.workerID] != registration.createdAt { + t.Fatalf("saved worker creation = %q, want %q", + before[0][registration.workerID], registration.createdAt) + } + heartbeat, err := strconv.ParseInt(before[1][registration.workerID], 10, 64) + if err != nil || heartbeat <= 0 { + t.Fatalf("first heartbeat = %q, want a positive Redis timestamp: %v", + before[1][registration.workerID], err) + } + creation, err := strconv.ParseInt(registration.createdAt, 10, 64) + if err != nil || heartbeat < creation { + t.Fatalf("heartbeat %d precedes worker creation %q: %v", heartbeat, registration.createdAt, err) + } + + hook.release() + got := awaitWorkerLifecycleResult(t, ctx, result) + if got.worker != nil || !errors.Is(got.err, replyError) || !errors.Is(got.err, destroyError) { + t.Fatalf("AddWorker = (%v, %v), want no worker and both injected errors", got.worker, got.err) + } + if attempts := hook.destroyAttempts.Load(); attempts != 1 { + t.Errorf("exact stream destruction attempts = %d, want 1", attempts) + } + if after := workerLifecycleMaps(t, ctx, rdb, node); !reflect.DeepEqual(before, after) { + t.Errorf("failed destruction changed discovery or revisions: before %v, after %v", before, after) + } + if streamAfter := readWorkerLifecycleStream(t, ctx, rdb, registration.lifecycleKey); !reflect.DeepEqual(streamBefore, streamAfter) { + t.Errorf("failed destruction changed the active stream: before %v, after %v", streamBefore, streamAfter) + } + requireWorkerLifecycleState(t, ctx, rdb, registration.lifecycleKey, original.Generation(), "active") + requireNoWorkerLifecycleAdmission(t, node, registration.workerID) + + // The failed constructor has no heartbeat writer. Make its saved heartbeat + // stale, then let the existing cleanup owner destroy the stream and records. + now, err := rdb.Time(ctx).Result() + if err != nil { + t.Fatalf("read Redis cleanup time: %v", err) + } + if err := rdb.HSet(ctx, rmapContentKey(node.resources.workerKeepAlive), + registration.workerID, strconv.FormatInt(now.Add(-2*node.workerTTL).UnixNano(), 10)).Err(); err != nil { + t.Fatalf("set synthetic stale worker heartbeat: %v", err) + } + lease, err := node.acquireWorkerCleanup(ctx, registration.workerID) + if err != nil || lease == nil { + t.Fatalf("acquire worker cleanup = (%v, %v), want an owned lease", lease, err) + } + complete := node.requeueWorkerJobs(ctx, lease) + if err := node.releaseWorkerCleanup(ctx, lease, complete); err != nil { + t.Fatalf("release worker cleanup: %v", err) + } + if !complete { + t.Fatal("existing cleanup did not finish the failed constructor's worker") + } + for index, entries := range workerLifecycleMaps(t, ctx, rdb, node) { + if _, exists := entries[registration.workerID]; exists { + t.Errorf("cleanup retained worker in map %d", index) + } + } + requireWorkerLifecycleState(t, ctx, rdb, registration.lifecycleKey, original.Generation(), "destroyed") + if err := original.Open(ctx); !errors.Is(err, streaming.ErrStreamDestroyed) { + t.Errorf("cleaned original stream Open = %v, want ErrStreamDestroyed", err) + } +} + +func TestWorkerRegistrationLifecycleLossRetiresNode(t *testing.T) { + for _, loss := range []string{"pool generation", "node cleanup"} { + t.Run(loss, func(t *testing.T) { + ctx, rdb := newWorkerLifecycleRedis(t) + node := newWorkerLifecycleNode(t, ctx, rdb) + hook := newWorkerLifecycleHook(t, rdb, node) + result := startWorkerLifecycleConstructor(ctx, node) + registration := awaitWorkerLifecycleRegistration(t, ctx, hook) + original := openWorkerLifecycleStream(t, ctx, rdb, registration) + before := workerLifecycleMaps(t, ctx, rdb, node) + for index, entries := range before { + if _, exists := entries[registration.workerID]; exists { + t.Fatalf("worker already registered in map %d before the paused command", index) + } + } + + // AddWorker already passed its pool and node checks. End the actual + // generation or clean up the stale node before registration executes. + switch loss { + case "pool generation": + if err := node.poolStream.Destroy(ctx); err != nil { + t.Fatalf("destroy original pool generation: %v", err) + } + case "node cleanup": + now, err := rdb.Time(ctx).Result() + if err != nil { + t.Fatalf("read Redis cleanup time: %v", err) + } + if err := rdb.HSet(ctx, rmapContentKey(node.resources.nodeKeepAlive), + node.ID, strconv.FormatInt(now.Add(-2*node.workerTTL).UnixNano(), 10)).Err(); err != nil { + t.Fatalf("set synthetic stale node heartbeat: %v", err) + } + cleaned, err := cleanupStalePoolNode(ctx, rdb, node.resources, + node.ID, newNodeCleanupOwner("synthetic-cleanup")) + if err != nil || !cleaned { + t.Fatalf("cleanup stale node = (%v, %v), want completed cleanup", cleaned, err) + } + requireWorkerLifecycleState(t, ctx, rdb, + "pulse:stream:"+node.nodeStream.Name+":lifecycle", node.nodeStream.Generation(), "destroyed") + } + hook.release() + got := awaitWorkerLifecycleResult(t, ctx, result) + if got.worker != nil || !errors.Is(got.err, ErrPoolGenerationLost) { + t.Fatalf("AddWorker = (%v, %v), want no worker and ErrPoolGenerationLost", got.worker, got.err) + } + select { + case <-node.closed: + case <-ctx.Done(): + t.Fatalf("node did not retire automatically after %s: %v", loss, ctx.Err()) + } + if !node.IsClosed() { + t.Error("node closure channel completed without IsClosed") + } + requireNoWorkerLifecycleAdmission(t, node, registration.workerID) + if after := workerLifecycleMaps(t, ctx, rdb, node); !reflect.DeepEqual(before, after) { + t.Errorf("rejected registration changed worker maps or revisions: before %v, after %v", before, after) + } + requireWorkerLifecycleState(t, ctx, rdb, registration.lifecycleKey, original.Generation(), "destroyed") + }) + } +} + +func TestWorkerRegistrationRejectsReplacedStreamGeneration(t *testing.T) { + ctx, rdb := newWorkerLifecycleRedis(t) + node := newWorkerLifecycleNode(t, ctx, rdb) + hook := newWorkerLifecycleHook(t, rdb, node) + result := startWorkerLifecycleConstructor(ctx, node) + registration := awaitWorkerLifecycleRegistration(t, ctx, hook) + original := openWorkerLifecycleStream(t, ctx, rdb, registration) + if err := original.Destroy(ctx); err != nil { + t.Fatalf("destroy original worker stream: %v", err) + } + successor, err := streaming.NewStream(original.Name, rdb) + if err != nil { + t.Fatalf("create successor stream handle: %v", err) + } + if _, err := successor.Add(ctx, "synthetic-successor", []byte("synthetic payload")); err != nil { + t.Fatalf("activate successor worker stream: %v", err) + } + if successor.Generation() == original.Generation() { + t.Fatalf("successor generation = %q, want a new generation", successor.Generation()) + } + before := workerLifecycleMaps(t, ctx, rdb, node) + for index, entries := range before { + if _, exists := entries[registration.workerID]; exists { + t.Fatalf("worker already registered in map %d before the paused command", index) + } + } + streamBefore := readWorkerLifecycleStream(t, ctx, rdb, registration.lifecycleKey) + requireWorkerLifecycleState(t, ctx, rdb, registration.lifecycleKey, successor.Generation(), "active") + + hook.release() + got := awaitWorkerLifecycleResult(t, ctx, result) + if got.worker != nil || !redis.HasErrorPrefix(got.err, "STREAMDESTROYED") { + t.Fatalf("AddWorker = (%v, %v), want the Redis STREAMDESTROYED registration rejection", got.worker, got.err) + } + if !errors.Is(got.err, streaming.ErrStreamDestroyed) { + t.Errorf("old stream rollback = %v, want ErrStreamDestroyed", got.err) + } + requireNoWorkerLifecycleAdmission(t, node, registration.workerID) + if after := workerLifecycleMaps(t, ctx, rdb, node); !reflect.DeepEqual(before, after) { + t.Errorf("old registration changed worker maps or revisions: before %v, after %v", before, after) + } + if streamAfter := readWorkerLifecycleStream(t, ctx, rdb, registration.lifecycleKey); !reflect.DeepEqual(streamBefore, streamAfter) { + t.Errorf("old constructor changed successor lifecycle or events: before %v, after %v", streamBefore, streamAfter) + } + if err := successor.Open(ctx); err != nil { + t.Errorf("successor is no longer active: %v", err) + } + if err := successor.Destroy(ctx); err != nil { + t.Errorf("destroy synthetic successor: %v", err) + } +} + +// newWorkerLifecycleRedis leases one database through the shared test helper. +// Each test has a deadline and flushes only the data in its leased database. +func newWorkerLifecycleRedis(t *testing.T) (context.Context, *redis.Client) { + t.Helper() + ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second) + t.Cleanup(cancel) + rdb := ptesting.NewRedisClient(t) + t.Cleanup(func() { + cleanupCtx, cleanupCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cleanupCancel() + if err := rdb.FlushDB(cleanupCtx).Err(); err != nil { + t.Errorf("flush leased test database: %v", err) + } + if err := rdb.Close(); err != nil { + t.Errorf("close leased test client: %v", err) + } + }) + return ctx, rdb +} + +// newWorkerLifecycleNode starts a public pool node with heartbeat intervals +// longer than the test deadline, so only the test changes synthetic liveness. +func newWorkerLifecycleNode(t *testing.T, ctx context.Context, rdb *redis.Client) *Node { + t.Helper() + node, err := AddNode(ctx, t.Name(), rdb, WithLogger(pulse.NoopLogger()), + WithWorkerTTL(time.Minute), WithJobSinkBlockDuration(testJobSinkBlockDuration)) + if err != nil { + t.Fatalf("AddNode: %v", err) + } + t.Cleanup(func() { + select { + case <-node.closed: + return + default: + } + cleanupCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + done := make(chan error, 1) + go func() { + done <- node.Shutdown(cleanupCtx) + }() + select { + case err := <-done: + if err != nil { + t.Errorf("shutdown test node: %v", err) + } + case <-cleanupCtx.Done(): + t.Errorf("shutdown test node did not finish: %v", cleanupCtx.Err()) + } + }) + return node +} + +// newWorkerLifecycleHook marks one constructor command for interruption and +// releases it during cleanup, including when an assertion stops the test. +func newWorkerLifecycleHook(t *testing.T, rdb *redis.Client, node *Node) *workerLifecycleHook { + t.Helper() + hook := &workerLifecycleHook{ + workersKey: rmapContentKey(node.resources.workers), + observed: make(chan workerLifecycleRegistration, 1), + resume: make(chan struct{}), + } + t.Cleanup(hook.release) + rdb.AddHook(hook) + return hook +} + +// startWorkerLifecycleConstructor calls the public AddWorker method with a +// marked context and returns its result through a buffered channel. +func startWorkerLifecycleConstructor(ctx context.Context, node *Node) <-chan workerLifecycleResult { + result := make(chan workerLifecycleResult, 1) + go func() { + worker, err := node.AddWorker( + context.WithValue(ctx, workerLifecycleContextKey{}, true), newMockJobHandler()) + result <- workerLifecycleResult{worker: worker, err: err} + }() + return result +} + +// awaitWorkerLifecycleRegistration returns the actual arguments intercepted +// during construction, or fails the test when its deadline expires. +func awaitWorkerLifecycleRegistration( + t *testing.T, ctx context.Context, hook *workerLifecycleHook, +) workerLifecycleRegistration { + t.Helper() + select { + case registration := <-hook.observed: + return registration + case <-ctx.Done(): + t.Fatalf("constructor did not reach registration: %v", ctx.Err()) + return workerLifecycleRegistration{} + } +} + +// awaitWorkerLifecycleResult returns the public constructor's outcome, or +// fails the test when construction does not finish before its deadline. +func awaitWorkerLifecycleResult( + t *testing.T, ctx context.Context, result <-chan workerLifecycleResult, +) workerLifecycleResult { + t.Helper() + select { + case got := <-result: + return got + case <-ctx.Done(): + t.Fatalf("AddWorker did not return before its deadline: %v", ctx.Err()) + return workerLifecycleResult{} + } +} + +// openWorkerLifecycleStream opens the stream named by the paused constructor +// and checks that its actual generation matches the registration arguments. +func openWorkerLifecycleStream( + t *testing.T, ctx context.Context, rdb *redis.Client, registration workerLifecycleRegistration, +) *streaming.Stream { + t.Helper() + stream, err := streaming.NewStream(workerStreamName(registration.workerID), rdb) + if err != nil { + t.Fatalf("create original worker stream handle: %v", err) + } + if err := stream.Open(ctx); err != nil { + t.Fatalf("open original worker stream: %v", err) + } + if stream.Generation() != registration.generation { + t.Fatalf("original generation = %q, registration names %q", stream.Generation(), registration.generation) + } + return stream +} + +// workerLifecycleMaps reads both complete worker hashes, including their +// revision fields, so a rejected constructor cannot hide a write and deletion. +func workerLifecycleMaps( + t *testing.T, ctx context.Context, rdb *redis.Client, node *Node, +) [2]map[string]string { + t.Helper() + var snapshot [2]map[string]string + for index, name := range []string{node.resources.workers, node.resources.workerKeepAlive} { + entries, err := rdb.HGetAll(ctx, rmapContentKey(name)).Result() + if err != nil { + t.Fatalf("read worker map %q: %v", name, err) + } + snapshot[index] = entries + } + return snapshot +} + +// readWorkerLifecycleStream reads the actual lifecycle and at most two events. +// The tests create one event, so missing or additional events fail immediately. +func readWorkerLifecycleStream( + t *testing.T, ctx context.Context, rdb *redis.Client, key string, +) workerLifecycleStreamSnapshot { + t.Helper() + lifecycle, err := rdb.HGetAll(ctx, key).Result() + if err != nil { + t.Fatalf("read stream lifecycle %q: %v", key, err) + } + physical := lifecycle["physical_key"] + if physical == "" { + t.Fatalf("stream lifecycle %q has no physical stream key", key) + } + events, err := rdb.XRangeN(ctx, physical, "-", "+", 2).Result() + if err != nil || len(events) != 1 { + t.Fatalf("read stream %q = (%v, %v), want one synthetic event", physical, events, err) + } + return workerLifecycleStreamSnapshot{lifecycle: lifecycle, events: events} +} + +// requireWorkerLifecycleState checks the saved stream version and state so a +// failed old constructor cannot silently destroy a newer stream. +func requireWorkerLifecycleState( + t *testing.T, ctx context.Context, rdb *redis.Client, key, generation, state string, +) { + t.Helper() + lifecycle, err := rdb.HGetAll(ctx, key).Result() + if err != nil { + t.Fatalf("read stream lifecycle %q: %v", key, err) + } + if lifecycle["generation"] != generation || lifecycle["state"] != state { + t.Errorf("stream lifecycle %q = %v, want generation %q and state %q", key, lifecycle, generation, state) + } +} + +// requireNoWorkerLifecycleAdmission checks both local records and the public +// worker list so failed construction cannot leave a worker available for jobs. +func requireNoWorkerLifecycleAdmission(t *testing.T, node *Node, workerID string) { + t.Helper() + if _, exists := node.localWorkers.Load(workerID); exists { + t.Error("failed constructor admitted a local worker") + } + if _, exists := node.workerStreams.Load(workerID); exists { + t.Error("failed constructor published a local worker stream") + } + if workers := node.Workers(); len(workers) != 0 { + t.Errorf("public Workers returned %d workers after failed construction", len(workers)) + } +} + +func (h *workerLifecycleHook) DialHook(next redis.DialHook) redis.DialHook { + return next +} + +// ProcessHook interrupts only the marked constructor's registration. A failed +// destroy matches that constructor's exact stream and generation; background +// Redis work and subsequent cleanup run normally. +func (h *workerLifecycleHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if ctx.Value(workerLifecycleContextKey{}) != true || + (cmd.Name() != "evalsha" && cmd.Name() != "eval") { + return next(ctx, cmd) + } + args := cmd.Args() + if h.destroyError != nil && len(args) == 13 && args[2] == 6 && + args[3] == h.registration.lifecycleKey && args[9] == "active" && + args[10] == h.registration.generation && args[11] == "destroyed" && + args[12] == "physical_key" { + h.destroyAttempts.Add(1) + return h.destroyError + } + if len(args) != 17 || args[2] != 8 || args[4] != h.workersKey { + return next(ctx, cmd) + } + var err error + if h.afterSave { + err = next(ctx, cmd) + if err != nil { + return err + } + } + if h.intercepted.CompareAndSwap(false, true) { + workerID, workerOK := args[13].(string) + createdAt, creationOK := args[14].(string) + generation, generationOK := args[16].(string) + lifecycleKey, keyOK := args[10].(string) + if !workerOK || !creationOK || !generationOK || !keyOK { + return fmt.Errorf("registration command contains unexpected argument types") + } + h.registration = workerLifecycleRegistration{ + workerID: workerID, createdAt: createdAt, + generation: generation, lifecycleKey: lifecycleKey, + } + h.observed <- h.registration + select { + case <-h.resume: + case <-ctx.Done(): + return ctx.Err() + } + if h.afterSave { + return h.replyError + } + } + if h.afterSave { + return err + } + return next(ctx, cmd) + } +} + +func (h *workerLifecycleHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return next +} + +func (h *workerLifecycleHook) release() { + h.releaseOnce.Do(func() { + close(h.resume) + }) +} diff --git a/pool/worker_registration_test.go b/pool/worker_registration_test.go new file mode 100644 index 0000000..60d81c2 --- /dev/null +++ b/pool/worker_registration_test.go @@ -0,0 +1,400 @@ +// These tests pause worker creation after its Redis registration is saved. +// Another pool node then checks that worker for cleanup, so a new worker must +// already have a current heartbeat before cleanup can observe it. +package pool + +import ( + "context" + "errors" + "strconv" + "sync" + "testing" + "time" + + redis "github.com/redis/go-redis/v9" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + "goa.design/pulse/pulse" + ptesting "goa.design/pulse/testing" +) + +type ( + workerRegistrationContextKey struct{} + workerRegistrationReplayContextKey struct{} + + workerRegistrationHook struct { + workersKey string + rdb *redis.Client + registered chan string + resume chan struct{} + once sync.Once + replyError error + args []any + replaying chan struct{} + replay chan struct{} + } + + workerRegistrationResult struct { + worker *Worker + err error + } +) + +func TestWorkerRegistrationIncludesFirstHeartbeat(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + other := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, other.Close(context.Background())) + assert.NoError(t, node.Shutdown(context.Background())) + }() + hook := &workerRegistrationHook{ + workersKey: rmapContentKey(node.resources.workers), + rdb: rdb, + registered: make(chan string, 1), + resume: make(chan struct{}), + } + rdb.AddHook(hook) + result := make(chan workerRegistrationResult, 1) + go func() { + worker, err := node.AddWorker( + context.WithValue(ctx, workerRegistrationContextKey{}, true), + newMockJobHandler(), + ) + result <- workerRegistrationResult{worker: worker, err: err} + }() + var workerID string + select { + case workerID = <-hook.registered: + case <-ctx.Done(): + close(hook.resume) + <-result + require.NoError(t, ctx.Err(), "worker registration was not observed") + } + lease, err := other.acquireWorkerCleanup(ctx, workerID) + assert.NoError(t, err) + assert.Nil(t, lease, "cleanup must not acquire a newly registered worker") + close(hook.resume) + got := <-result + assert.NoError(t, got.err) + assert.NotNil(t, got.worker) + if lease != nil { + assert.NoError(t, other.releaseWorkerCleanup(ctx, lease, false)) + } +} + +func TestWorkerRegistrationCannotReviveRetiredWorker(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, node.Shutdown(context.Background())) + }() + worker := newTestWorker(t, ctx, node) + + // Stop local heartbeat writers while keeping the stream active, matching + // the constructor's state when a Redis operation is repeated. + require.NoError(t, worker.stopLocal()) + require.NoError(t, node.registerWorker(ctx, worker.ID, worker.CreatedAt, worker.stream)) + createdAt, err := rdb.HGet(ctx, rmapContentKey(node.resources.workers), worker.ID).Result() + require.NoError(t, err) + require.Equal(t, strconv.FormatInt(worker.CreatedAt.UnixNano(), 10), createdAt) + + require.NoError(t, node.RemoveWorker(ctx, worker)) + err = node.registerWorker(ctx, worker.ID, worker.CreatedAt, worker.stream) + require.ErrorContains(t, err, "STREAMDESTROYED") + registered, err := rdb.HExists(ctx, rmapContentKey(node.resources.workers), worker.ID).Result() + require.NoError(t, err) + assert.False(t, registered) +} + +func TestWorkerRegistrationRejectsCleanupInProgress(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, node.Shutdown(context.Background())) + }() + worker := newTestWorker(t, ctx, node) + lease, err := node.acquireGracefulRequeue(ctx, worker.ID) + require.NoError(t, err) + require.NotNil(t, lease) + err = node.registerWorker(ctx, worker.ID, worker.CreatedAt, worker.stream) + assert.ErrorContains(t, err, "WORKERCLEANUPLOST") + assert.NoError(t, node.releaseWorkerCleanup(ctx, lease, false)) +} + +func TestWorkerRegistrationFailedReplyRemovesBothEntries(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, node.Shutdown(context.Background())) + }() + replyError := errors.New("registration reply unavailable") + hook := &workerRegistrationHook{ + workersKey: rmapContentKey(node.resources.workers), + rdb: rdb, + registered: make(chan string, 1), + resume: make(chan struct{}), + replyError: replyError, + } + close(hook.resume) + rdb.AddHook(hook) + + worker, err := node.AddWorker( + context.WithValue(ctx, workerRegistrationContextKey{}, true), + newMockJobHandler(), + ) + require.ErrorIs(t, err, replyError) + assert.Nil(t, worker) + workerID := <-hook.registered + for _, name := range []string{node.resources.workers, node.resources.workerKeepAlive} { + saved, err := rdb.HExists(ctx, rmapContentKey(name), workerID).Result() + require.NoError(t, err) + assert.False(t, saved, "failed creation must remove %s", name) + } + state, err := rdb.HGet(ctx, "pulse:stream:"+workerStreamName(workerID)+":lifecycle", "state").Result() + require.NoError(t, err) + assert.Equal(t, "destroyed", state) +} + +func TestWorkerRegistrationRollbackPreservesForeignCleanupClaim(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + other := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, other.Close(context.Background())) + assert.NoError(t, node.Shutdown(context.Background())) + }() + replyError := errors.New("registration reply unavailable") + hook := &workerRegistrationHook{ + workersKey: rmapContentKey(node.resources.workers), + rdb: rdb, + registered: make(chan string, 1), + resume: make(chan struct{}), + replyError: replyError, + } + rdb.AddHook(hook) + result := make(chan workerRegistrationResult, 1) + go func() { + worker, err := node.AddWorker( + context.WithValue(ctx, workerRegistrationContextKey{}, true), + newMockJobHandler(), + ) + result <- workerRegistrationResult{worker: worker, err: err} + }() + var workerID string + select { + case workerID = <-hook.registered: + case <-ctx.Done(): + close(hook.resume) + <-result + require.NoError(t, ctx.Err(), "worker registration was not observed") + } + now, err := rdb.Time(ctx).Result() + assert.NoError(t, err) + assert.NoError(t, rdb.HSet( + ctx, + rmapContentKey(node.resources.workerKeepAlive), + workerID, + strconv.FormatInt(now.Add(-2*node.workerTTL).UnixNano(), 10), + ).Err()) + lease, err := other.acquireWorkerCleanup(ctx, workerID) + assert.NoError(t, err) + assert.NotNil(t, lease) + close(hook.resume) + got := <-result + assert.ErrorIs(t, got.err, replyError) + assert.ErrorContains(t, got.err, "WORKERCLEANUPLOST") + assert.Nil(t, got.worker) + for _, name := range []string{node.resources.workers, node.resources.workerKeepAlive} { + saved, err := rdb.HExists(ctx, rmapContentKey(name), workerID).Result() + assert.NoError(t, err) + assert.True(t, saved, "another cleanup owner still needs %s", name) + } + if lease != nil { + assert.NoError(t, other.releaseWorkerCleanup(ctx, lease, false)) + } +} + +func TestWorkerRegistrationDelayedCommandCannotReviveWorker(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + node := newTestNode(t, ctx, rdb, t.Name()) + defer func() { + assert.NoError(t, node.Shutdown(context.Background())) + }() + hook := &workerRegistrationHook{ + workersKey: rmapContentKey(node.resources.workers), + rdb: rdb, + registered: make(chan string, 1), + resume: make(chan struct{}), + replaying: make(chan struct{}, 1), + replay: make(chan struct{}), + } + close(hook.resume) + rdb.AddHook(hook) + worker, err := node.AddWorker( + context.WithValue(ctx, workerRegistrationContextKey{}, true), + newMockJobHandler(), + ) + require.NoError(t, err) + require.Equal(t, worker.ID, <-hook.registered) + replayed := make(chan error, 1) + go func() { + replayed <- rdb.Do( + context.WithValue(ctx, workerRegistrationReplayContextKey{}, true), + hook.args..., + ).Err() + }() + select { + case <-hook.replaying: + case <-ctx.Done(): + close(hook.replay) + <-replayed + require.NoError(t, ctx.Err(), "registration replay was not observed") + } + removeErr := node.RemoveWorker(ctx, worker) + close(hook.replay) + replayErr := <-replayed + require.NoError(t, removeErr) + assert.ErrorContains(t, replayErr, "STREAMDESTROYED") + for _, name := range []string{node.resources.workers, node.resources.workerKeepAlive} { + saved, err := rdb.HExists(ctx, rmapContentKey(name), worker.ID).Result() + require.NoError(t, err) + assert.False(t, saved) + } +} + +func TestWorkerRegistrationRepeatedCommandPreservesConstructorHeartbeat(t *testing.T) { + rdb := ptesting.NewRedisClient(t) + defer ptesting.CleanupRedis(t, rdb, false, "") + ctx, cancel := context.WithTimeout(ptesting.NewTestContext(t), 20*time.Second) + defer cancel() + // Keep worker expiry longer than this test's deadline so an overwritten + // startup timestamp cannot turn the constructor wait into a cleanup test. + node, err := AddNode(ctx, t.Name(), rdb, + WithLogger(pulse.NoopLogger()), + WithWorkerTTL(time.Minute), + WithJobSinkBlockDuration(testJobSinkBlockDuration), + ) + require.NoError(t, err) + defer func() { + assert.NoError(t, node.Shutdown(context.Background())) + }() + hook := &workerRegistrationHook{ + workersKey: rmapContentKey(node.resources.workers), + rdb: rdb, + registered: make(chan string, 1), + resume: make(chan struct{}), + } + rdb.AddHook(hook) + result := make(chan workerRegistrationResult, 1) + go func() { + worker, err := node.AddWorker( + context.WithValue(ctx, workerRegistrationContextKey{}, true), + newMockJobHandler(), + ) + result <- workerRegistrationResult{worker: worker, err: err} + }() + var workerID string + select { + case workerID = <-hook.registered: + case <-ctx.Done(): + close(hook.resume) + <-result + require.NoError(t, ctx.Err(), "worker registration was not observed") + } + heartbeatKey := rmapContentKey(node.resources.workerKeepAlive) + before, err := rdb.HGetAll(ctx, heartbeatKey).Result() + assert.NoError(t, err) + + // Hold the successful constructor response while an identical older + // command runs. The constructor must still observe its returned timestamp. + heartbeat, err := rdb.Do(ctx, hook.args...).Text() + assert.NoError(t, err) + after, err := rdb.HGetAll(ctx, heartbeatKey).Result() + assert.NoError(t, err) + assert.Equal(t, before, after, "a repeated registration must not replace the heartbeat") + assert.Equal(t, before[workerID], heartbeat) + assert.NoError(t, waitPoolMapValue(ctx, node.workerKeepAliveMap, workerID, heartbeat)) + close(hook.resume) + got := <-result + assert.NoError(t, got.err) + assert.NotNil(t, got.worker) +} + +func (h *workerRegistrationHook) DialHook(next redis.DialHook) redis.DialHook { + return next +} + +// ProcessHook pauses the marked AddWorker call after Redis saves its worker +// entry. Other node operations continue, exposing the actual cleanup race. +// A marked replay waits until the test retires the worker, then reaches Redis +// so the stored stream state decides whether the old command can run. +func (h *workerRegistrationHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook { + return func(ctx context.Context, cmd redis.Cmder) error { + if ctx.Value(workerRegistrationReplayContextKey{}) == true { + h.replaying <- struct{}{} + select { + case <-h.replay: + case <-ctx.Done(): + return ctx.Err() + } + } + err := next(ctx, cmd) + if err != nil || ctx.Value(workerRegistrationContextKey{}) != true || + (cmd.Name() != "evalsha" && cmd.Name() != "eval") { + return err + } + for _, arg := range cmd.Args() { + if arg != h.workersKey { + continue + } + h.once.Do(func() { + h.args = append([]any(nil), cmd.Args()...) + entries, readErr := h.rdb.HKeys(ctx, h.workersKey).Result() + if readErr != nil { + err = readErr + return + } + for _, entry := range entries { + if entry[0] != '=' { + h.registered <- entry + select { + case <-h.resume: + case <-ctx.Done(): + err = ctx.Err() + } + if err == nil { + err = h.replyError + } + return + } + } + }) + break + } + return err + } +} + +func (h *workerRegistrationHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook { + return next +}