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 <mudler@localai.io>
This commit is contained in:
Ettore Di Giacinto committed 2026-08-31 10:31:51 +00:00
1 parent 1974bc1ea0
commit f0fa4a7b1f
3 files changed
+89

No files matched your search

@@ -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