diff --git a/core/application/distributed.go b/core/application/distributed.go index 96b67f9ba..c782fa654 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -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. diff --git a/core/application/distributed_test.go b/core/application/distributed_test.go index a822e33ff..6002a9104 100644 --- a/core/application/distributed_test.go +++ b/core/application/distributed_test.go @@ -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 diff --git a/core/services/messaging/match_test.go b/core/services/messaging/match_test.go index 4804768ca..c532dadb7 100644 --- a/core/services/messaging/match_test.go +++ b/core/services/messaging/match_test.go @@ -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()) diff --git a/core/services/pgbus/bus.go b/core/services/pgbus/bus.go index f16d664c2..3f635d0e5 100644 --- a/core/services/pgbus/bus.go +++ b/core/services/pgbus/bus.go @@ -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 { diff --git a/core/services/pgbus/bus_test.go b/core/services/pgbus/bus_test.go index 0585677ee..899536c3b 100644 --- a/core/services/pgbus/bus_test.go +++ b/core/services/pgbus/bus_test.go @@ -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")) - }) -}) diff --git a/core/services/pgbus/listen_ordering_internal_test.go b/core/services/pgbus/listen_ordering_internal_test.go new file mode 100644 index 000000000..c7fc5657a --- /dev/null +++ b/core/services/pgbus/listen_ordering_internal_test.go @@ -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()) + }) +}) diff --git a/core/services/pgbus/spill_test.go b/core/services/pgbus/spill_test.go index 16bafb2e2..0545908e2 100644 --- a/core/services/pgbus/spill_test.go +++ b/core/services/pgbus/spill_test.go @@ -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")) }) }) diff --git a/go.mod b/go.mod index 30bc2f521..b3ccf8dc9 100644 --- a/go.mod +++ b/go.mod @@ -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