mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
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:
1 parent
1974bc1ea0
commit
f0fa4a7b1f
3 files changed
+89
No files matched your search
@@ -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
|
||||
}
|
||||
@@ -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.
|
||||
|
||||
Reference in new issue
Block a user