mirror of
https://github.com/mudler/LocalAI.git
synced 2026-09-12 22:33:54 -04:00
fix(distributed): order LISTEN and UNLISTEN on one lock
Unsubscribe decided a channel had lost its last subscriber under one lock and issued the UNLISTEN after releasing it. A Subscribe on the same root could decide to LISTEN in that window, and the two reached the connection in the wrong order: the root ended up not listened with a live subscription on it. It does not heal, because the next Subscribe sees the registration already there and never re-LISTENs, so the whole root stays deaf on that replica until the connection drops. The decision and the statement it implies now happen under one lock, held across both, at both call sites. A second lock and not the registration lock: issuing waits on the listener goroutine, delivery takes the registration lock, and holding that across the wait deadlocks the carrier. The race is spec'd through a barrier seam rather than by racing goroutines. The natural window is microseconds wide, and a spec that waits for it to open passes by luck; the seam scripts the interleaving, so the spec decides in both directions. Resolving a spilled message moved off the listener. PostgreSQL keeps undelivered notifications in a shared, fixed-size queue, so a listener that stops draining it can block COMMIT for every publisher on the server, not only this one. The listener now only drains; one resolver goroutine reads the row back and dispatches, which also keeps a spilled message and an inline one on the same subject in the order they were published. Three wiring lines that could be deleted with the suite staying green: the sweeper's start is now pinned by a Config interval, and the two lines that carry the bus into the deployment now refuse to boot when either is missing. A subscription can also report what it dropped, so the party that missed a message is the party that can see it. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
This commit is contained in:
1 parent
cf619fc91b
commit
8f71c08d94
8 files changed
+528
-42
No files matched your search
@@ -642,7 +642,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
modelAdapter := nodes.NewModelRouterAdapter(router)
|
||||
|
||||
success = true
|
||||
return &DistributedServices{
|
||||
ds := &DistributedServices{
|
||||
Nats: natsClient,
|
||||
Store: store,
|
||||
Registry: registry,
|
||||
@@ -667,7 +667,38 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
|
||||
WorkerDialer: workerDialer,
|
||||
BackendClients: backendClients,
|
||||
Bus: bus,
|
||||
}, nil
|
||||
}
|
||||
// Checked once, here, on the assembled struct. See requireBroadcastCarrier.
|
||||
if err := requireBroadcastCarrier(ds); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ds, nil
|
||||
}
|
||||
|
||||
// requireBroadcastCarrier refuses to hand back a distributed deployment whose
|
||||
// broadcast carrier is missing.
|
||||
//
|
||||
// The carrier reaches the deployment over two lines: the newBroadcastBus call
|
||||
// in initDistributed, and the Bus field in the twenty-three field literal
|
||||
// above. Deleting either one compiles and leaves every suite in this repository
|
||||
// green, and the two failures are different. Without the construction, nothing
|
||||
// can ever be published between replicas. Without the assignment the carrier is
|
||||
// opened and connected but Shutdown cannot see it, so every restart leaves a
|
||||
// pinned PostgreSQL session and its goroutines behind until the server runs out
|
||||
// of connections, and the operator sees the failure land on whatever connects
|
||||
// next rather than on LocalAI.
|
||||
//
|
||||
// Neither line can be reddened by a spec today: initDistributed opens NATS
|
||||
// before it reaches any of this, so it cannot be called from a unit test, and a
|
||||
// pointer field left out of a struct literal is not a compile error. What this
|
||||
// converts both omissions into is a deployment that refuses to start and names
|
||||
// what is missing, which is as far as they can be pinned until initDistributed
|
||||
// is testable. The guard itself is spec'd.
|
||||
func requireBroadcastCarrier(ds *DistributedServices) error {
|
||||
if ds == nil || ds.Bus == nil {
|
||||
return fmt.Errorf("distributed mode was initialized without a broadcast carrier: nothing could be published between replicas, and the PostgreSQL session it pins could not be closed on shutdown")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// newBroadcastBus opens the deployment's fan-out carrier on the auth database.
|
||||
|
||||
@@ -61,6 +61,36 @@ var _ = Describe("opening the deployment's broadcast carrier", func() {
|
||||
})
|
||||
})
|
||||
|
||||
// The partial pin on two wiring lines that cannot be reddened by a spec: the
|
||||
// newBroadcastBus call, and `Bus: bus` in the returned literal. Neither is a
|
||||
// compile error when deleted and initDistributed cannot be unit tested while it
|
||||
// opens NATS first, so what is available is a boot refusal, and this is what
|
||||
// keeps that refusal honest.
|
||||
var _ = Describe("refusing a deployment with no broadcast carrier", func() {
|
||||
It("accepts services that carry one", func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
cfg := &config.ApplicationConfig{}
|
||||
cfg.Auth.DatabaseURL = dsn
|
||||
bus, err := newBroadcastBus(context.Background(), cfg, db)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(bus.Close)
|
||||
|
||||
Expect(requireBroadcastCarrier(&DistributedServices{Bus: bus})).To(Succeed())
|
||||
})
|
||||
|
||||
It("refuses services whose carrier was never assigned, and says what it costs", func() {
|
||||
err := requireBroadcastCarrier(&DistributedServices{})
|
||||
|
||||
Expect(err).To(HaveOccurred())
|
||||
Expect(err.Error()).To(ContainSubstring("published between replicas"))
|
||||
Expect(err.Error()).To(ContainSubstring("shutdown"))
|
||||
})
|
||||
|
||||
It("refuses a nil deployment rather than dereferencing it", func() {
|
||||
Expect(requireBroadcastCarrier(nil)).ToNot(Succeed())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("shutting the distributed services down", func() {
|
||||
It("closes the broadcast carrier", func() {
|
||||
// A pinned PostgreSQL session and the goroutine parked on it, per
|
||||
|
||||
@@ -1,6 +1,8 @@
|
||||
package messaging_test
|
||||
|
||||
import (
|
||||
"errors"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
@@ -58,7 +60,11 @@ var _ = DescribeTable("ValidFilter",
|
||||
func(filter string, wantErr bool) {
|
||||
err := messaging.ValidFilter(filter)
|
||||
if wantErr {
|
||||
Expect(err).To(HaveOccurred())
|
||||
// The CLASS, not merely "an error". Carriers match on
|
||||
// ErrUnsupportedFilter to tell "this caller asked for something we
|
||||
// do not implement" apart from "the store is unhappy", so the
|
||||
// definition has to pin what the callers match on.
|
||||
Expect(errors.Is(err, messaging.ErrUnsupportedFilter)).To(BeTrue(), "got %v", err)
|
||||
return
|
||||
}
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
+193
-17
@@ -80,6 +80,23 @@ const subscribeTimeout = 30 * time.Second
|
||||
// wedging the carrier silently is not.
|
||||
const deliveryQueueDepth = 256
|
||||
|
||||
// notificationQueueDepth is how many notifications may be waiting to be
|
||||
// resolved and dispatched before the carrier starts dropping them.
|
||||
//
|
||||
// It exists so that resolving a spilled broadcast, which is one SELECT, never
|
||||
// happens on the goroutine that drains PostgreSQL's notification stream. That
|
||||
// goroutine falling behind does not merely delay this replica: PostgreSQL holds
|
||||
// undelivered notifications in a shared, fixed-size async queue, and a listener
|
||||
// that stops draining it can fill that queue and block COMMIT for every
|
||||
// publisher on the SERVER, LocalAI's or not. Fourteen traffic types are moving
|
||||
// onto this carrier, several of which spill by construction, so this is a
|
||||
// hazard the carrier has to own rather than one to leave to its adopters.
|
||||
//
|
||||
// Dropping locally when the resolver falls this far behind is the right trade
|
||||
// against that: a lost broadcast is recoverable and loud, a stalled server is
|
||||
// neither.
|
||||
const notificationQueueDepth = 1024
|
||||
|
||||
// broadcastRoots is the closed set of subject roots this carrier serves.
|
||||
// A subject whose first token is not here is REFUSED at publish and at
|
||||
// subscribe, rather than being mapped to a channel of its own.
|
||||
@@ -152,6 +169,15 @@ type Config struct {
|
||||
DSN string
|
||||
// DB is the pooled handle NOTIFY and the spill table are written on.
|
||||
DB *gorm.DB
|
||||
// SweepInterval is how often this carrier retires spilled broadcasts that
|
||||
// have aged out. Zero means spillSweepInterval.
|
||||
//
|
||||
// It is a field rather than a constant so that the sweeper being STARTED is
|
||||
// a testable fact. SweepSpill and SpillSweepSQL can both be exercised
|
||||
// directly, and neither of them proves the carrier ever calls them: a
|
||||
// deleted `go b.sweep()` left the whole suite green and bus_messages
|
||||
// growing forever.
|
||||
SweepInterval time.Duration
|
||||
}
|
||||
|
||||
// notification is what travels in a pg_notify payload. The keys are one byte
|
||||
@@ -185,14 +211,43 @@ type Bus struct {
|
||||
connected atomic.Bool
|
||||
|
||||
cmds chan listenCmd
|
||||
inbound chan inbound
|
||||
listenerDone chan struct{}
|
||||
resolverDone chan struct{}
|
||||
closeOnce sync.Once
|
||||
|
||||
// listenMu orders a channel's refcount decision with the LISTEN or
|
||||
// UNLISTEN that decision implies, as ONE step.
|
||||
//
|
||||
// Deciding under mu and issuing outside it is a real defect and not a
|
||||
// theoretical one. A last Unsubscribe that has decided to UNLISTEN can be
|
||||
// overtaken by a Subscribe that has decided to LISTEN; the two reach the
|
||||
// connection in that order; the channel ends up not listened with a live
|
||||
// subscription on it. It does not self-heal, because the next Subscribe on
|
||||
// that root sees first == false and never re-LISTENs, so the whole root is
|
||||
// silently deaf on this replica until a connection drop triggers relisten.
|
||||
//
|
||||
// A second lock rather than mu, because command waits on the listener
|
||||
// goroutine and delivery takes mu: holding mu across that wait would
|
||||
// deadlock the carrier. Neither the listener nor the resolver ever takes
|
||||
// this one.
|
||||
listenMu sync.Mutex
|
||||
|
||||
// listenBarrier is a test seam and is nil in production. See barrier.
|
||||
listenBarrier func(stage, op string)
|
||||
|
||||
mu sync.Mutex
|
||||
nextID uint64
|
||||
subs map[string]map[uint64]*subscription
|
||||
}
|
||||
|
||||
// inbound is one notification as it came off the connection, before it is
|
||||
// decoded, resolved and dispatched.
|
||||
type inbound struct {
|
||||
channel string
|
||||
payload string
|
||||
}
|
||||
|
||||
// New opens the carrier: one pinned LISTEN connection, and the pooled handle
|
||||
// publishes travel on.
|
||||
func New(ctx context.Context, cfg Config) (*Bus, error) {
|
||||
@@ -225,11 +280,14 @@ func New(ctx context.Context, cfg Config) (*Bus, error) {
|
||||
ctx: busCtx,
|
||||
cancel: cancel,
|
||||
cmds: make(chan listenCmd),
|
||||
inbound: make(chan inbound, notificationQueueDepth),
|
||||
listenerDone: make(chan struct{}),
|
||||
resolverDone: make(chan struct{}),
|
||||
subs: map[string]map[uint64]*subscription{},
|
||||
}
|
||||
b.connected.Store(true)
|
||||
go b.listen(conn)
|
||||
go b.resolve()
|
||||
go b.sweep()
|
||||
return b, nil
|
||||
}
|
||||
@@ -274,6 +332,7 @@ func (b *Bus) Close() {
|
||||
b.closeOnce.Do(func() {
|
||||
b.cancel()
|
||||
<-b.listenerDone
|
||||
<-b.resolverDone
|
||||
b.connected.Store(false)
|
||||
})
|
||||
}
|
||||
@@ -335,6 +394,12 @@ func (b *Bus) Subscribe(subject string, handler func([]byte)) (messaging.Subscri
|
||||
return nil, fmt.Errorf("pgbus: no handler for %q", subject)
|
||||
}
|
||||
|
||||
// Before the lock, so a spec can observe that this registration has
|
||||
// started even when the lock is what stops it going any further.
|
||||
b.barrier("enter", "LISTEN")
|
||||
b.listenMu.Lock()
|
||||
defer b.listenMu.Unlock()
|
||||
|
||||
b.mu.Lock()
|
||||
b.nextID++
|
||||
sub := &subscription{
|
||||
@@ -344,6 +409,7 @@ func (b *Bus) Subscribe(subject string, handler func([]byte)) (messaging.Subscri
|
||||
id: b.nextID,
|
||||
queue: make(chan []byte, deliveryQueueDepth),
|
||||
stop: make(chan struct{}),
|
||||
done: make(chan struct{}),
|
||||
}
|
||||
first := len(b.subs[channel]) == 0
|
||||
if first {
|
||||
@@ -355,18 +421,54 @@ func (b *Bus) Subscribe(subject string, handler func([]byte)) (messaging.Subscri
|
||||
go sub.run(handler)
|
||||
|
||||
if first {
|
||||
b.barrier("issue", "LISTEN")
|
||||
if err := b.command("LISTEN " + pgx.Identifier{channel}.Sanitize()); err != nil {
|
||||
_ = sub.Unsubscribe()
|
||||
// Not Unsubscribe: the ordering lock is already held here, and the
|
||||
// connection was never listening on this channel, so there is
|
||||
// nothing to UNLISTEN.
|
||||
b.forget(sub)
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
return sub, nil
|
||||
}
|
||||
|
||||
// forget removes a registration without touching the channel's LISTEN state.
|
||||
func (b *Bus) forget(sub *subscription) {
|
||||
sub.once.Do(func() {
|
||||
b.mu.Lock()
|
||||
delete(b.subs[sub.channel], sub.id)
|
||||
if len(b.subs[sub.channel]) == 0 {
|
||||
delete(b.subs, sub.channel)
|
||||
}
|
||||
b.mu.Unlock()
|
||||
close(sub.stop)
|
||||
})
|
||||
}
|
||||
|
||||
// barrier is a test seam and does nothing in production.
|
||||
//
|
||||
// It exists because the ordering listenMu enforces cannot be observed from
|
||||
// outside this package and cannot be provoked from outside it either: the
|
||||
// natural window is microseconds wide, and it was measured at zero hits in
|
||||
// forty attempts while being ten out of ten once widened. A spec that waits for
|
||||
// that window to open is a spec that passes by luck, which is worse than no
|
||||
// spec at all for a defect that leaves a whole subject root deaf.
|
||||
func (b *Bus) barrier(stage, op string) {
|
||||
if b.listenBarrier != nil {
|
||||
b.listenBarrier(stage, op)
|
||||
}
|
||||
}
|
||||
|
||||
// command hands a LISTEN or UNLISTEN to the goroutine that owns the connection
|
||||
// and waits for it, so a Subscribe that has returned is a registration the
|
||||
// server has already acknowledged. Returning before that would lose every
|
||||
// message published in the gap.
|
||||
// and waits for it. Returning before the server had acknowledged it would lose
|
||||
// every message published in the gap.
|
||||
//
|
||||
// A Subscribe that has returned is therefore always a registration the server
|
||||
// has acknowledged, including the case where this Subscribe issued nothing
|
||||
// because the channel was already listened: listenMu means the Subscribe that
|
||||
// DID issue the LISTEN had already been acknowledged before this one could see
|
||||
// its registration.
|
||||
func (b *Bus) command(sql string) error {
|
||||
cmd := listenCmd{sql: sql, done: make(chan error, 1)}
|
||||
timeout := time.NewTimer(subscribeTimeout)
|
||||
@@ -439,7 +541,34 @@ func (b *Bus) listen(conn *pgx.Conn) {
|
||||
conn = replacement
|
||||
continue
|
||||
}
|
||||
b.deliver(n.Channel, n.Payload)
|
||||
b.offer(n.Channel, n.Payload)
|
||||
}
|
||||
}
|
||||
|
||||
// offer hands a notification to the resolver. It never blocks: the listener's
|
||||
// only job is to keep PostgreSQL's async queue draining.
|
||||
func (b *Bus) offer(channel, payload string) {
|
||||
select {
|
||||
case b.inbound <- inbound{channel: channel, payload: payload}:
|
||||
default:
|
||||
xlog.Error("Broadcast carrier dropped a notification: the resolver is not keeping up",
|
||||
"channel", channel, "depth", notificationQueueDepth)
|
||||
}
|
||||
}
|
||||
|
||||
// resolve is where a spilled broadcast is read back and where every
|
||||
// notification is dispatched. Both are off the listener on purpose, and both
|
||||
// are on ONE goroutine, so a spilled message and an inline one on the same
|
||||
// subject keep the order they were published in.
|
||||
func (b *Bus) resolve() {
|
||||
defer close(b.resolverDone)
|
||||
for {
|
||||
select {
|
||||
case <-b.ctx.Done():
|
||||
return
|
||||
case in := <-b.inbound:
|
||||
b.deliver(in.channel, in.payload)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -500,6 +629,8 @@ func (b *Bus) relisten(conn *pgx.Conn) error {
|
||||
// deliver resolves one notification and hands it to the subscribers whose
|
||||
// filters match.
|
||||
func (b *Bus) deliver(channel, payload string) {
|
||||
b.barrier("enter", "DELIVER")
|
||||
|
||||
var n notification
|
||||
if err := json.Unmarshal([]byte(payload), &n); err != nil {
|
||||
xlog.Error("Broadcast carrier received an undecodable notification", "channel", channel, "error", err)
|
||||
@@ -507,6 +638,14 @@ func (b *Bus) deliver(channel, payload string) {
|
||||
}
|
||||
|
||||
data := []byte(n.Data)
|
||||
if n.SpillID == "" && len(data) == 0 {
|
||||
// Publish cannot produce this: json.Marshal never returns an empty
|
||||
// encoding. It is reachable only if something other than this carrier
|
||||
// notifies on a localai_ channel, and delivering nil to a handler
|
||||
// would be a message that says nothing rather than no message at all.
|
||||
xlog.Error("Broadcast carrier received a notification with no payload", "channel", channel, "subject", n.Subject)
|
||||
return
|
||||
}
|
||||
if n.SpillID != "" {
|
||||
resolved, err := b.resolveSpill(n.SpillID)
|
||||
if err != nil {
|
||||
@@ -544,12 +683,37 @@ type subscription struct {
|
||||
filter string
|
||||
id uint64
|
||||
|
||||
queue chan []byte
|
||||
stop chan struct{}
|
||||
once sync.Once
|
||||
queue chan []byte
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
dropped atomic.Uint64
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
// DropCounter is the optional interface a Subscription satisfies when it can
|
||||
// report how many broadcasts it lost.
|
||||
//
|
||||
// It exists because the drop is otherwise invisible to the party that needs to
|
||||
// know. The error log lands on the receiving replica, there is no sequence
|
||||
// number and no gap signal, so "anything that must survive a gap belongs in a
|
||||
// table, with the broadcast as a hint to go and look" cannot be acted on by the
|
||||
// subscriber that missed the hint. A subscriber whose subject has no successor
|
||||
// message, a job result rather than a progress tick, can type-assert to this
|
||||
// and go read the row.
|
||||
//
|
||||
// It is not on messaging.Subscription: the NATS client's subscription cannot
|
||||
// answer it, and widening that interface would make every existing consumer
|
||||
// claim a guarantee it does not have.
|
||||
type DropCounter interface {
|
||||
Dropped() uint64
|
||||
}
|
||||
|
||||
// Dropped reports how many broadcasts this subscription lost because its
|
||||
// handler was too far behind. It only ever grows.
|
||||
func (s *subscription) Dropped() uint64 { return s.dropped.Load() }
|
||||
|
||||
func (s *subscription) run(handler func([]byte)) {
|
||||
defer close(s.done)
|
||||
for {
|
||||
select {
|
||||
case data := <-s.queue:
|
||||
@@ -566,8 +730,9 @@ func (s *subscription) enqueue(subject string, data []byte) {
|
||||
select {
|
||||
case s.queue <- data:
|
||||
default:
|
||||
s.dropped.Add(1)
|
||||
xlog.Error("Broadcast carrier dropped a message: subscriber is not keeping up",
|
||||
"filter", s.filter, "subject", subject, "depth", deliveryQueueDepth)
|
||||
"filter", s.filter, "subject", subject, "depth", deliveryQueueDepth, "dropped", s.dropped.Load())
|
||||
}
|
||||
}
|
||||
|
||||
@@ -576,17 +741,24 @@ func (s *subscription) enqueue(subject string, data []byte) {
|
||||
func (s *subscription) Unsubscribe() error {
|
||||
var err error
|
||||
s.once.Do(func() {
|
||||
s.bus.mu.Lock()
|
||||
delete(s.bus.subs[s.channel], s.id)
|
||||
last := len(s.bus.subs[s.channel]) == 0
|
||||
b := s.bus
|
||||
b.barrier("enter", "UNLISTEN")
|
||||
// The whole decision AND its issuance, as one step. See listenMu.
|
||||
b.listenMu.Lock()
|
||||
defer b.listenMu.Unlock()
|
||||
|
||||
b.mu.Lock()
|
||||
delete(b.subs[s.channel], s.id)
|
||||
last := len(b.subs[s.channel]) == 0
|
||||
if last {
|
||||
delete(s.bus.subs, s.channel)
|
||||
delete(b.subs, s.channel)
|
||||
}
|
||||
s.bus.mu.Unlock()
|
||||
b.mu.Unlock()
|
||||
|
||||
close(s.stop)
|
||||
if last && s.bus.ctx.Err() == nil {
|
||||
err = s.bus.command("UNLISTEN " + pgx.Identifier{s.channel}.Sanitize())
|
||||
if last && b.ctx.Err() == nil {
|
||||
b.barrier("issue", "UNLISTEN")
|
||||
err = b.command("UNLISTEN " + pgx.Identifier{s.channel}.Sanitize())
|
||||
}
|
||||
})
|
||||
return err
|
||||
@@ -612,7 +784,11 @@ func (b *Bus) resolveSpill(id string) ([]byte, error) {
|
||||
|
||||
// sweep retires spilled rows that every replica has had time to read.
|
||||
func (b *Bus) sweep() {
|
||||
ticker := time.NewTicker(spillSweepInterval)
|
||||
interval := b.cfg.SweepInterval
|
||||
if interval <= 0 {
|
||||
interval = spillSweepInterval
|
||||
}
|
||||
ticker := time.NewTicker(interval)
|
||||
defer ticker.Stop()
|
||||
for {
|
||||
select {
|
||||
|
||||
@@ -5,7 +5,6 @@ package pgbus_test
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"path/filepath"
|
||||
"time"
|
||||
|
||||
@@ -209,6 +208,34 @@ var _ = Describe("the PostgreSQL broadcast carrier", func() {
|
||||
})
|
||||
})
|
||||
|
||||
Describe("a subscriber that falls behind", func() {
|
||||
It("counts what it lost, so the subscriber can go and read the row", func() {
|
||||
// The drop is loud in the log of the replica that took it, and that
|
||||
// is the wrong party: the subscriber is the one that has to decide
|
||||
// to go and read the table instead. Task 12 moves jobs.*.result
|
||||
// onto this carrier, and a result has no successor message, so an
|
||||
// invisible loss there is a job whose answer silently never
|
||||
// arrives.
|
||||
block := make(chan struct{})
|
||||
DeferCleanup(func() { close(block) })
|
||||
s, err := sub.Subscribe("jobs.slow.consumer", func([]byte) { <-block })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
counter, ok := s.(pgbus.DropCounter)
|
||||
Expect(ok).To(BeTrue(), "a subscription must be able to report its losses")
|
||||
Expect(counter.Dropped()).To(BeZero())
|
||||
|
||||
// Comfortably past the queue depth, stated absolutely rather than
|
||||
// as deliveryQueueDepth+n, so a change to the depth moves the
|
||||
// behaviour and not the expectation.
|
||||
for i := 0; i < 400; i++ {
|
||||
Expect(pub.Publish("jobs.slow.consumer", map[string]int{"i": i})).To(Succeed())
|
||||
}
|
||||
|
||||
Eventually(counter.Dropped, 30*time.Second).Should(BeNumerically(">", 0))
|
||||
})
|
||||
})
|
||||
|
||||
Describe("losing the LISTEN connection", func() {
|
||||
It("reconnects, restores its registrations and delivers again", func() {
|
||||
// The transport failure a fake cannot produce. Terminating the
|
||||
@@ -328,13 +355,3 @@ var _ = Describe("constructing the carrier", func() {
|
||||
Expect(err).To(HaveOccurred())
|
||||
})
|
||||
})
|
||||
|
||||
var _ = Describe("the interface the carrier is interchangeable through", func() {
|
||||
It("is satisfied by the bus", func() {
|
||||
// Compile-time, not behavioural: the point is that a signature change
|
||||
// on either carrier fails here rather than in whichever call site is
|
||||
// migrated next.
|
||||
var b messaging.Broadcaster = (*pgbus.Bus)(nil)
|
||||
Expect(fmt.Sprintf("%T", b)).To(Equal("*pgbus.Bus"))
|
||||
})
|
||||
})
|
||||
@@ -0,0 +1,178 @@
|
||||
// SPDX-License-Identifier: MIT
|
||||
|
||||
package pgbus
|
||||
|
||||
import (
|
||||
"context"
|
||||
"sync"
|
||||
|
||||
. "github.com/onsi/ginkgo/v2"
|
||||
. "github.com/onsi/gomega"
|
||||
|
||||
"github.com/mudler/LocalAI/core/services/testutil"
|
||||
)
|
||||
|
||||
// Internal on purpose. What these specs are about is the ORDER in which two
|
||||
// registrations reach the connection, and that order is not observable from
|
||||
// outside the package: a channel that is subscribed but not listened behaves
|
||||
// exactly like a deployment where nobody is publishing.
|
||||
//
|
||||
// They drive that order through the barrier seam rather than by racing
|
||||
// goroutines and hoping. The natural window here was measured at zero hits in
|
||||
// forty attempts and ten out of ten once widened by hand, so a spec that races
|
||||
// for it would pass by luck on a defect that leaves a whole subject root deaf
|
||||
// until the connection drops.
|
||||
var _ = Describe("ordering LISTEN and UNLISTEN on one root", func() {
|
||||
var b *Bus
|
||||
|
||||
newBus := func(cfg Config) *Bus {
|
||||
GinkgoHelper()
|
||||
built, err := New(context.Background(), cfg)
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(built.Close)
|
||||
return built
|
||||
}
|
||||
|
||||
BeforeEach(func() {
|
||||
db, dsn := testutil.SetupTestDBWithDSN()
|
||||
Expect(Migrate(context.Background(), db)).To(Succeed())
|
||||
b = newBus(Config{DSN: dsn, DB: db})
|
||||
})
|
||||
|
||||
It("does not let a new LISTEN overtake the UNLISTEN it races", func() {
|
||||
leaving, err := b.Subscribe("jobs.race.leaving", func([]byte) {})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
unlistenDecided := make(chan struct{})
|
||||
release := make(chan struct{})
|
||||
arrived := make(chan struct{})
|
||||
listenIssued := make(chan struct{}, 1)
|
||||
// Set after the first Subscribe, so only the racing pair is observed.
|
||||
b.listenBarrier = func(stage, op string) {
|
||||
switch {
|
||||
case op == "UNLISTEN" && stage == "issue":
|
||||
close(unlistenDecided)
|
||||
<-release
|
||||
case op == "LISTEN" && stage == "enter":
|
||||
close(arrived)
|
||||
case op == "LISTEN" && stage == "issue":
|
||||
listenIssued <- struct{}{}
|
||||
}
|
||||
}
|
||||
|
||||
unsubscribed := make(chan error, 1)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
unsubscribed <- leaving.Unsubscribe()
|
||||
}()
|
||||
// The last unsubscribe has decided to UNLISTEN and has not issued it.
|
||||
// Eventually rather than a bare receive: a carrier that never reaches
|
||||
// the decision at all must fail this spec, not hang it.
|
||||
Eventually(unlistenDecided).Should(BeClosed())
|
||||
|
||||
delivered := make(chan []byte, 1)
|
||||
subscribed := make(chan error, 1)
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
_, err := b.Subscribe("jobs.race.arriving", func(d []byte) { delivered <- d })
|
||||
subscribed <- err
|
||||
}()
|
||||
// The competing subscribe has started. Waiting for this rather than for
|
||||
// a duration is what makes the assertion below decide something.
|
||||
Eventually(arrived).Should(BeClosed())
|
||||
|
||||
// The bite, and it is deterministic in both directions. Unserialized,
|
||||
// this LISTEN is on the wire within microseconds of the line above and
|
||||
// the UNLISTEN then undoes it. Serialized, it provably cannot be
|
||||
// issued while the unsubscribe holds the ordering lock.
|
||||
Consistently(listenIssued).ShouldNot(Receive())
|
||||
|
||||
close(release)
|
||||
Eventually(unsubscribed).Should(Receive(BeNil()))
|
||||
Eventually(subscribed).Should(Receive(BeNil()))
|
||||
|
||||
// The consequence, stated as delivery: the root must not be deaf.
|
||||
Expect(b.Publish("jobs.race.arriving", map[string]string{"m": "after"})).To(Succeed())
|
||||
Eventually(delivered).Should(Receive(MatchJSON(`{"m":"after"}`)))
|
||||
})
|
||||
|
||||
It("listens once for a root and stops listening when its last subscriber leaves", func() {
|
||||
// The half of Unsubscribe that the delivery spec cannot see. Without
|
||||
// the registration being removed from the map, the second unsubscribe
|
||||
// never decides it was the last, and the connection keeps listening on
|
||||
// a root nothing is subscribed to for the life of the process.
|
||||
issued := make(chan string, 4)
|
||||
b.listenBarrier = func(stage, op string) {
|
||||
if stage == "issue" {
|
||||
issued <- op
|
||||
}
|
||||
}
|
||||
|
||||
one, err := b.Subscribe("jobs.refcount.a", func([]byte) {})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
two, err := b.Subscribe("jobs.refcount.b", func([]byte) {})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(one.Unsubscribe()).To(Succeed())
|
||||
Expect(two.Unsubscribe()).To(Succeed())
|
||||
|
||||
Expect(issued).To(HaveLen(2))
|
||||
Expect(<-issued).To(Equal("LISTEN"))
|
||||
Expect(<-issued).To(Equal("UNLISTEN"))
|
||||
})
|
||||
|
||||
It("waits for the resolver as well as the listener when it closes", func() {
|
||||
// Close returning while the resolver is still dispatching hands the
|
||||
// caller a bus whose handlers can still fire after shutdown, against a
|
||||
// database handle the process is about to drop.
|
||||
//
|
||||
// Asserted as "Close has not returned yet" rather than as "the done
|
||||
// channel is closed afterwards": the resolver exits on its own when the
|
||||
// context is cancelled, so the second shape passes whether Close waits
|
||||
// or not.
|
||||
dispatching := make(chan struct{}, 1)
|
||||
release := make(chan struct{})
|
||||
var releaseOnce sync.Once
|
||||
// Registered inside the It, so it runs BEFORE the bus teardown the
|
||||
// BeforeEach registered. A failing assertion below must not leave the
|
||||
// resolver parked on a channel that Close is waiting for.
|
||||
DeferCleanup(func() { releaseOnce.Do(func() { close(release) }) })
|
||||
b.listenBarrier = func(stage, op string) {
|
||||
if op == "DELIVER" && stage == "enter" {
|
||||
select {
|
||||
case dispatching <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
<-release
|
||||
}
|
||||
}
|
||||
_, err := b.Subscribe("jobs.close.pending", func([]byte) {})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
Expect(b.Publish("jobs.close.pending", map[string]string{"m": "x"})).To(Succeed())
|
||||
Eventually(dispatching).Should(Receive())
|
||||
|
||||
closed := make(chan struct{})
|
||||
go func() {
|
||||
defer GinkgoRecover()
|
||||
b.Close()
|
||||
close(closed)
|
||||
}()
|
||||
|
||||
Consistently(closed).ShouldNot(BeClosed())
|
||||
releaseOnce.Do(func() { close(release) })
|
||||
Eventually(closed).Should(BeClosed())
|
||||
})
|
||||
|
||||
It("stops the delivery goroutine of a subscription that has left", func() {
|
||||
// The other half. A subscription that is deregistered but whose runner
|
||||
// is never stopped leaks a goroutine per Unsubscribe, and the delivery
|
||||
// spec cannot tell the two halves apart because either one alone makes
|
||||
// the handler silent.
|
||||
s, err := b.Subscribe("jobs.refcount.runner", func([]byte) {})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(s.Unsubscribe()).To(Succeed())
|
||||
|
||||
Eventually(s.(*subscription).done).Should(BeClosed())
|
||||
})
|
||||
})
|
||||
@@ -139,6 +139,25 @@ var _ = Describe("broadcasts too large for a notification", func() {
|
||||
Expect(row.CreatedAt).ToNot(BeZero())
|
||||
})
|
||||
|
||||
It("delivers nothing for a notification that carries no payload at all", func() {
|
||||
// Publish cannot produce this, because json.Marshal is never empty. It
|
||||
// is reachable if anything else ever notifies on a localai_ channel,
|
||||
// and a handler called with nil is a message that says nothing rather
|
||||
// than no message, which is exactly the confusion this carrier must
|
||||
// never create. The sentinel proves it kept carrying.
|
||||
out := make(chan []byte, 4)
|
||||
_, err := sub.Subscribe(subject, func(b []byte) { out <- b })
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
|
||||
Expect(db.Exec("SELECT pg_notify(?, ?)", "localai_jobs",
|
||||
fmt.Sprintf(`{"s":%q}`, subject)).Error).To(Succeed())
|
||||
Expect(pub.Publish(subject, map[string]string{"m": "sentinel"})).To(Succeed())
|
||||
|
||||
var got []byte
|
||||
Eventually(out, 10*time.Second).Should(Receive(&got))
|
||||
Expect(got).To(MatchJSON(`{"m":"sentinel"}`))
|
||||
})
|
||||
|
||||
It("delivers nothing, and keeps carrying, when a spilled row cannot be found", func() {
|
||||
// A notification whose row is gone is a lost message, not an empty one:
|
||||
// handing a handler nil would let a consumer read a carrier failure as
|
||||
@@ -159,13 +178,31 @@ var _ = Describe("broadcasts too large for a notification", func() {
|
||||
})
|
||||
|
||||
var _ = Describe("retiring spilled rows", func() {
|
||||
var db *gorm.DB
|
||||
var (
|
||||
db *gorm.DB
|
||||
dsn string
|
||||
)
|
||||
|
||||
BeforeEach(func() {
|
||||
db, _ = testutil.SetupTestDBWithDSN()
|
||||
db, dsn = testutil.SetupTestDBWithDSN()
|
||||
Expect(pgbus.Migrate(context.Background(), db)).To(Succeed())
|
||||
})
|
||||
|
||||
// aged writes one spilled row and backdates it on the DATABASE clock, which
|
||||
// is the clock the sweep compares against.
|
||||
aged := func(id string) {
|
||||
GinkgoHelper()
|
||||
Expect(db.Create(&pgbus.BusMessage{ID: id, Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed())
|
||||
Expect(db.Exec("UPDATE bus_messages SET created_at = now() - interval '1 hour' WHERE id = ?", id).Error).To(Succeed())
|
||||
}
|
||||
|
||||
rows := func() []string {
|
||||
GinkgoHelper()
|
||||
var ids []string
|
||||
Expect(db.Model(&pgbus.BusMessage{}).Pluck("id", &ids).Error).To(Succeed())
|
||||
return ids
|
||||
}
|
||||
|
||||
It("leaves the cutoff to the database clock", func() {
|
||||
// Pinned as a statement shape rather than as behaviour on purpose. The
|
||||
// test container shares this host's clock, so a cutoff computed in Go
|
||||
@@ -177,18 +214,29 @@ var _ = Describe("retiring spilled rows", func() {
|
||||
Expect(pgbus.SpillSweepSQL).ToNot(ContainSubstring("created_at < ?"))
|
||||
})
|
||||
|
||||
It("runs the sweep on its own, without anyone asking it to", func() {
|
||||
// SweepSpill and SpillSweepSQL are both spec'd directly, and neither of
|
||||
// them proves the carrier ever CALLS them. Deleting the sweeper's `go`
|
||||
// statement left the whole suite green and bus_messages growing
|
||||
// forever, which is the unpinned-wiring shape this programme exists to
|
||||
// remove. The interval is a Config field so that this spec exists.
|
||||
aged("swept")
|
||||
Expect(db.Create(&pgbus.BusMessage{ID: "kept", Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed())
|
||||
|
||||
b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db, SweepInterval: 20 * time.Millisecond})
|
||||
Expect(err).ToNot(HaveOccurred())
|
||||
DeferCleanup(b.Close)
|
||||
|
||||
Eventually(rows, 30*time.Second).Should(ConsistOf("kept"))
|
||||
})
|
||||
|
||||
It("deletes rows past the retention and keeps the rest", func() {
|
||||
old := pgbus.BusMessage{ID: "old", Subject: "jobs.x", Payload: []byte(`{}`)}
|
||||
fresh := pgbus.BusMessage{ID: "fresh", Subject: "jobs.x", Payload: []byte(`{}`)}
|
||||
Expect(db.Create(&old).Error).To(Succeed())
|
||||
Expect(db.Create(&fresh).Error).To(Succeed())
|
||||
Expect(db.Exec("UPDATE bus_messages SET created_at = now() - interval '1 hour' WHERE id = 'old'").Error).To(Succeed())
|
||||
aged("old")
|
||||
Expect(db.Create(&pgbus.BusMessage{ID: "fresh", Subject: "jobs.x", Payload: []byte(`{}`)}).Error).To(Succeed())
|
||||
|
||||
Expect(pgbus.SweepSpill(context.Background(), db, 5*time.Minute)).To(Succeed())
|
||||
|
||||
var ids []string
|
||||
Expect(db.Model(&pgbus.BusMessage{}).Pluck("id", &ids).Error).To(Succeed())
|
||||
Expect(ids).To(ConsistOf("fresh"))
|
||||
Expect(rows()).To(ConsistOf("fresh"))
|
||||
})
|
||||
})
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ require (
|
||||
github.com/gpustack/gguf-parser-go v0.25.0
|
||||
github.com/hpcloud/tail v1.0.0
|
||||
github.com/ipfs/go-log v1.0.5
|
||||
github.com/jackc/pgx/v5 v5.9.2
|
||||
github.com/jaypipes/ghw v0.24.0
|
||||
github.com/joho/godotenv v1.5.1
|
||||
github.com/klauspost/cpuid/v2 v2.3.0
|
||||
@@ -245,7 +246,6 @@ require (
|
||||
github.com/google/go-querystring v1.1.0 // indirect
|
||||
github.com/jackc/pgpassfile v1.0.0 // indirect
|
||||
github.com/jackc/pgservicefile v0.0.0-20240606120523-5a60cdf6a761 // indirect
|
||||
github.com/jackc/pgx/v5 v5.9.2
|
||||
github.com/jackc/puddle/v2 v2.2.2 // indirect
|
||||
github.com/jbenet/go-context v0.0.0-20150711004518-d14ea06fba99 // indirect
|
||||
github.com/jung-kurt/gofpdf v1.16.2 // indirect
|
||||
|
||||
Reference in new issue
Block a user