mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
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:
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{
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user