From f0fa4a7b1fdd67e3bb0fa951a91af21b21820122 Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Mon, 31 Aug 2026 10:31:51 +0000 Subject: [PATCH] test(distributed): wait for the log subscriber instead of racing it The WebSocket log handler writes its "initial" batch before it calls Subscribe, so a line appended the instant that batch arrives lands in the circular buffer with no subscriber to receive it. Three backend-logs specs append exactly there and then wait out a 5s read deadline; once a gorilla read hits its deadline the connection is unusable, so the spec cannot retry. `--focus='Worker WebSocket log streaming' --repeat=25` failed on attempt 17 with nothing else running, which is far too often to wire into CI. Add BackendLogStore.SubscriberCount, resolving a model ID by the same exact-key and replica-prefix rules Subscribe uses, and have the specs poll it until the handler has attached. Nothing in production calls it and no assertion is weakened; the handler's own snapshot/subscribe window is left as it is, being a production streaming question rather than a test one. Verified with 60 repeats of the WebSocket specs and three consecutive --randomize-all runs of the whole distributed suite, all at --flake-attempts 1: 239 of 240 specs pass in about 80 seconds. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- pkg/model/backend_log_store.go | 39 ++++++++++++++++++++++ pkg/model/backend_log_store_test.go | 32 ++++++++++++++++++ tests/e2e/distributed/backend_logs_test.go | 18 ++++++++++ 3 files changed, 89 insertions(+) diff --git a/pkg/model/backend_log_store.go b/pkg/model/backend_log_store.go index c5b5253dd..30268fde4 100644 --- a/pkg/model/backend_log_store.go +++ b/pkg/model/backend_log_store.go @@ -344,3 +344,42 @@ func (s *BackendLogStore) Subscribe(modelID string) (chan BackendLogLine, func() return ch, unsubscribe } + +// SubscriberCount reports how many live subscriptions exist for modelID, +// resolving the ID with the same exact-key / replica-prefix rules as Subscribe. +// +// Streaming handlers send a GetLines snapshot before they call Subscribe, so a +// line appended between those two calls reaches the buffer but no channel. A +// caller that has to observe a line it appends itself must therefore wait for +// the subscription to exist rather than assume the handler got there first. +func (s *BackendLogStore) SubscriberCount(modelID string) int { + s.mu.RLock() + exact, exactOK := s.buffers[modelID] + var replicas []*backendLogBuffer + if !strings.Contains(modelID, replicaSeparator) { + prefix := modelID + replicaSeparator + for k, b := range s.buffers { + if strings.HasPrefix(k, prefix) { + replicas = append(replicas, b) + } + } + } + s.mu.RUnlock() + + // Counted after releasing s.mu: no other path takes s.mu and a buffer lock + // together, and keeping it that way costs nothing here. + count := func(buf *backendLogBuffer) int { + buf.mu.Lock() + defer buf.mu.Unlock() + return len(buf.subscribers) + } + + total := 0 + if exactOK { + total += count(exact) + } + for _, b := range replicas { + total += count(b) + } + return total +} diff --git a/pkg/model/backend_log_store_test.go b/pkg/model/backend_log_store_test.go index 775e07cdb..593bcf5a2 100644 --- a/pkg/model/backend_log_store_test.go +++ b/pkg/model/backend_log_store_test.go @@ -76,6 +76,38 @@ var _ = Describe("BackendLogStore", func() { }) }) + Describe("SubscriberCount", func() { + It("reports zero before anyone subscribes and drops back after unsubscribe", func() { + s.AppendLine("model-a", "stderr", "preload") + Expect(s.SubscriberCount("model-a")).To(Equal(0)) + + _, unsubscribe := s.Subscribe("model-a") + Expect(s.SubscriberCount("model-a")).To(Equal(1)) + + unsubscribe() + Expect(s.SubscriberCount("model-a")).To(Equal(0)) + }) + + // Subscribe resolves a bare model ID across every replica buffer, so the + // count has to follow the same rule or a caller waiting on it would give + // up while a perfectly good subscription was in place. + It("sums the replica buffers a bare model ID resolves to", func() { + s.AppendLine("model-a#0", "stderr", "preload-r0") + s.AppendLine("model-a#1", "stderr", "preload-r1") + + _, unsubscribe := s.Subscribe("model-a") + defer unsubscribe() + + Expect(s.SubscriberCount("model-a")).To(Equal(2)) + Expect(s.SubscriberCount("model-a#0")).To(Equal(1)) + Expect(s.SubscriberCount("model-b")).To(Equal(0)) + }) + + It("returns zero for a model that has no buffer at all", func() { + Expect(s.SubscriberCount("never-seen")).To(Equal(0)) + }) + }) + Describe("Subscribe", func() { // Confirms the WebSocket streaming path (the live tail UI) receives // lines from every replica when the caller subscribes by bare modelID. diff --git a/tests/e2e/distributed/backend_logs_test.go b/tests/e2e/distributed/backend_logs_test.go index 79dea3902..473dad54f 100644 --- a/tests/e2e/distributed/backend_logs_test.go +++ b/tests/e2e/distributed/backend_logs_test.go @@ -25,6 +25,21 @@ import ( "gorm.io/gorm/logger" ) +// waitForLogSubscriber blocks until the worker's WebSocket log handler has +// registered its subscription on the store. +// +// The handler writes the "initial" batch first and subscribes only afterwards, +// so a line appended the instant that batch lands is buffered but never +// streamed, and the spec then waits out its full read deadline. Measured at +// roughly one run in seventeen with `--repeat`, which is far too often for CI. +// Waiting on the subscription removes the race from the spec; the handler's own +// snapshot/subscribe window is a separate production question. +func waitForLogSubscriber(logStore *model.BackendLogStore, modelID string) { + GinkgoHelper() + Eventually(func() int { return logStore.SubscriberCount(modelID) }, "10s", "5ms"). + Should(BeNumerically(">", 0), "the WebSocket handler never subscribed to %q", modelID) +} + var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func() { Context("Worker HTTP log endpoints", func() { @@ -212,6 +227,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(initialLines[1].Text).To(Equal("line-2")) // Now append a new line and verify it arrives via WebSocket + waitForLogSubscriber(logStore, "ws-model") logStore.AppendLine("ws-model", "stdout", "line-3-realtime") conn.SetReadDeadline(time.Now().Add(5 * time.Second)) @@ -280,6 +296,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(conn.ReadJSON(&initialMsg)).To(Succeed()) // Append line to a different model + waitForLogSubscriber(logStore, "ws-model") logStore.AppendLine("other-model", "stdout", "should not appear") // Append line to our model logStore.AppendLine("ws-model", "stdout", "should appear") @@ -475,6 +492,7 @@ var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func Expect(initialLines[0].Text).To(Equal("initial line from worker")) // Append a new line on the worker's log store + waitForLogSubscriber(logStore, "proxy-model") logStore.AppendLine("proxy-model", "stderr", "realtime via proxy") // Read the streamed line through the proxy