From a30645f9fe193b64b636bd397923e198c012e7a9 Mon Sep 17 00:00:00 2001 From: shrutiyam-glitch Date: Tue, 15 Sep 2026 09:56:50 -0700 Subject: [PATCH] ateapi: apply a replica's own worker writes to its cache at commit (#1583) Fixes #1370 **Summary** `atepg` keeps a registry of the live `WatchWorkers` channels and publishes each committed worker event onto them at commit, about one poll interval ahead of the outbox row carrying the same event. Replicas see their own writes sooner without a second notification API. **Why** Replicas previously learned of their own writes only from the outbox watch, ~50ms later. In that window the scheduler re-picks a worker it just assigned and retries on the precondition check, and capacity a replica just freed stays invisible to its own placement. Cross-replica propagation is unchanged. **Key changes** - `WatchWorkers` enrolls its channel; the poller deregisters before closing it, which is what makes a send on a closed channel impossible. - `writeAndAppendEvent` publishes only after `tx.Commit` returns. - The commit-time copy is decoded from the same payload the outbox stores, so both deliveries are identical, including unknown-field pruning. - Sends are non-blocking and never advance the poll cursor, so a full buffer costs a watcher nothing and gap-free delivery is unchanged. - Consumers now also see their own writes out of xid order and again on the poll. `workercache` already fences creates and updates on version, so duplicates and stale events are no-ops. **Testing** `TestLocalPublishReachesWatchers`, `TestLocalPublishSurvivesWatchClose`, and `TestCache_WatchEventsAreFenced` are new. `TestWatchWorkers_DeliveryFencedByOldestTransaction` now consumes the commit-time copy first and asserts the xmin fence on the outbox copy; the fence governs the poll, and a local publish bypasses it by construction. Full `atepg` and `workercache` suites pass under `-race`. - [X] Tests pass - [X] Appropriate changes to documentation are included in the PR --- cmd/ateapi/internal/store/atepg/atepg.go | 51 ++++++- cmd/ateapi/internal/store/atepg/outbox.go | 18 ++- .../internal/store/atepg/outbox_test.go | 143 +++++++++++++++++- .../internal/workercache/workercache_test.go | 38 +++++ 4 files changed, 243 insertions(+), 7 deletions(-) diff --git a/cmd/ateapi/internal/store/atepg/atepg.go b/cmd/ateapi/internal/store/atepg/atepg.go index a30ccef29..c383a6bef 100644 --- a/cmd/ateapi/internal/store/atepg/atepg.go +++ b/cmd/ateapi/internal/store/atepg/atepg.go @@ -26,6 +26,7 @@ import ( "errors" "fmt" "log/slog" + "sync" "time" "github.com/agent-substrate/substrate/cmd/ateapi/internal/store" @@ -58,6 +59,46 @@ type Persistence struct { pollFailureCloseAfter time.Duration stopMaintenance context.CancelFunc maintenanceDone chan struct{} + // watchMu guards watchers, the live WatchWorkers channels. + watchMu sync.Mutex + watchers map[chan store.WorkerEvent]struct{} +} + +// addWatcher enrolls a WatchWorkers channel to receive locally published events. +func (p *Persistence) addWatcher(ch chan store.WorkerEvent) { + p.watchMu.Lock() + defer p.watchMu.Unlock() + p.watchers[ch] = struct{}{} +} + +// removeWatcher unenrolls a channel. The caller must call it before closing +// the channel: once it returns, publishLocally can no longer send on it. +func (p *Persistence) removeWatcher(ch chan store.WorkerEvent) { + p.watchMu.Lock() + defer p.watchMu.Unlock() + delete(p.watchers, ch) +} + +// publishLocally hands a committed event to this process's watchers a poll +// interval ahead of the outbox, one copy each. Sends are non-blocking: a +// watcher with a full buffer is skipped and gets the event from the outbox. +func (p *Persistence) publishLocally(ctx context.Context, payload []byte) { + p.watchMu.Lock() + defer p.watchMu.Unlock() + if len(p.watchers) == 0 { + return + } + event, err := unmarshalWorkerEvent(payload) + if err != nil { + slog.ErrorContext(ctx, "decoding locally published worker event failed", slog.Any("err", err)) + return + } + for ch := range p.watchers { + select { + case ch <- store.WorkerEvent{Type: event.Type, Worker: proto.Clone(event.Worker).(*ateapipb.Worker)}: + default: + } + } } var _ store.Interface = (*Persistence)(nil) @@ -175,7 +216,15 @@ func newPersistence(ctx context.Context, pool, watchPool *pgxpool.Pool) (*Persis return nil, err } maintenanceCtx, stopMaintenance := context.WithCancel(context.Background()) - p := &Persistence{pool: pool, watchPool: watchPool, leaseTTL: defaultLeaseTTL, pollFailureCloseAfter: outboxPollFailureCloseAfter, stopMaintenance: stopMaintenance, maintenanceDone: make(chan struct{})} + p := &Persistence{ + pool: pool, + watchPool: watchPool, + leaseTTL: defaultLeaseTTL, + pollFailureCloseAfter: outboxPollFailureCloseAfter, + stopMaintenance: stopMaintenance, + maintenanceDone: make(chan struct{}), + watchers: make(map[chan store.WorkerEvent]struct{}), + } // Cover the partition lead before accepting writes; from then on the // maintenance loop keeps partitions ahead of the clock (and the // DEFAULT partition catches writes if it ever falls behind). diff --git a/cmd/ateapi/internal/store/atepg/outbox.go b/cmd/ateapi/internal/store/atepg/outbox.go index 8f43f4591..c20e8646a 100644 --- a/cmd/ateapi/internal/store/atepg/outbox.go +++ b/cmd/ateapi/internal/store/atepg/outbox.go @@ -88,8 +88,9 @@ func (p *Persistence) writeAndAppendEvent(ctx context.Context, eventType store.W return nil, err } + var payload []byte if worker != nil { - payload, err := marshalWorkerEvent(eventType, worker) + payload, err = marshalWorkerEvent(eventType, worker) if err != nil { return nil, fmt.Errorf("marshaling worker event: %w", err) } @@ -101,6 +102,9 @@ func (p *Persistence) writeAndAppendEvent(ctx context.Context, eventType store.W if err := tx.Commit(ctx); err != nil { return nil, fmt.Errorf("committing transaction: %w", err) } + if payload != nil { + p.publishLocally(ctx, payload) + } return worker, nil } @@ -439,7 +443,9 @@ func (p *Persistence) truncateWorkerOutboxDefault(ctx context.Context, q querier // to guarantee gap-free delivery. Note that a long-running transaction anywhere // in the database will stall delivery. // -// Events are delivered in xid order, so consumers must reconcile worker versions. +// This process's own writes are published at commit, out of xid order, and again +// on the poll, so consumers must reconcile versions and tolerate duplicates. +// // If the watcher detects missed events—either by lagging behind retention drops // or if a database restart truncates the UNLOGGED partitions—it closes the channel // to force the consumer to resync from the primary tables. @@ -462,8 +468,14 @@ func (p *Persistence) WatchWorkers(ctx context.Context) (*store.WorkerWatch, err } ch := make(chan store.WorkerEvent, 128) + // Committed writes in this process are published straight onto ch, ahead + // of the poll that would carry them. + p.addWatcher(ch) go func() { - defer close(ch) + defer func() { + p.removeWatcher(ch) + close(ch) + }() ticker := time.NewTicker(outboxPollInterval) defer ticker.Stop() // failingSince limits how long consumers serve stale state during an outage. diff --git a/cmd/ateapi/internal/store/atepg/outbox_test.go b/cmd/ateapi/internal/store/atepg/outbox_test.go index 24a84d889..22a9f9f24 100644 --- a/cmd/ateapi/internal/store/atepg/outbox_test.go +++ b/cmd/ateapi/internal/store/atepg/outbox_test.go @@ -385,8 +385,8 @@ func TestWorkerEvents_OneRowPerTransaction(t *testing.T) { } // TestWatchWorkers_DeliveryFencedByOldestTransaction documents the xmin -// fence's real bound: one old transaction anywhere holds back delivery of -// everything committed after it, for as long as it lives. +// fence's real bound: one old transaction anywhere holds back outbox delivery +// of everything committed after it, for as long as it lives. func TestWatchWorkers_DeliveryFencedByOldestTransaction(t *testing.T) { s := setupPostgresPersistence(t) ctx := context.Background() @@ -403,6 +403,7 @@ func TestWatchWorkers_DeliveryFencedByOldestTransaction(t *testing.T) { t.Fatalf("Begin blocker failed: %v", err) } defer blocker.Rollback(ctx) //nolint:errcheck // released below + if _, err := blocker.Exec(ctx, `SELECT pg_current_xact_id()`); err != nil { t.Fatalf("assigning blocker xid failed: %v", err) } @@ -417,9 +418,17 @@ func TestWatchWorkers_DeliveryFencedByOldestTransaction(t *testing.T) { t.Fatalf("CreateWorker failed: %v", err) } + // The commit-time copy is published in-process and never sees the fence. + // The fence governs the outbox copy that follows it. + select { + case <-watch.Events: + case <-time.After(2 * time.Second): + t.Fatal("commit-time copy not delivered") + } + select { case event := <-watch.Events: - t.Fatalf("event %+v delivered through the fence while an older transaction was in flight", event) + t.Fatalf("outbox copy of %+v delivered through the fence while an older transaction was in flight", event) case <-time.After(600 * time.Millisecond): // Expected: committed but fenced behind the blocker's xid. } @@ -1017,3 +1026,131 @@ func TestWatchWorkers_ClosesOnCorruptPayload(t *testing.T) { t.Fatal("watch stayed open past a corrupt payload (silent skip)") } } + +// TestLocalPublishReachesWatchers pins the local fast-path contract: every +// worker event — create, update, and delete — is on an active watcher's +// channel by the time the write call returns, carrying the committed state. +// Each read is non-blocking, so it can only be satisfied by the commit-time +// publish; the poller needs at least one outboxPollInterval tick. +func TestLocalPublishReachesWatchers(t *testing.T) { + requirePool(t) + ctx := context.Background() + + p, err := Connect(ctx, containerDSN, "public") + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer p.pool.Close() + defer p.Close() + clearAll(t, p) + + watch, err := p.WatchWorkers(ctx) + if err != nil { + t.Fatalf("WatchWorkers failed: %v", err) + } + defer watch.Close() + + // nextEvent takes the next buffered event of type want, skipping outbox + // copies of earlier writes. It never blocks, so only the commit-time + // publish can satisfy it: the poller needs at least one tick to run. + nextEvent := func(what string, want store.WorkerEventType) store.WorkerEvent { + t.Helper() + for { + select { + case ev, ok := <-watch.Events: + if !ok { + t.Fatalf("after %s: watch closed", what) + } + if ev.Type == want { + return ev + } + default: + t.Fatalf("after %s: no %v on the watch channel; the commit-time publish did not reach the watcher", what, want) + return store.WorkerEvent{} + } + } + } + + created, err := p.CreateWorker(ctx, &ateapipb.Worker{ + Metadata: &ateapipb.ResourceMetadata{Name: "local-publish-worker"}, + WorkerNamespace: "ns", + WorkerPool: "pool", + WorkerPod: "pod", + }) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + ev := nextEvent("create", store.WorkerEventCreated) + if got, want := ev.Worker.GetMetadata().GetVersion(), created.GetMetadata().GetVersion(); got != want { + t.Errorf("created event version = %d, want committed version %d", got, want) + } + + updated, err := p.UpdateWorker(ctx, created.GetMetadata().GetName(), store.PreconditionFrom(created), func(toUpdate *ateapipb.Worker) error { + toUpdate.Ip = "10.0.0.9" + return nil + }) + if err != nil { + t.Fatalf("UpdateWorker failed: %v", err) + } + ev = nextEvent("update", store.WorkerEventUpdated) + if got, want := ev.Worker.GetMetadata().GetVersion(), updated.GetMetadata().GetVersion(); got != want { + t.Errorf("updated event version = %d, want committed version %d", got, want) + } + if ev.Worker.GetIp() != "10.0.0.9" { + t.Errorf("updated event carries Ip %q, want the committed mutation", ev.Worker.GetIp()) + } + // The watcher's copy must be isolated from the caller's returned Worker. + if ev.Worker == updated { + t.Error("locally published Worker aliases the caller's returned Worker; it must be a copy") + } + + if _, err := p.DeleteWorker(ctx, created.GetMetadata().GetName(), store.DeletePreconditions{}); err != nil { + t.Fatalf("DeleteWorker failed: %v", err) + } + ev = nextEvent("delete", store.WorkerEventDeleted) + if ev.Worker.GetMetadata().GetName() != created.GetMetadata().GetName() { + t.Errorf("deleted event names worker %q, want %q", ev.Worker.GetMetadata().GetName(), created.GetMetadata().GetName()) + } +} + +// TestLocalPublishSurvivesWatchClose covers the close race the watcher +// registry exists to prevent: a write publishing concurrently with a watch +// shutting down must not send on a closed channel. +func TestLocalPublishSurvivesWatchClose(t *testing.T) { + requirePool(t) + ctx := context.Background() + + p, err := Connect(ctx, containerDSN, "public") + if err != nil { + t.Fatalf("Connect failed: %v", err) + } + defer p.pool.Close() + defer p.Close() + clearAll(t, p) + + for i := range 20 { + watch, err := p.WatchWorkers(ctx) + if err != nil { + t.Fatalf("WatchWorkers failed: %v", err) + } + done := make(chan struct{}) + go func() { + defer close(done) + watch.Close() + }() + name := fmt.Sprintf("close-race-worker-%d", i) + w, err := p.CreateWorker(ctx, &ateapipb.Worker{ + Metadata: &ateapipb.ResourceMetadata{Name: name}, + WorkerNamespace: "ns", + WorkerPool: "pool", + WorkerPod: "pod", + }) + if err != nil { + t.Fatalf("CreateWorker failed: %v", err) + } + if _, err := p.DeleteWorker(ctx, w.GetMetadata().GetName(), store.DeletePreconditions{}); err != nil { + t.Fatalf("DeleteWorker failed: %v", err) + } + <-done + } +} diff --git a/cmd/ateapi/internal/workercache/workercache_test.go b/cmd/ateapi/internal/workercache/workercache_test.go index 1a36273a0..24e325c68 100644 --- a/cmd/ateapi/internal/workercache/workercache_test.go +++ b/cmd/ateapi/internal/workercache/workercache_test.go @@ -470,3 +470,41 @@ func eventually(t *testing.T, condition func() bool, timeout time.Duration) { t.Fatal("condition not met within timeout") } } + +// The watch may replay an event or deliver one out of order, so an event +// older than what the cache already holds must not regress it. +func TestCache_WatchEventsAreFenced(t *testing.T) { + w := makeWorker("ns", "pod1", 1) + fs := newFakeStore(w) + c := workercache.New(fs, time.Hour) + if err := c.Start(t.Context()); err != nil { + t.Fatalf("Start: %v", err) + } + + updated := makeWorker("ns", "pod1", 2) + updated.Status.Allocated = &ateapipb.WorkerResources{Actors: 1} + fs.send(store.WorkerEvent{Type: store.WorkerEventUpdated, Worker: updated}) + eventually(t, func() bool { + w, err := c.Worker(workerName("ns", "pod1")) + return err == nil && w.GetMetadata().GetVersion() == 2 && w.GetStatus().GetAllocated().GetActors() == 1 + }, 2*time.Second) + + // A replayed and an older event must both be no-ops. + fs.send(store.WorkerEvent{Type: store.WorkerEventUpdated, Worker: updated}) + fs.send(store.WorkerEvent{Type: store.WorkerEventUpdated, Worker: makeWorker("ns", "pod1", 1)}) + // Ordering on the channel means observing a later event proves the two + // above were already applied. + fs.send(store.WorkerEvent{Type: store.WorkerEventCreated, Worker: makeWorker("ns", "pod2", 1)}) + eventually(t, func() bool { + _, err := c.Worker(workerName("ns", "pod2")) + return err == nil + }, 2*time.Second) + + got, err := c.Worker(workerName("ns", "pod1")) + if err != nil { + t.Fatalf("Worker: %v", err) + } + if got.GetMetadata().GetVersion() != 2 || got.GetStatus().GetAllocated().GetActors() != 1 { + t.Fatalf("stale event regressed the cache to version %d", got.GetMetadata().GetVersion()) + } +}