test(distributed): scope the log-subscriber wait and mark the race it works around

Three corrections from review of the previous commit.

The lock-order comment on SubscriberCount claimed no path takes s.mu and a
buffer lock together. Subscribe does exactly that, holding s.mu.RLock across
replica registrations that take buf.mu. State the rule that is actually true —
s.mu precedes any buffer lock, so counting after releasing it preserves the
order — and say what follows from it: the total is a sample, not a snapshot.

waitForLogSubscriber read as general-purpose but unblocks on the first
registered subscription. Subscribe attaches the exact-key buffer and each
replica buffer one at a time, so for a replicated model the count goes positive
while later replicas are still unattached and the race survives. Rename it
waitForSingleLogSubscriber, document that it holds only where Subscribe
resolves to one buffer, and assert on exactly 1: misuse then fails loudly on
the count rather than going quietly back to being flaky. Taking the expected
count as a parameter was the alternative, but that makes callers predict a
store-internal number and an under-count fails the same silent way as the
original bug.

The snapshot-then-subscribe race had no artifact outside a report, and review
found a second site carrying it. Mark both handlers identically, including the
point that swapping the two calls duplicates rather than drops and so is not
the fix. The race itself is left alone; this branch stays test infrastructure.

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:42:49 +00:00
1 parent f0fa4a7b1f
commit 53639c4df3
4 files changed
+36 -9

No files matched your search

@@ -121,6 +121,13 @@ func BackendLogsWebSocketEndpoint(ml *model.ModelLoader) echo.HandlerFunc {
conn := &backendLogsConn{Conn: ws}
// KNOWN RACE: the snapshot is sent before the subscription is registered, so
// a line appended in that window is never streamed. A viewer attaching while
// a model loads (when a backend is at its noisiest) can silently miss lines;
// they stay in the buffer, so a reload shows them. Fixing it needs an atomic
// snapshot-plus-subscribe under the store lock, not a reorder of these two
// calls, which would duplicate instead of drop.
// Send existing lines as initial batch
existingLines := ml.BackendLogs().GetLines(modelID)
initialMsg := map[string]any{
@@ -839,6 +839,13 @@ func handleBackendLogsWS(w http.ResponseWriter, r *http.Request, logStore *model
conn := &backendLogsWSConn{Conn: ws}
// KNOWN RACE: the snapshot is sent before the subscription is registered, so
// a line appended in that window is never streamed. A viewer attaching while
// a model loads (when a backend is at its noisiest) can silently miss lines;
// they stay in the buffer, so a reload shows them. Fixing it needs an atomic
// snapshot-plus-subscribe under the store lock, not a reorder of these two
// calls, which would duplicate instead of drop.
// Send existing lines as initial batch
existingLines := logStore.GetLines(modelID)
initialMsg := map[string]any{
+5 -2
View File
@@ -366,8 +366,11 @@ func (s *BackendLogStore) SubscriberCount(modelID string) int {
}
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.
// Lock order in this type is always s.mu before any buffer lock — Subscribe
// holds s.mu.RLock across its replica registrations, which take buf.mu — so
// counting after releasing s.mu keeps that order rather than inverting it.
// The total is therefore a sample, not a snapshot: a concurrent Subscribe
// can register a further buffer while this loop runs.
count := func(buf *backendLogBuffer) int {
buf.mu.Lock()
defer buf.mu.Unlock()
+17 -7
View File
@@ -25,7 +25,7 @@ import (
"gorm.io/gorm/logger"
)
// waitForLogSubscriber blocks until the worker's WebSocket log handler has
// waitForSingleLogSubscriber 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,
@@ -33,11 +33,21 @@ import (
// 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) {
// snapshot/subscribe window is a separate production question, marked at both
// production sites.
//
// Only valid where BackendLogStore.Subscribe resolves modelID to exactly ONE
// buffer: a bare model ID with no "<modelID>#N" replica buffers in the store, or
// a full process key. Subscribe registers the exact-key buffer and each replica
// buffer one at a time, so for a model that does have replicas the count goes
// positive while later replicas are still unattached and the race survives.
// Hence the assertion is on exactly 1 rather than "at least 1": a spec that
// misapplies this to a replicated model fails loudly on the count instead of
// going quietly back to being flaky.
func waitForSingleLogSubscriber(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)
Should(Equal(1), "the WebSocket handler never subscribed to %q exactly once", modelID)
}
var _ = Describe("Distributed Backend Log Streaming", Label("Distributed"), func() {
@@ -227,7 +237,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")
waitForSingleLogSubscriber(logStore, "ws-model")
logStore.AppendLine("ws-model", "stdout", "line-3-realtime")
conn.SetReadDeadline(time.Now().Add(5 * time.Second))
@@ -296,7 +306,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")
waitForSingleLogSubscriber(logStore, "ws-model")
logStore.AppendLine("other-model", "stdout", "should not appear")
// Append line to our model
logStore.AppendLine("ws-model", "stdout", "should appear")
@@ -492,7 +502,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")
waitForSingleLogSubscriber(logStore, "proxy-model")
logStore.AppendLine("proxy-model", "stderr", "realtime via proxy")
// Read the streamed line through the proxy