From 0cb7738bf5684f4e166ee22b534828620d23a5ae Mon Sep 17 00:00:00 2001 From: NekoPunch Date: Fri, 14 Aug 2026 02:32:46 -0700 Subject: [PATCH] fix(ateapi): drop worker watch events instead of stalling pub/sub A blocking send on the full watch buffer stops the pub/sub socket being drained, so Valkey evicts the subscriber for overrunning its output buffer and every event in flight is lost silently. Overflow now drops events, counts them, and signals an invalidation so the worker cache relists instead of serving a stale snapshot until the periodic relist. Claude-Session: https://claude.ai/code/session_01XuQqkwLf5Zx6CSZFHSC6hb --- .../internal/controlapi/functional_test.go | 2 +- .../controlapi/workflow_suspend_test.go | 2 +- cmd/ateapi/internal/store/atepg/atepg.go | 3 +- .../internal/store/ateredis/ateredis.go | 43 ++++- .../internal/store/ateredis/ateredis_test.go | 95 +++++++++- cmd/ateapi/internal/store/store.go | 11 +- .../internal/store/storetest/storetest.go | 2 +- .../internal/workercache/workercache.go | 38 ++++ .../internal/workercache/workercache_test.go | 165 +++++++++++++++++- cmd/ateapi/main.go | 2 +- 10 files changed, 337 insertions(+), 26 deletions(-) diff --git a/cmd/ateapi/internal/controlapi/functional_test.go b/cmd/ateapi/internal/controlapi/functional_test.go index 0087b0da1..80bcc33af 100644 --- a/cmd/ateapi/internal/controlapi/functional_test.go +++ b/cmd/ateapi/internal/controlapi/functional_test.go @@ -295,7 +295,7 @@ func setupTest(t *testing.T, ns string) *testContext { rdb := redis.NewClusterClient(&redis.ClusterOptions{ Addrs: []string{mr.Addr()}, }) - persistence := ateredis.NewPersistence(rdb) + persistence := ateredis.NewPersistence(rdb, nil) // 2. Initialize Clientsets using global cfg k8sClient, err := kubernetes.NewForConfig(cfg) diff --git a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go index 38c42c318..daa708960 100644 --- a/cmd/ateapi/internal/controlapi/workflow_suspend_test.go +++ b/cmd/ateapi/internal/controlapi/workflow_suspend_test.go @@ -226,7 +226,7 @@ func newTestPersistence(t *testing.T) store.Interface { t.Cleanup(mr.Close) rdb := redis.NewClusterClient(&redis.ClusterOptions{Addrs: []string{mr.Addr()}}) t.Cleanup(func() { rdb.Close() }) //nolint:errcheck // test cleanup - return ateredis.NewPersistence(rdb) + return ateredis.NewPersistence(rdb, nil) } // newDanglingDialer returns a dialer whose informer cache has no pods, so diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index ead0b22f8..d17544bf4 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -1506,7 +1506,8 @@ func (p *Persistence) WatchWorkers(ctx context.Context) (*store.WorkerWatch, err } } }() - return store.NewWorkerWatch(ch, cancel), nil + // nil invalidation: this watch blocks on a full buffer rather than dropping. + return store.NewWorkerWatch(ch, nil, cancel), nil } // --- Workflow locks --- diff --git a/cmd/ateapi/internal/store/ateredis/ateredis.go b/cmd/ateapi/internal/store/ateredis/ateredis.go index a27962754..3309b776a 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis.go +++ b/cmd/ateapi/internal/store/ateredis/ateredis.go @@ -57,6 +57,7 @@ import ( "github.com/agent-substrate/substrate/pkg/proto/ateapipb" "github.com/google/uuid" "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel/metric" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/types/known/timestamppb" @@ -79,18 +80,34 @@ type redisClient interface { // Persistence is a service that stores information about applications in Redis. type Persistence struct { - rdb redisClient - lockTTL time.Duration + rdb redisClient + lockTTL time.Duration + droppedEvents metric.Int64Counter } var _ store.Interface = (*Persistence)(nil) -// NewPersistence creates a new Persistence. -func NewPersistence(redisClient *redis.ClusterClient) *Persistence { - return &Persistence{ +const droppedEventsMetric = "ate.store.worker_watch.dropped_events" + +// NewPersistence creates a new Persistence. meter may be nil to disable metrics. +func NewPersistence(redisClient *redis.ClusterClient, meter metric.Meter) *Persistence { + p := &Persistence{ rdb: redisClient, lockTTL: defaultLockTTL, } + if meter != nil { + counter, err := meter.Int64Counter( + droppedEventsMetric, + metric.WithUnit("{event}"), + metric.WithDescription("Worker watch events dropped because the consumer buffer was full."), + ) + if err != nil { + slog.Error("Failed to register worker watch dropped-events counter", "metric", droppedEventsMetric, "error", err) + } else { + p.droppedEvents = counter + } + } + return p } // actorDBKey returns the Redis key an actor is stored under. The encoding is @@ -659,6 +676,7 @@ func (s *Persistence) WatchWorkers(ctx context.Context) (*store.WorkerWatch, err return nil, fmt.Errorf("while confirming worker subscription: %w", err) } ch := make(chan store.WorkerEvent, 128) + invalidated := make(chan struct{}, 1) go func() { defer close(ch) defer pubsub.Close() @@ -678,13 +696,22 @@ func (s *Persistence) WatchWorkers(ctx context.Context) (*store.WorkerWatch, err } select { case ch <- event: - case <-watchCtx.Done(): - return + default: + // Blocking here would stall the pub/sub socket and drop the + // connection; drop the event and coalesce an invalidation. + if s.droppedEvents != nil { + s.droppedEvents.Add(ctx, 1) + } + select { + case invalidated <- struct{}{}: + slog.WarnContext(ctx, "worker watch buffer full; dropping events, invalidation signaled") + default: + } } } } }() - return store.NewWorkerWatch(ch, cancel), nil + return store.NewWorkerWatch(ch, invalidated, cancel), nil } // DebugClearAll flushes all data from Redis. diff --git a/cmd/ateapi/internal/store/ateredis/ateredis_test.go b/cmd/ateapi/internal/store/ateredis/ateredis_test.go index 23b9a2536..350976a5c 100644 --- a/cmd/ateapi/internal/store/ateredis/ateredis_test.go +++ b/cmd/ateapi/internal/store/ateredis/ateredis_test.go @@ -28,6 +28,9 @@ import ( "github.com/alicebob/miniredis/v2" "github.com/google/go-cmp/cmp" "github.com/redis/go-redis/v9" + "go.opentelemetry.io/otel/metric" + sdkmetric "go.opentelemetry.io/otel/sdk/metric" + "go.opentelemetry.io/otel/sdk/metric/metricdata" "google.golang.org/protobuf/encoding/protojson" "google.golang.org/protobuf/proto" "google.golang.org/protobuf/testing/protocmp" @@ -38,6 +41,10 @@ import ( ) func setupTest(t *testing.T) (*miniredis.Miniredis, *Persistence, context.Context) { + return setupTestWithMeter(t, nil) +} + +func setupTestWithMeter(t *testing.T, meter metric.Meter) (*miniredis.Miniredis, *Persistence, context.Context) { mr, err := miniredis.Run() if err != nil { t.Fatalf("failed to start miniredis: %v", err) @@ -50,7 +57,7 @@ func setupTest(t *testing.T) (*miniredis.Miniredis, *Persistence, context.Contex Addrs: []string{mr.Addr()}, }) t.Cleanup(func() { rdb.Close() }) - return mr, NewPersistence(rdb), t.Context() + return mr, NewPersistence(rdb, meter), t.Context() } // testAtespace is the atespace used by tests that create a single actor. Actors @@ -3236,3 +3243,89 @@ func TestDeleteActorTemplate_VersionInOtherAtespace_NotBlocking(t *testing.T) { t.Errorf("DeleteActorTemplate = %v, want nil: the only version lives in team-b", err) } } + +func droppedEventsCount(t *testing.T, reader *sdkmetric.ManualReader) int64 { + t.Helper() + var rm metricdata.ResourceMetrics + if err := reader.Collect(context.Background(), &rm); err != nil { + t.Fatalf("collecting metrics: %v", err) + } + for _, sm := range rm.ScopeMetrics { + for _, m := range sm.Metrics { + if m.Name != droppedEventsMetric { + continue + } + var total int64 + for _, dp := range m.Data.(metricdata.Sum[int64]).DataPoints { + total += dp.Value + } + return total + } + } + return 0 +} + +func TestWatchWorkersOverflowSignalsInvalidation(t *testing.T) { + reader := sdkmetric.NewManualReader() + meter := sdkmetric.NewMeterProvider(sdkmetric.WithReader(reader)).Meter("test") + _, s, ctx := setupTestWithMeter(t, meter) + + watch, err := s.WatchWorkers(ctx) + if err != nil { + t.Fatalf("WatchWorkers failed: %v", err) + } + defer watch.Close() + + for i := 0; i < 200; i++ { + if err := s.CreateWorker(ctx, &ateapipb.Worker{ + WorkerNamespace: "default", + WorkerPool: "pool-1", + WorkerPod: fmt.Sprintf("pod-%d", i), + }); err != nil { + t.Fatalf("CreateWorker %d failed: %v", i, err) + } + } + + select { + case <-watch.Invalidated: + case <-time.After(10 * time.Second): + t.Fatal("no invalidation signal after overflowing the watch buffer") + } + + if got := droppedEventsCount(t, reader); got <= 0 { + t.Errorf("%s = %d, want > 0 after overflow", droppedEventsMetric, got) + } + + for draining := true; draining; { + select { + case _, ok := <-watch.Events: + if !ok { + t.Fatal("watch events channel closed after overflow; the subscription must survive") + } + case <-time.After(200 * time.Millisecond): + draining = false + } + } + + if err := s.CreateWorker(ctx, &ateapipb.Worker{ + WorkerNamespace: "default", + WorkerPool: "pool-1", + WorkerPod: "pod-after-overflow", + }); err != nil { + t.Fatalf("CreateWorker after overflow failed: %v", err) + } + deadline := time.After(10 * time.Second) + for { + select { + case event, ok := <-watch.Events: + if !ok { + t.Fatal("watch events channel closed after overflow") + } + if event.Worker.GetWorkerPod() == "pod-after-overflow" { + return + } + case <-deadline: + t.Fatal("event published after overflow never delivered") + } + } +} diff --git a/cmd/ateapi/internal/store/store.go b/cmd/ateapi/internal/store/store.go index 54157bc14..f376dab3a 100644 --- a/cmd/ateapi/internal/store/store.go +++ b/cmd/ateapi/internal/store/store.go @@ -280,15 +280,18 @@ type WorkerEvent struct { type WorkerWatch struct { // Events delivers worker state changes until the watch is torn down. Events <-chan WorkerEvent + // Invalidated signals that events were dropped and the consumer should + // relist. Signals coalesce; nil for stores that never drop events. + Invalidated <-chan struct{} // stop releases the subscription backing Events. It is a context.CancelFunc, // so it is safe to call multiple times. stop context.CancelFunc } -// NewWorkerWatch builds a WorkerWatch from an events channel and the cancel -// func that tears down its subscription. -func NewWorkerWatch(events <-chan WorkerEvent, stop context.CancelFunc) *WorkerWatch { - return &WorkerWatch{Events: events, stop: stop} +// NewWorkerWatch builds a WorkerWatch from an events channel, an optional +// invalidation channel, and the cancel func that tears down its subscription. +func NewWorkerWatch(events <-chan WorkerEvent, invalidated <-chan struct{}, stop context.CancelFunc) *WorkerWatch { + return &WorkerWatch{Events: events, Invalidated: invalidated, stop: stop} } // Close releases the subscription. Safe to call multiple times. diff --git a/cmd/ateapi/internal/store/storetest/storetest.go b/cmd/ateapi/internal/store/storetest/storetest.go index e92f3e5f5..2a4297e8b 100644 --- a/cmd/ateapi/internal/store/storetest/storetest.go +++ b/cmd/ateapi/internal/store/storetest/storetest.go @@ -37,7 +37,7 @@ func SetupTestStore(t *testing.T) (store.Interface, func()) { Addrs: []string{mr.Addr()}, }) - persistence := ateredis.NewPersistence(rdb) + persistence := ateredis.NewPersistence(rdb, nil) cleanup := func() { rdb.Close() diff --git a/cmd/ateapi/internal/workercache/workercache.go b/cmd/ateapi/internal/workercache/workercache.go index 25846e4f8..0c9bd9ffb 100644 --- a/cmd/ateapi/internal/workercache/workercache.go +++ b/cmd/ateapi/internal/workercache/workercache.go @@ -34,6 +34,11 @@ import ( // relistPageSize is the page size used for the relist. const relistPageSize = 1000 +const ( + minInvalidationRetry = 100 * time.Millisecond + maxInvalidationRetry = 5 * time.Second +) + // Cache maintains an in-memory snapshot of all workers. // // TODO: add metrics — at minimum a gauge for worker count, a counter for @@ -143,6 +148,30 @@ func (c *Cache) relist(ctx context.Context) error { func (c *Cache) watchEvents(ctx context.Context, watch *store.WorkerWatch) { ticker := time.NewTicker(c.relistInterval) defer ticker.Stop() + + // Signal consumed and backlog discarded: a failed relist has nothing to retrigger it, so retry ourselves. + var retry <-chan time.Time + retryDelay := minInvalidationRetry + correct := func() { + // Discard events buffered before the drop so a stale event can't roll + // back the relisted snapshot. + for draining := true; draining; { + select { + case _, ok := <-watch.Events: + draining = ok + default: + draining = false + } + } + if err := c.relist(ctx); err != nil { + slog.WarnContext(ctx, "worker cache: invalidation relist failed, retrying", slog.Any("err", err), slog.Duration("in", retryDelay)) + retry = time.After(retryDelay) + retryDelay = min(retryDelay*2, maxInvalidationRetry) + return + } + retry, retryDelay = nil, minInvalidationRetry + } + for { select { case event, ok := <-watch.Events: @@ -157,13 +186,22 @@ func (c *Cache) watchEvents(ctx context.Context, watch *store.WorkerWatch) { if watch == nil { return // context cancelled } + // resync already relisted; a pending correction would only drain the new watch. + retry, retryDelay = nil, minInvalidationRetry c.ready.Store(true) } else { c.applyEvent(event) } + case <-watch.Invalidated: + slog.WarnContext(ctx, "worker cache: watch events dropped, relisting") + correct() + case <-retry: + correct() case <-ticker.C: if err := c.relist(ctx); err != nil { slog.WarnContext(ctx, "worker cache: periodic relist failed", slog.Any("err", err)) + } else { + retry, retryDelay = nil, minInvalidationRetry } case <-ctx.Done(): c.ready.Store(false) diff --git a/cmd/ateapi/internal/workercache/workercache_test.go b/cmd/ateapi/internal/workercache/workercache_test.go index 09d77d018..b350b38f0 100644 --- a/cmd/ateapi/internal/workercache/workercache_test.go +++ b/cmd/ateapi/internal/workercache/workercache_test.go @@ -383,24 +383,27 @@ func TestCache_Relist_FailureIsNonFatal(t *testing.T) { type fakeStore struct { store.Interface - mu sync.Mutex - workers []*ateapipb.Worker - watchCh chan store.WorkerEvent - listErr error // if set, ListWorkers returns it - closes int // number of times a returned watch was Closed + mu sync.Mutex + workers []*ateapipb.Worker + watchCh chan store.WorkerEvent + invalidateCh chan struct{} + listErr error // if set, ListWorkers returns it + failListOnce bool // if set, the next ListWorkers fails and clears it + closes int // number of times a returned watch was Closed } func newFakeStore(workers ...*ateapipb.Worker) *fakeStore { return &fakeStore{ - workers: workers, - watchCh: make(chan store.WorkerEvent, 16), + workers: workers, + watchCh: make(chan store.WorkerEvent, 16), + invalidateCh: make(chan struct{}, 1), } } func (f *fakeStore) WatchWorkers(_ context.Context) (*store.WorkerWatch, error) { f.mu.Lock() defer f.mu.Unlock() - return store.NewWorkerWatch(f.watchCh, func() { + return store.NewWorkerWatch(f.watchCh, f.invalidateCh, func() { f.mu.Lock() f.closes++ f.mu.Unlock() @@ -410,6 +413,10 @@ func (f *fakeStore) WatchWorkers(_ context.Context) (*store.WorkerWatch, error) func (f *fakeStore) ListWorkers(_ context.Context, _ store.ListOptions) (store.ListResponse[*ateapipb.Worker], error) { f.mu.Lock() defer f.mu.Unlock() + if f.failListOnce { + f.failListOnce = false + return store.ListResponse[*ateapipb.Worker]{}, errors.New("transient list failure") + } if f.listErr != nil { return store.ListResponse[*ateapipb.Worker]{}, f.listErr } @@ -465,3 +472,145 @@ func eventually(t *testing.T, condition func() bool, timeout time.Duration) { t.Fatal("condition not met within timeout") } } + +func TestInvalidationTriggersImmediateRelist(t *testing.T) { + ctx := t.Context() + w1 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w1"} + fs := newFakeStore(w1) + c := workercache.New(fs, time.Hour) + if err := c.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + w2 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w2"} + fs.mu.Lock() + fs.workers = []*ateapipb.Worker{w1, w2} + fs.mu.Unlock() + + fs.invalidateCh <- struct{}{} + + deadline := time.Now().Add(5 * time.Second) + for { + workers, err := c.Workers() + if err != nil { + t.Fatalf("Workers() = %v; the cache must stay ready through an invalidation relist", err) + } + if len(workers) == 2 { + return + } + if time.Now().After(deadline) { + t.Fatalf("cache never picked up the dropped worker via relist; have %d workers", len(workers)) + } + time.Sleep(10 * time.Millisecond) + } +} + +// The invalidation signal is consumed once, so a failed corrective relist must +// be retried on its own schedule — the periodic ticker is an hour away here. +func TestInvalidationRelistRetriesAfterFailure(t *testing.T) { + ctx := t.Context() + w1 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w1"} + fs := newFakeStore(w1) + c := workercache.New(fs, time.Hour) + if err := c.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + w2 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w2"} + fs.mu.Lock() + fs.workers = []*ateapipb.Worker{w1, w2} + fs.failListOnce = true + fs.mu.Unlock() + + fs.invalidateCh <- struct{}{} + + eventually(t, func() bool { + workers, err := c.Workers() + if err != nil { + t.Errorf("Workers() = %v; the cache must keep serving its stale snapshot while the relist retries", err) + return true + } + return len(workers) == 2 + }, 5*time.Second) + + if workers, err := c.Workers(); err != nil || len(workers) != 2 { + t.Fatalf("Workers() = %d workers, %v; want 2 after the retried relist", len(workers), err) + } +} + +// A resync relists against a fresh watch, satisfying any pending correction. +// A retry left armed from before would drain that new watch and relist over +// state the resync just established. +func TestResyncCancelsPendingInvalidationRetry(t *testing.T) { + ctx := t.Context() + w1 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w1"} + fs := newFakeStore(w1) + c := workercache.New(fs, time.Hour) + if err := c.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + fs.mu.Lock() + fs.failListOnce = true + fs.mu.Unlock() + fs.invalidateCh <- struct{}{} + + eventually(t, func() bool { + fs.mu.Lock() + defer fs.mu.Unlock() + return !fs.failListOnce + }, 2*time.Second) + + fs.disconnect() + eventually(t, func() bool { + workers, err := c.Workers() + return err == nil && len(workers) == 1 + }, 2*time.Second) + + // w2 exists only as an event, never in the store snapshot: a relist the + // resync did not ask for would drop it back out of the cache. + w2 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w2"} + fs.send(store.WorkerEvent{Type: store.WorkerEventCreated, Worker: w2}) + eventually(t, func() bool { + workers, err := c.Workers() + return err == nil && len(workers) == 2 + }, 2*time.Second) + + time.Sleep(400 * time.Millisecond) + if workers, err := c.Workers(); err != nil || len(workers) != 2 { + t.Fatalf("Workers() = %d workers, %v; want 2: the stale retry relisted over the resynced watch", len(workers), err) + } +} + +func TestInvalidationDiscardsStaleBacklog(t *testing.T) { + ctx := t.Context() + w1 := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "w1"} + fs := newFakeStore(w1) + c := workercache.New(fs, time.Hour) + if err := c.Start(ctx); err != nil { + t.Fatalf("Start: %v", err) + } + + ghost := &ateapipb.Worker{WorkerNamespace: "ns", WorkerPool: "pool", WorkerPod: "ghost"} + for range 16 { + fs.send(store.WorkerEvent{Type: store.WorkerEventCreated, Worker: ghost}) + } + fs.invalidateCh <- struct{}{} + + deadline := time.Now().Add(5 * time.Second) + for stable := 0; stable < 5; { + workers, err := c.Workers() + if err != nil { + t.Fatalf("Workers() = %v; the cache must stay ready through an invalidation relist", err) + } + if len(workers) == 1 && workers[0].GetWorkerPod() == "w1" { + stable++ + } else { + stable = 0 + if time.Now().After(deadline) { + t.Fatalf("ghost worker still in cache after invalidation relist; have %d workers", len(workers)) + } + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/cmd/ateapi/main.go b/cmd/ateapi/main.go index 5d1d3e0d5..0d2e18d9f 100644 --- a/cmd/ateapi/main.go +++ b/cmd/ateapi/main.go @@ -326,7 +326,7 @@ func connectStore(ctx context.Context) (store.Interface, error) { if err != nil { return nil, fmt.Errorf("setting up Redis/Valkey: %w", err) } - return ateredis.NewPersistence(redisClient), nil + return ateredis.NewPersistence(redisClient, otel.Meter("ateapi")), nil case "postgres": if *postgresConnectionString == "" { return nil, fmt.Errorf("--store-backend=postgres requires --postgres-connection-string")