From 245010f2f6f7dc124c2eed4b27aa541a63b4870d Mon Sep 17 00:00:00 2001 From: Ettore Di Giacinto Date: Thu, 3 Sep 2026 22:50:54 +0000 Subject: [PATCH] feat(distributed): carry the state.*.delta families on PostgreSQL syncstate.Config held one carrier field typed as the NATS client, so a pgbus.Bus could not be handed to a SyncedMap at all: it satisfies messaging.Broadcaster and not MessagingClient. The durable re-hydration path built for the responses map therefore had a NATS-only consumer and nothing in the build said so. The field becomes Bus messaging.Broadcaster, SubscribeJSON moves to its own file and relaxes its parameter to Broadcaster, and the four adopters fan out over PostgreSQL LISTEN/NOTIFY: fine-tune jobs, quantization jobs, agent tasks with their per-tenant children, and Open Responses metadata. A new spec proves it on a real database, over two Bus instances on two pinned listener connections: a Set and a Delete carry, a payload past the 8000-byte notification cap comes back byte identical through the spill row, two families sharing one LISTEN channel stay separate, and a terminated listener re-hydrates a row written while it was gone. The five sites that each chose a carrier for an adopter are collapsed into one DistributedServices.Broadcast() accessor. Five field reads were five chances to leave one family on NATS with nothing failing, because messaging.Client satisfies Broadcaster too. The accessor also refuses to hand out a nil pgbus.Bus wrapped in a non-nil interface, which every adopter would read as "broadcast" and dereference on the first Set. SetTaskSyncNATS and SetJobSyncNATS are renamed to SetTaskSyncBus and SetJobSyncBus so a missed wiring site fails to compile. The response metadata table gains a retention of its own, defaulting to 24 hours. It inherited the Open Responses store TTL, which defaults to 0 meaning no expiration. Zero is defensible for a map that dies with the process and is not for a table: the table grew for the life of the deployment and a restarting replica re-hydrated every response the cluster had ever created. A row that names its own expiry is still judged on that column alone, and "this row is dead" now has one SQL spelling that PurgeExpired deletes by and ListUnexpired is the negation of, so a hydrate cannot resurrect what a sweep has already retired. Assisted-by: Claude Opus 5 [claude-code] Signed-off-by: Ettore Di Giacinto --- core/application/agent_jobs.go | 7 +- core/application/application.go | 2 +- core/application/distributed.go | 27 ++ core/application/distributed_test.go | 63 ++++ core/application/startup.go | 4 +- core/http/app.go | 27 +- .../http/endpoints/localai/agent_jobs_test.go | 4 +- .../endpoints/openresponses/metadata_store.go | 13 +- core/http/endpoints/openresponses/store.go | 2 +- core/http/endpoints/openresponses/sync.go | 23 +- .../http/endpoints/openresponses/sync_test.go | 27 ++ core/http/routes/openresponses.go | 2 +- core/services/agentpool/agent_jobs.go | 45 +-- core/services/agentpool/task_sync_test.go | 93 +++++- core/services/agentpool/user_services.go | 25 +- core/services/distributed/responses.go | 87 +++++- core/services/distributed/responses_test.go | 122 +++++++- core/services/finetune/service.go | 17 +- core/services/finetune/service_test.go | 32 +- core/services/messaging/client.go | 13 - core/services/messaging/json.go | 26 ++ core/services/quantization/service.go | 17 +- core/services/quantization/service_test.go | 26 +- core/services/syncstate/syncstate.go | 57 +++- core/services/syncstate/syncstate_pg_test.go | 293 ++++++++++++++++++ core/services/syncstate/syncstate_test.go | 57 +++- core/services/testutil/fakebus.go | 20 ++ docs/content/features/distributed-mode.md | 43 ++- .../distributed/syncstate_distributed_test.go | 43 ++- 29 files changed, 1051 insertions(+), 166 deletions(-) create mode 100644 core/services/messaging/json.go create mode 100644 core/services/syncstate/syncstate_pg_test.go diff --git a/core/application/agent_jobs.go b/core/application/agent_jobs.go index f380b0750..f8afe8c02 100644 --- a/core/application/agent_jobs.go +++ b/core/application/agent_jobs.go @@ -37,8 +37,11 @@ func (a *Application) RestartAgentJobService() error { if d.JobStore != nil { agentJobService.SetDistributedJobStore(d.JobStore) } - // Keep agent tasks consistent across replicas (same client the dispatcher uses). - agentJobService.SetTaskSyncNATS(d.Nats) + // Keep agent tasks consistent across replicas, on the deployment's + // broadcast carrier. This is the restart path and it is a second site + // for the same rule: a fix applied only in startup.go leaves every + // service the settings UI restarts on whatever carrier it picked here. + agentJobService.SetTaskSyncBus(d.Broadcast()) } // Start the service diff --git a/core/application/application.go b/core/application/application.go index b49851d47..6a4f99f32 100644 --- a/core/application/application.go +++ b/core/application/application.go @@ -680,7 +680,7 @@ func (a *Application) StartAgentPool() { } // Keep per-user agent tasks consistent across replicas (nil in standalone). if d := a.Distributed(); d != nil { - usm.SetJobSyncNATS(d.Nats) + usm.SetJobSyncBus(d.Broadcast()) } aps.SetUserServicesManager(usm) diff --git a/core/application/distributed.go b/core/application/distributed.go index 08747a451..5f85bc4df 100644 --- a/core/application/distributed.go +++ b/core/application/distributed.go @@ -733,6 +733,33 @@ func requireBroadcastCarrier(ds *DistributedServices) error { return nil } +// Broadcast is the ONE place a wiring site gets the deployment's fan-out +// carrier, and it exists so that "this family travels on the broadcast carrier +// and not on NATS" is decided once instead of at every adopter. +// +// It was five sites before this: the fine-tune service, the quantization +// service, the agent-task setter (twice, on two startup paths), the per-user +// services manager and the Open Responses store. Every one of them takes a +// messaging.Broadcaster, and *messaging.Client satisfies that interface too, so +// a site left holding ds.Nats compiles, starts, publishes and is delivered - +// onto a carrier the deployment is being taken off. Nothing would fail until +// NATS went away. Collapsing the choice to one function makes it a fact a spec +// can pin, which five scattered field reads were not. +// +// The return is the interface and not *pgbus.Bus on purpose: handing a nil +// *pgbus.Bus to an adopter would produce a non-nil interface wrapping a nil +// pointer, and every adopter reads a nil carrier as "standalone, do not +// broadcast". A typed nil would instead panic on the first Set. initDistributed +// already refuses to return a deployment with no carrier (see +// requireBroadcastCarrier), so the nil branch here is belt and braces for a +// zero-valued struct in a test. +func (ds *DistributedServices) Broadcast() messaging.Broadcaster { + if ds == nil || ds.Bus == nil { + return nil + } + return ds.Bus +} + // newBroadcastBus opens the deployment's fan-out carrier on the auth database. // // The DSN is cfg.Auth.DatabaseURL and it may never be anything else. A second diff --git a/core/application/distributed_test.go b/core/application/distributed_test.go index 6002a9104..edcf7115a 100644 --- a/core/application/distributed_test.go +++ b/core/application/distributed_test.go @@ -9,7 +9,9 @@ import ( . "github.com/onsi/gomega" "github.com/mudler/LocalAI/core/config" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/syncstate" "github.com/mudler/LocalAI/core/services/testutil" ) @@ -108,3 +110,64 @@ var _ = Describe("shutting the distributed services down", func() { Expect(bus.IsConnected()).To(BeFalse()) }) }) + +// The one place the four state.*.delta families are told which carrier they +// travel on. +// +// It was five field reads before this: the fine-tune service, the quantization +// service, the agent-task setter on two startup paths, the per-user services +// manager and the Open Responses store. Every one of them takes a +// messaging.Broadcaster, which *messaging.Client satisfies too, so a site left +// holding ds.Nats compiled, started, published and was delivered onto the +// carrier the deployment is being taken off, and nothing failed until NATS did. +// Collapsing the choice into one function is what makes it a fact these specs +// can hold. +var _ = Describe("handing the broadcast carrier to its adopters", func() { + It("returns the carrier the deployment opened and never the NATS client", 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) + + // A NATS client is present on the struct, exactly as it is in a real + // deployment for as long as the request/reply and queue halves survive. + // Identity, not "is a Broadcaster": both fields satisfy that interface. + ds := &DistributedServices{Nats: &messaging.Client{}, Bus: bus} + + Expect(ds.Broadcast()).To(BeIdenticalTo(messaging.Broadcaster(bus))) + }) + + It("returns an interface that reads as absent, not a typed nil, when there is no carrier", func() { + // Every adopter branches on `bus == nil` to mean standalone. A nil + // *pgbus.Bus placed in an interface is NOT nil, so that branch would be + // skipped and the first Set would panic on a request rather than at + // boot. + // + // Compared with == and not with BeNil(). Gomega's BeNil reports a nil + // POINTER inside an interface as nil, so it passes on exactly the value + // this spec exists to reject; the first draft of this spec did, and the + // mutation that removed the guard stayed green. + var ds *DistributedServices + Expect(ds.Broadcast() == nil).To(BeTrue(), "a nil deployment must yield an interface that is itself nil") + Expect((&DistributedServices{}).Broadcast() == nil).To(BeTrue(), + "a deployment with no carrier must yield an interface that is itself nil, not one wrapping a nil *pgbus.Bus") + }) + + It("gives an adopter a carrier-less map rather than one that panics on the first write", func() { + // The consequence, driven through the component every adopter builds. + // A typed nil satisfies `!= nil`, so Start subscribes on it and Set + // publishes on it, and both dereference a nil *pgbus.Bus on a request + // path rather than at boot. + m := syncstate.New(syncstate.Config[string, string]{ + Name: "test.jobs", + Key: func(v string) string { return v }, + Bus: (&DistributedServices{}).Broadcast(), + }) + Expect(m.Start(context.Background())).To(Succeed()) + DeferCleanup(func() { Expect(m.Close()).To(Succeed()) }) + + Expect(func() { Expect(m.Set(context.Background(), "v")).To(Succeed()) }).ToNot(Panic()) + }) +}) diff --git a/core/application/startup.go b/core/application/startup.go index c790a1403..b985d7e55 100644 --- a/core/application/startup.go +++ b/core/application/startup.go @@ -328,8 +328,8 @@ func New(opts ...config.AppOption) (*Application, error) { application.agentJobService.SetDistributedBackends(distSvc.Dispatcher) application.agentJobService.SetDistributedJobStore(distSvc.JobStore) // Keep agent tasks consistent across replicas (jobs already sync via the - // dispatcher + DB read-through). Same NATS client the dispatcher uses. - application.agentJobService.SetTaskSyncNATS(distSvc.Nats) + // dispatcher + DB read-through), on the deployment's broadcast carrier. + application.agentJobService.SetTaskSyncBus(distSvc.Broadcast()) } // Wire skill store into AgentPoolService (wired at pool start time via closure) // The actual wiring happens in StartAgentPool since the pool doesn't exist yet. diff --git a/core/http/app.go b/core/http/app.go index ec2b864ee..e736d4cfd 100644 --- a/core/http/app.go +++ b/core/http/app.go @@ -523,13 +523,15 @@ func API(application *application.Application) (*echo.Echo, error) { routes.RegisterAgentPoolRoutes(e, application, agentsMw, skillsMw, collectionsMw) // Fine-tuning routes fineTuningMw := auth.RequireFeature(application.AuthDB(), auth.FeatureFineTuning) - // In distributed mode pass the shared NATS client + PostgreSQL store so - // fine-tune jobs stay consistent across replicas (the SyncedMap broadcasts - // mutations and hydrates from the DB); standalone passes nil for both. - var ftNats messaging.MessagingClient + // In distributed mode pass the deployment's broadcast carrier + PostgreSQL + // store so fine-tune jobs stay consistent across replicas (the SyncedMap + // broadcasts mutations and hydrates from the DB); standalone passes nil for + // both. The carrier comes from Broadcast() and never from a field read here: + // see the comment on that method for why the choice is made in one place. + var ftBus messaging.Broadcaster var ftStore *distributed.FineTuneStore if d := application.Distributed(); d != nil { - ftNats = d.Nats + ftBus = d.Broadcast() if d.DistStores != nil && d.DistStores.FineTune != nil { ftStore = d.DistStores.FineTune } @@ -538,20 +540,21 @@ func API(application *application.Application) (*echo.Echo, error) { application.ApplicationConfig(), application.ModelLoader(), application.ModelConfigLoader(), - ftNats, + ftBus, ftStore, ) routes.RegisterFineTuningRoutes(e, ftService, application.ApplicationConfig(), application, fineTuningMw) // Quantization routes quantizationMw := auth.RequireFeature(application.AuthDB(), auth.FeatureQuantization) - // In distributed mode pass the shared NATS client + PostgreSQL store so - // quantization jobs stay consistent across replicas (the SyncedMap broadcasts - // mutations and hydrates from the DB); standalone passes nil for both. - var quantNats messaging.MessagingClient + // In distributed mode pass the deployment's broadcast carrier + PostgreSQL + // store so quantization jobs stay consistent across replicas (the SyncedMap + // broadcasts mutations and hydrates from the DB); standalone passes nil for + // both. Same rule and same single source as the fine-tune wiring above. + var quantBus messaging.Broadcaster var quantStore *distributed.QuantStore if d := application.Distributed(); d != nil { - quantNats = d.Nats + quantBus = d.Broadcast() if d.DistStores != nil && d.DistStores.Quant != nil { quantStore = d.DistStores.Quant } @@ -560,7 +563,7 @@ func API(application *application.Application) (*echo.Echo, error) { application.ApplicationConfig(), application.ModelLoader(), application.ModelConfigLoader(), - quantNats, + quantBus, quantStore, ) routes.RegisterQuantizationRoutes(e, qService, application.ApplicationConfig(), application, quantizationMw) diff --git a/core/http/endpoints/localai/agent_jobs_test.go b/core/http/endpoints/localai/agent_jobs_test.go index 9e34eb6af..122bfae14 100644 --- a/core/http/endpoints/localai/agent_jobs_test.go +++ b/core/http/endpoints/localai/agent_jobs_test.go @@ -58,7 +58,7 @@ var _ = Describe("DELETE /api/agent/tasks/:id across tenants", func() { mine := app.AgentJobService() Expect(mine).ToNot(BeNil()) mine.SetUserID("u1") - mine.SetTaskSyncNATS(bus) + mine.SetTaskSyncBus(bus) Expect(mine.LoadTasksFromFile()).To(Succeed()) otherDir := GinkgoT().TempDir() @@ -70,7 +70,7 @@ var _ = Describe("DELETE /api/agent/tasks/:id across tenants", func() { other = agentpool.NewAgentJobServiceWithPaths(otherCfg, nil, nil, nil, filepath.Join(otherDir, "tasks.json"), filepath.Join(otherDir, "jobs.json")) other.SetUserID("u2") - other.SetTaskSyncNATS(bus) + other.SetTaskSyncBus(bus) Expect(other.LoadTasksFromFile()).To(Succeed()) taskID, err = other.CreateTask(schema.Task{Name: "u2 only", Model: "m", Prompt: "p"}) diff --git a/core/http/endpoints/openresponses/metadata_store.go b/core/http/endpoints/openresponses/metadata_store.go index ba4833f37..53afc51da 100644 --- a/core/http/endpoints/openresponses/metadata_store.go +++ b/core/http/endpoints/openresponses/metadata_store.go @@ -52,11 +52,14 @@ func (a *responseMetadataStoreAdapter) List(ctx context.Context) ([]*syncedRespo // purge filter on. // // ExpiresAt is lifted out of the payload into its own column because -// ListUnexpired and PurgeExpired compare against it on the database clock. An -// adapter that left it null would make every response immortal in this table: -// the map would still expire it in memory, so nothing would look broken until -// the table had grown without bound and a restarted replica re-hydrated -// responses that died hours earlier. +// ListUnexpired and PurgeExpired compare against it on the database clock. +// +// It is null whenever the deployment runs the default Open Responses TTL of 0, +// which is the ordinary case and not an error: the store then falls back to +// distributed.DefaultResponseMetadataRetention, so the row is still swept and +// still drops out of a hydrate. What the column buys is the other direction. A +// deployment that DOES configure a TTL gets that TTL honoured here, rather than +// having its responses outlive the map they mirror or die before it. func (a *responseMetadataStoreAdapter) Upsert(ctx context.Context, v *syncedResponse) error { if v == nil { return fmt.Errorf("replicating response metadata: nil value") diff --git a/core/http/endpoints/openresponses/store.go b/core/http/endpoints/openresponses/store.go index ce671b4b9..ba3b4eee6 100644 --- a/core/http/endpoints/openresponses/store.go +++ b/core/http/endpoints/openresponses/store.go @@ -56,7 +56,7 @@ type ResponseStore struct { // (see sync.go), which is how a standalone deployment keeps exactly the // previous process-local behaviour. Guarded by mu. synced *syncstate.SyncedMap[string, *syncedResponse] - nats messaging.MessagingClient + bus messaging.Broadcaster cancelSub messaging.Subscription replicaID string lifeCtx context.Context diff --git a/core/http/endpoints/openresponses/sync.go b/core/http/endpoints/openresponses/sync.go index 261affbd3..09b237e21 100644 --- a/core/http/endpoints/openresponses/sync.go +++ b/core/http/endpoints/openresponses/sync.go @@ -14,7 +14,8 @@ import ( ) // syncStateName is the syncstate namespace for replicated response metadata. -// It becomes the NATS subject "state.responses.metadata.delta". +// It becomes the broadcast subject "state.responses-metadata.delta" (the '.' in +// the name is sanitized to '-', so the subject keeps three tokens). const syncStateName = "responses.metadata" // ErrResponseNotLocal is returned by the stream-resume accessors when the @@ -93,12 +94,12 @@ type responseCancelEvent struct { // and not a deployment shape; tolerating it would silently restore the // deltas-only map this parameter exists to replace, and a non-nil interface // wrapping a nil pointer would instead surface as a panic on a request. -func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.MessagingClient, +func (s *ResponseStore) EnableDistributed(ctx context.Context, bus messaging.Broadcaster, replicaID string, store *distributed.ResponseMetadataStore) error { if store == nil { return errors.New("enabling cross-replica Open Responses: the store parameter is nil, so a reconnecting replica would have nothing to re-hydrate from") } - if nats == nil { + if bus == nil { return nil } @@ -110,7 +111,7 @@ func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.Me synced := syncstate.New(syncstate.Config[string, *syncedResponse]{ Name: syncStateName, Key: func(v *syncedResponse) string { return v.ID }, - Nats: nats, + Bus: bus, Store: &responseMetadataStoreAdapter{store: store}, }) if err := synced.Start(ctx); err != nil { @@ -122,12 +123,12 @@ func (s *ResponseStore) EnableDistributed(ctx context.Context, nats messaging.Me // the store, so everything it reads has to be in place first. s.mu.Lock() s.replicaID = replicaID - s.nats = nats + s.bus = bus s.lifeCtx, s.lifeCancel = lifeCtx, lifeCancel s.synced = synced s.mu.Unlock() - sub, err := messaging.SubscribeJSON(nats, messaging.SubjectResponseCancelWildcard, s.applyRemoteCancel) + sub, err := messaging.SubscribeJSON(bus, messaging.SubjectResponseCancelWildcard, s.applyRemoteCancel) if err != nil { if cerr := s.Close(); cerr != nil { xlog.Warn("failed to tear down response metadata sync after subscribe error", "error", cerr) @@ -245,14 +246,14 @@ func (s *ResponseStore) syncMap() *syncstate.SyncedMap[string, *syncedResponse] // distributed returns the replication handles as a consistent snapshot. Every // path that broadcasts reads them through here so a concurrent Close cannot be // observed half-applied. A nil map means standalone mode. -func (s *ResponseStore) distributed() (*syncstate.SyncedMap[string, *syncedResponse], context.Context, messaging.MessagingClient, string) { +func (s *ResponseStore) distributed() (*syncstate.SyncedMap[string, *syncedResponse], context.Context, messaging.Broadcaster, string) { s.mu.RLock() defer s.mu.RUnlock() ctx := s.lifeCtx if ctx == nil { ctx = context.Background() } - return s.synced, ctx, s.nats, s.replicaID + return s.synced, ctx, s.bus, s.replicaID } // replicaIdentity returns this process's replica ID (empty in standalone mode). @@ -373,9 +374,9 @@ func (s *ResponseStore) delegateCancel(v *syncedResponse) (*schema.ORResponseRes return v.Response, nil } - m, ctx, nats, replicaID := s.distributed() - if nats != nil { - if err := nats.Publish(messaging.SubjectResponseCancel(v.ID), + m, ctx, bus, replicaID := s.distributed() + if bus != nil { + if err := bus.Publish(messaging.SubjectResponseCancel(v.ID), responseCancelEvent{ResponseID: v.ID, Origin: replicaID}); err != nil { xlog.Warn("failed to broadcast Open Responses cancel", "response_id", v.ID, "error", err) } diff --git a/core/http/endpoints/openresponses/sync_test.go b/core/http/endpoints/openresponses/sync_test.go index 2495d1043..608a798c0 100644 --- a/core/http/endpoints/openresponses/sync_test.go +++ b/core/http/endpoints/openresponses/sync_test.go @@ -301,6 +301,33 @@ var _ = Describe("ResponseStore cross-replica", func() { }) Describe("wiring", func() { + It("publishes response metadata on the carrier it was handed, on the responses family's subject", func() { + // S5 in the wiring table. EnableDistributed takes a + // messaging.Broadcaster, and route registration hands it the + // deployment's broadcast carrier; the subject is asserted by name + // because every state.* family shares one LISTEN channel and the + // subject is the only thing that separates them. + const id = "resp_subject" + replicaA.Store(id, &schema.OpenResponsesRequest{Model: "test-model"}, newResponse(id, schema.ORStatusCompleted)) + + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta(syncStateName))).To(BeNumerically(">=", 1)) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(0)) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("agent.tasks"))).To(Equal(0)) + }) + + It("subscribes for delegated cancels on that same carrier", func() { + // The second leg EnableDistributed wires. It rides the same carrier + // as the metadata map, and a cancel that lands on the wrong replica + // reaches nothing without it. + solo := NewResponseStore(0) + own := testutil.NewFakeBus() + Expect(solo.EnableDistributed(ctx, own, "replica-solo", store)).To(Succeed()) + DeferCleanup(func() { Expect(solo.Close()).To(Succeed()) }) + + Expect(own.Subscribers()).To(BeNumerically(">=", 2), + "one subscription for the metadata map and one for the cancel wildcard") + }) + It("refuses to enable replication without a durable store", func() { // A nil store here is a wiring bug, not a deployment shape: this is // reached only from the distributed branch of route registration. diff --git a/core/http/routes/openresponses.go b/core/http/routes/openresponses.go index 1d3977d28..4cf243864 100644 --- a/core/http/routes/openresponses.go +++ b/core/http/routes/openresponses.go @@ -33,7 +33,7 @@ func RegisterOpenResponsesRoutes(app *echo.Echo, responseStore = d.DistStores.Responses } if err := openresponses.GetGlobalStore().EnableDistributed( - application.ApplicationConfig().Context, d.Nats, application.InstanceID(), responseStore); err != nil { + application.ApplicationConfig().Context, d.Broadcast(), application.InstanceID(), responseStore); err != nil { xlog.Error("Failed to enable cross-replica Open Responses store", "error", err) } } diff --git a/core/services/agentpool/agent_jobs.go b/core/services/agentpool/agent_jobs.go index 2ba9f9a90..8c310a75e 100644 --- a/core/services/agentpool/agent_jobs.go +++ b/core/services/agentpool/agent_jobs.go @@ -46,15 +46,17 @@ type AgentJobService struct { evaluator *templates.Evaluator // tasks is the cross-replica task store: an in-memory map kept consistent - // across replicas via NATS, with read-through to the configured persister + // across replicas over the deployment's fan-out carrier, with read-through + // to the configured persister // (file in standalone, PostgreSQL in distributed). Unlike jobs - which already // converge via the dispatcher + DB read-through - tasks previously read // in-memory only, so ListTasks went stale on non-originating replicas. tasks *syncstate.SyncedMap[string, schema.Task] - // taskNats is the distributed NATS client backing the tasks SyncedMap. It is - // not available at construction time, so it is injected via SetTaskSyncNATS - // during distributed wiring; nil keeps tasks in-memory-only (standalone). - taskNats messaging.MessagingClient + // taskBus is the deployment's broadcast carrier backing the tasks SyncedMap. + // It is not available at construction time, so it is injected via + // SetTaskSyncBus during distributed wiring; nil keeps tasks in-memory-only + // (standalone). + taskBus messaging.Broadcaster // Storage (in-memory primary, persister for secondary persistence) jobs *xsync.SyncedMap[string, schema.Job] @@ -101,8 +103,8 @@ func (s *AgentJobService) SetDistributedBackends(dispatcher DistributedDispatche // // The rebuild is what makes the two setters order-independent. Without it the // map keeps whichever tenant it was built with, so wiring that happened to call -// SetTaskSyncNATS first would publish this user's tasks on the CLUSTER-WIDE -// subject and every other tenant would apply them. Like SetTaskSyncNATS, this +// SetTaskSyncBus first would publish this user's tasks on the CLUSTER-WIDE +// subject and every other tenant would apply them. Like SetTaskSyncBus, this // is only ever called before Start / hydrate, while the map is still empty, so // rebuilding loses no state. func (s *AgentJobService) SetUserID(id string) { @@ -117,27 +119,32 @@ func (s *AgentJobService) SetDistributedJobStore(store *jobs.JobStore) { s.persister = &dbJobPersister{store: store} } -// SetTaskSyncNATS wires the distributed NATS client used to keep agent *tasks* -// consistent across replicas (jobs already converge via the dispatcher + DB -// read-through, so they are left untouched). The client is not available when the -// service is constructed, so it is injected here during distributed wiring and the -// tasks SyncedMap is rebuilt to pick it up. It is always called before Start / +// SetTaskSyncBus wires the deployment's broadcast carrier used to keep agent +// *tasks* consistent across replicas (jobs already converge via the dispatcher + +// DB read-through, so they are left untouched). The carrier is not available when +// the service is constructed, so it is injected here during distributed wiring and +// the tasks SyncedMap is rebuilt to pick it up. It is always called before Start / // hydrate, while the map is still empty, so rebuilding loses no state. Passing nil // (standalone) keeps the map in-memory-only with no broadcast. -func (s *AgentJobService) SetTaskSyncNATS(nats messaging.MessagingClient) { - s.taskNats = nats +// +// The parameter is messaging.Broadcaster, so in distributed mode this family +// travels on PostgreSQL LISTEN/NOTIFY. The name says Bus and not NATS because +// the two are no longer the same thing and a stale name here would be the only +// documentation a wiring site reads. +func (s *AgentJobService) SetTaskSyncBus(bus messaging.Broadcaster) { + s.taskBus = bus s.buildTasksMap() } // buildTasksMap (re)constructs the cross-replica tasks SyncedMap from the current -// taskNats. The Store adapter reads s.persister/s.userID live, so a persister swap -// (SetDistributedJobStore) needs no rebuild; only the NATS client, fixed at -// New-time, forces one - hence SetTaskSyncNATS calls this. +// taskBus. The Store adapter reads s.persister/s.userID live, so a persister swap +// (SetDistributedJobStore) needs no rebuild; only the carrier, fixed at map-build +// time, forces one - hence SetTaskSyncBus calls this. func (s *AgentJobService) buildTasksMap() { s.tasks = syncstate.New(syncstate.Config[string, schema.Task]{ Name: "agent.tasks", Key: func(t schema.Task) string { return t.ID }, - Nats: s.taskNats, + Bus: s.taskBus, Store: &taskStoreAdapter{svc: s}, // There is one AgentJobService per user, so this map is per-tenant and // its deltas must not reach another tenant's copy. The empty userID is @@ -276,7 +283,7 @@ func NewAgentJobServiceWithPaths( cronEntries: xsync.NewSyncedMap[string, cron.EntryID](), retentionDays: retentionDays, } - // Build the cross-replica tasks map standalone (nil NATS); SetTaskSyncNATS + // Build the cross-replica tasks map standalone (nil carrier); SetTaskSyncBus // rebuilds it with the distributed client once that is available, before Start. s.buildTasksMap() return s diff --git a/core/services/agentpool/task_sync_test.go b/core/services/agentpool/task_sync_test.go index dc65c4097..15bb4c92f 100644 --- a/core/services/agentpool/task_sync_test.go +++ b/core/services/agentpool/task_sync_test.go @@ -26,7 +26,7 @@ import ( // newTaskSyncService builds an AgentJobService wired to the given bus and a // throwaway data dir (so the file persister has somewhere to write). Model/config // loaders are nil because the task sync paths under test never touch them. -func newTaskSyncService(bus messaging.MessagingClient) *AgentJobService { +func newTaskSyncService(bus messaging.Broadcaster) *AgentJobService { tmpDir := GinkgoT().TempDir() sysState := &system.SystemState{} sysState.Model.ModelsPath = tmpDir @@ -40,7 +40,7 @@ func newTaskSyncService(bus messaging.MessagingClient) *AgentJobService { // Distinct per-replica files so the file persister write-through never // crosses replicas: convergence here must be proven via the bus alone. tmpDir+"/tasks.json", tmpDir+"/jobs.json") - svc.SetTaskSyncNATS(bus) + svc.SetTaskSyncBus(bus) return svc } @@ -168,7 +168,7 @@ func newTaskSyncServiceForTenant(bus messaging.MessagingClient, userID string) * svc := NewAgentJobServiceWithPaths(appConfig, nil, nil, nil, tmpDir+"/tasks.json", tmpDir+"/jobs.json") svc.SetUserID(userID) - svc.SetTaskSyncNATS(bus) + svc.SetTaskSyncBus(bus) return svc } @@ -293,9 +293,41 @@ var _ = Describe("AgentJobService same-tenant replicas", func() { Expect(got.Name).To(Equal("Shared")) }) + It("publishes the cluster-wide view on the carrier SetTaskSyncBus was handed", func() { + // S3 in the wiring table. Two startup paths call this one setter + // (startup.go and the settings-driven restart in agent_jobs.go), so the + // service must broadcast on whatever it was given and on the agent-task + // family's own subject, never on another family's. + bus := testutil.NewFakeBus() + svc := newTaskSyncService(bus) + Expect(svc.Start(context.Background())).To(Succeed()) + defer func() { Expect(svc.Stop()).To(Succeed()) }() + + _, err := svc.CreateTask(schema.Task{Name: "Cluster", Model: "m", Prompt: "p"}) + Expect(err).NotTo(HaveOccurred()) + + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("agent.tasks"))).To(Equal(1)) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(0)) + }) + + It("broadcasts nothing when it is handed no carrier at all", func() { + // The standalone half of S3: a single-binary deployment passes nil and + // must get a map that still works and never reaches for a carrier. + bus := testutil.NewFakeBus() + svc := newTaskSyncService(nil) + Expect(svc.Start(context.Background())).To(Succeed()) + defer func() { Expect(svc.Stop()).To(Succeed()) }() + + id, err := svc.CreateTask(schema.Task{Name: "Solo", Model: "m", Prompt: "p"}) + Expect(err).NotTo(HaveOccurred()) + _, err = svc.GetTask(id) + Expect(err).NotTo(HaveOccurred(), "a carrier-less service must still serve its own reads") + Expect(bus.Subscribers()).To(Equal(0)) + }) + It("scopes the subject whichever order the setters are called in", func() { // UserServicesManager.GetJobs happens to call SetUserID before - // SetTaskSyncNATS. Nothing enforces that order, and a map built from a + // SetTaskSyncBus. Nothing enforces that order, and a map built from a // user id that had not been set yet publishes cluster-wide - the leak, // reintroduced by a line move. bus := testutil.NewFakeBus() @@ -311,3 +343,56 @@ var _ = Describe("AgentJobService same-tenant replicas", func() { Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("agent.tasks"))).To(Equal(0)) }) }) + +// The per-user manager is the wiring site that fails silently. +// +// S4 in the table, and it is the reason the table has five rows rather than +// four: the global AgentJobService and every per-user one are different +// objects, so a carrier handed only to the global service leaves every tenant's +// tasks unreplicated with nothing failing anywhere. The manager stores the +// carrier once and every service it builds afterwards inherits it. +var _ = Describe("UserServicesManager task carrier propagation", func() { + newManager := func() *UserServicesManager { + GinkgoHelper() + tmpDir := GinkgoT().TempDir() + sysState := &system.SystemState{} + sysState.Model.ModelsPath = tmpDir + appConfig := config.NewApplicationConfig( + config.WithDynamicConfigDir(tmpDir), + config.WithContext(context.Background()), + ) + appConfig.SystemState = sysState + return NewUserServicesManager(NewUserScopedStorage(tmpDir, tmpDir), appConfig, nil, nil, nil) + } + + It("hands the carrier it was given to each per-user service, on that tenant's subject", func() { + bus := testutil.NewFakeBus() + m := newManager() + m.SetJobSyncBus(bus) + + svc, err := m.GetJobs("u1") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(svc.Stop()).To(Succeed()) }) + + _, err = svc.CreateTask(schema.Task{Name: "Tenant", Model: "m", Prompt: "p"}) + Expect(err).NotTo(HaveOccurred()) + + Expect(bus.PublishCount(messaging.SubjectSyncStateTenantDelta("agent.tasks", "u1"))).To(Equal(1), + "a per-user service that inherited no carrier leaves that tenant unreplicated, silently") + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("agent.tasks"))).To(Equal(0), + "and one that inherited it must not publish a tenant's tasks cluster-wide") + }) + + It("leaves per-user services carrier-less when the manager was given none", func() { + bus := testutil.NewFakeBus() + m := newManager() + + svc, err := m.GetJobs("u2") + Expect(err).NotTo(HaveOccurred()) + DeferCleanup(func() { Expect(svc.Stop()).To(Succeed()) }) + + _, err = svc.CreateTask(schema.Task{Name: "Tenant", Model: "m", Prompt: "p"}) + Expect(err).NotTo(HaveOccurred()) + Expect(bus.Subscribers()).To(Equal(0)) + }) +}) diff --git a/core/services/agentpool/user_services.go b/core/services/agentpool/user_services.go index 56d19e0fc..c0cbcfef7 100644 --- a/core/services/agentpool/user_services.go +++ b/core/services/agentpool/user_services.go @@ -29,9 +29,9 @@ type UserServicesManager struct { // Shared distributed backends (set once, inherited by per-user job services) jobDispatcher DistributedDispatcher jobDBStore *jobs.JobStore - // jobNats keeps per-user agent tasks consistent across replicas (nil in + // jobBus keeps per-user agent tasks consistent across replicas (nil in // standalone). Inherited by each per-user AgentJobService. - jobNats messaging.MessagingClient + jobBus messaging.Broadcaster } // NewUserServicesManager creates a new UserServicesManager. @@ -166,10 +166,14 @@ func (m *UserServicesManager) GetJobs(userID string) (*AgentJobService, error) { if m.jobDispatcher != nil { svc.SetDistributedBackends(m.jobDispatcher) } - // Inherit the NATS client so per-user tasks broadcast across replicas. Must be - // set before the hydrate below (LoadFromDB / LoadTasksFromFile) so the tasks - // SyncedMap is rebuilt with the client while it is still empty. - svc.SetTaskSyncNATS(m.jobNats) + // Inherit the broadcast carrier so per-user tasks fan out across replicas. + // Must be set before the hydrate below (LoadFromDB / LoadTasksFromFile) so the + // tasks SyncedMap is rebuilt with the carrier while it is still empty. + // + // This is a second wiring site for the same rule, and it is the one that is + // invisible: fixing the global service alone leaves every tenant's map on + // whatever carrier this manager was handed, with nothing failing. + svc.SetTaskSyncBus(m.jobBus) if m.jobDBStore != nil { svc.SetDistributedJobStore(m.jobDBStore) // Load tasks/jobs from DB immediately (per-user services skip Start()) @@ -197,10 +201,11 @@ func (m *UserServicesManager) SetJobDBStore(s *jobs.JobStore) { m.jobDBStore = s } -// SetJobSyncNATS sets the NATS client used to keep per-user agent tasks consistent -// across replicas. -func (m *UserServicesManager) SetJobSyncNATS(nats messaging.MessagingClient) { - m.jobNats = nats +// SetJobSyncBus sets the broadcast carrier used to keep per-user agent tasks +// consistent across replicas. Every per-user service built afterwards inherits +// it; see the call in the builder above. +func (m *UserServicesManager) SetJobSyncBus(bus messaging.Broadcaster) { + m.jobBus = bus } // ListAllUserIDs returns all user IDs that have scoped data directories. diff --git a/core/services/distributed/responses.go b/core/services/distributed/responses.go index 5d93d11de..ede01803b 100644 --- a/core/services/distributed/responses.go +++ b/core/services/distributed/responses.go @@ -47,8 +47,52 @@ func (ResponseMetadataRecord) TableName() string { return "response_metadata" } // confused them would blank the map on a transient outage. type ResponseMetadataStore struct { db *gorm.DB + + // retention bounds how long a row that carries NO expiry of its own stays + // in this table. See DefaultResponseMetadataRetention. + retention time.Duration } +// DefaultResponseMetadataRetention is how long a response row with no expiry of +// its own is kept. +// +// It is INDEPENDENT of the Open Responses store TTL, and that is the whole +// point. That TTL defaults to 0, documented as "no expiration", and 0 there is +// defensible: the in-memory map it governs dies with the process, so unbounded +// means "bounded by this process's lifetime and by memory an operator can see". +// A TABLE has neither bound. Inheriting that 0 here made every replicated +// response immortal on disk, so the table grew for the life of the deployment +// and a restarting replica re-hydrated its map with every response the cluster +// had ever created. Neither is a condition anything reports. +// +// So the durable projection gets a retention of its own. It is a floor on +// CROSS-REPLICA visibility and never on the response: the replica that owns a +// response keeps it in memory for exactly as long as the configured TTL says, +// and an operator who sets a TTL longer than this gets that TTL honoured, +// because a row with an expires_at is judged on that column alone. +// +// Twenty-four hours is chosen against what the row is for. It exists so a peer +// replica, or a restarted one, can answer a GET or resolve a +// previous_response_id it did not create. That is a lookup a client makes +// within minutes of the response, not a day later, and the row is metadata +// rather than the generation itself. +const DefaultResponseMetadataRetention = 24 * time.Hour + +// deadResponseMetadata is the ONE definition of "this row is no longer live", +// evaluated on the DATABASE clock so every replica agrees. +// +// PurgeExpired deletes exactly the rows it matches and ListUnexpired returns +// exactly the rows it does not, spelled as NOT of this same string. Two +// separate spellings would let a hydrate resurrect a row a sweep had already +// decided was dead, or leave a row that is invisible to every reader sitting in +// the table forever, and neither reports anything. +// +// The bind parameter is the retention in seconds, and it is the only one: +// make_interval(secs => ?) turns it into an interval the SERVER subtracts from +// its own now(), so no process's clock enters the comparison. +const deadResponseMetadata = `((expires_at IS NOT NULL AND expires_at <= now()) OR ` + + `(expires_at IS NULL AND created_at <= now() - make_interval(secs => ?)))` + // NewResponseMetadataStore creates a ResponseMetadataStore and migrates its // table. // @@ -60,9 +104,23 @@ type ResponseMetadataStore struct { // The migration runs under the same advisory lock NewFineTuneStore uses, because // several replicas start at once and concurrent AutoMigrate races. func NewResponseMetadataStore(db *gorm.DB) (*ResponseMetadataStore, error) { + return NewResponseMetadataStoreWithRetention(db, DefaultResponseMetadataRetention) +} + +// NewResponseMetadataStoreWithRetention is NewResponseMetadataStore with an +// explicit bound on rows that carry no expiry of their own. +// +// A non-positive retention is refused rather than taken to mean "keep +// everything": unbounded is the shape this store had, it grows a table nothing +// ever sweeps, and a deployment that wants a different bound should say which +// one rather than switch the bound off. +func NewResponseMetadataStoreWithRetention(db *gorm.DB, retention time.Duration) (*ResponseMetadataStore, error) { if db == nil { return nil, fmt.Errorf("response metadata store: no database handle") } + if retention <= 0 { + return nil, fmt.Errorf("response metadata store: retention must be positive, got %s; a row with no expiry of its own would never be swept and the table would grow for the life of the deployment", retention) + } if name := db.Dialector.Name(); name != "postgres" { return nil, fmt.Errorf("response metadata store requires PostgreSQL, this deployment runs on %q", name) } @@ -71,9 +129,13 @@ func NewResponseMetadataStore(db *gorm.DB) (*ResponseMetadataStore, error) { }); err != nil { return nil, fmt.Errorf("migrating response_metadata: %w", err) } - return &ResponseMetadataStore{db: db}, nil + return &ResponseMetadataStore{db: db, retention: retention}, nil } +// Retention is the effective bound on rows with no expiry of their own. +// Reported so a deployment can log what it actually got. +func (s *ResponseMetadataStore) Retention() time.Duration { return s.retention } + // Upsert idempotently inserts or replaces one row by primary key. // // created_at is deliberately NOT in the update set: the row is rewritten on @@ -100,19 +162,20 @@ func (s *ResponseMetadataStore) Delete(ctx context.Context, id string) error { return s.db.WithContext(ctx).Where("id = ?", id).Delete(&ResponseMetadataRecord{}).Error } -// ListUnexpired returns every row whose ExpiresAt is null or still in the -// future, measured on the DATABASE clock. +// ListUnexpired returns every row that is still live: its ExpiresAt is in the +// future, or it has none and is younger than the retention. Both legs are +// measured on the DATABASE clock. // // The clock is the database's because every replica hydrating from this table // must agree on which rows are live, and a Go-side cutoff makes that a property -// of whichever process asked. It is spelled `expires_at IS NULL OR expires_at > -// now()` and it is dialect-guarded in the constructor, because now() on the -// SQLite single-binary path reads as a missing migration rather than as an -// error. +// of whichever process asked. The predicate is deadResponseMetadata negated, so +// this and PurgeExpired cannot disagree about a row; it is dialect-guarded in +// the constructor, because now() and make_interval on the SQLite single-binary +// path read as a missing migration rather than as an error. func (s *ResponseMetadataStore) ListUnexpired(ctx context.Context) ([]ResponseMetadataRecord, error) { var out []ResponseMetadataRecord if err := s.db.WithContext(ctx). - Where("expires_at IS NULL OR expires_at > now()"). + Where("NOT "+deadResponseMetadata, s.retention.Seconds()). Order("created_at"). Find(&out).Error; err != nil { return nil, fmt.Errorf("listing unexpired response metadata: %w", err) @@ -120,12 +183,12 @@ func (s *ResponseMetadataStore) ListUnexpired(ctx context.Context) ([]ResponseMe return out, nil } -// PurgeExpired deletes rows whose ExpiresAt has passed and returns how many. -// It is the reason a table of ephemeral state does not grow forever, and it -// runs on the same DATABASE clock as ListUnexpired. +// PurgeExpired deletes every row ListUnexpired would refuse to return and +// reports how many. It is the reason a table of ephemeral state does not grow +// forever, and it runs on the same DATABASE clock and the same predicate. func (s *ResponseMetadataStore) PurgeExpired(ctx context.Context) (int64, error) { res := s.db.WithContext(ctx). - Where("expires_at IS NOT NULL AND expires_at <= now()"). + Where(deadResponseMetadata, s.retention.Seconds()). Delete(&ResponseMetadataRecord{}) if res.Error != nil { return 0, fmt.Errorf("purging expired response metadata: %w", res.Error) diff --git a/core/services/distributed/responses_test.go b/core/services/distributed/responses_test.go index ffe77c1b0..f529959ae 100644 --- a/core/services/distributed/responses_test.go +++ b/core/services/distributed/responses_test.go @@ -240,6 +240,89 @@ var _ = Describe("ResponseMetadataStore", func() { }) }) + // The bound on rows that carry no expiry of their own. + // + // The Open Responses store TTL defaults to 0, documented as "no + // expiration", and every response then reaches this table with a null + // expires_at. Zero is defensible for the in-memory map it governs, which + // dies with the process; a table has no such bound, so it grew for the life + // of the deployment and a restarting replica re-hydrated its map with every + // response the cluster had ever created. Neither is a condition anything + // reports, which is why it needed a default of its own rather than a note. + Describe("retention for rows with no expiry of their own", func() { + aged := func(id string, age time.Duration, expiresAt *time.Time) *distributed.ResponseMetadataRecord { + r := newRecord(id, expiresAt) + r.CreatedAt = time.Now().Add(-age) + return r + } + + It("stops listing and then sweeps an untagged row older than the retention", func() { + Expect(store.Upsert(ctx, aged("resp_stale", distributed.DefaultResponseMetadataRetention+time.Hour, nil))).To(Succeed()) + Expect(store.Upsert(ctx, aged("resp_recent", time.Minute, nil))).To(Succeed()) + + recs, err := store.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(ids(recs)).To(ConsistOf("resp_recent"), + "a row past the retention must not re-hydrate a replica's map") + + n, err := store.PurgeExpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(Equal(int64(1)), "and the sweep must actually retire it, or the table still grows forever") + + recs, err = store.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(ids(recs)).To(ConsistOf("resp_recent")) + }) + + It("honours a row's own expiry over the retention, in both directions", func() { + // The retention is a floor on rows that named no expiry, never a + // ceiling on rows that did. A deployment that configured a longer + // TTL must get it, and one that configured a shorter one must not + // have its responses kept alive by this default. + far := time.Now().Add(distributed.DefaultResponseMetadataRetention * 10) + past := time.Now().Add(-time.Minute) + + Expect(store.Upsert(ctx, aged("resp_long_ttl", distributed.DefaultResponseMetadataRetention+time.Hour, &far))).To(Succeed()) + Expect(store.Upsert(ctx, aged("resp_short_ttl", time.Minute, &past))).To(Succeed()) + + recs, err := store.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(ids(recs)).To(ConsistOf("resp_long_ttl")) + + n, err := store.PurgeExpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(n).To(Equal(int64(1))) + }) + + It("takes an explicit retention and reports the one it is using", func() { + short, err := distributed.NewResponseMetadataStoreWithRetention(db, time.Hour) + Expect(err).ToNot(HaveOccurred()) + Expect(short.Retention()).To(Equal(time.Hour)) + Expect(store.Retention()).To(Equal(distributed.DefaultResponseMetadataRetention)) + + Expect(short.Upsert(ctx, aged("resp_two_hours", 2*time.Hour, nil))).To(Succeed()) + + recs, err := short.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(recs).To(BeEmpty(), "the configured retention, not the default, decides") + + // The same row is still live to a store with the default retention, + // which is what proves the bound is the store's and not the row's. + recs, err = store.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + Expect(ids(recs)).To(ConsistOf("resp_two_hours")) + }) + + It("refuses a non-positive retention rather than restoring an unbounded table", func() { + _, err := distributed.NewResponseMetadataStoreWithRetention(db, 0) + Expect(err).To(HaveOccurred()) + Expect(err.Error()).To(ContainSubstring("grow")) + + _, err = distributed.NewResponseMetadataStoreWithRetention(db, -time.Hour) + Expect(err).To(HaveOccurred()) + }) + }) + // The expiry cutoff is the database's clock, not the asking process's. The // container shares this host's clock, so no behavioural spec above can tell // a Go-side time.Now() bind parameter from now(); these pin the statement @@ -260,11 +343,11 @@ var _ = Describe("ResponseMetadataStore", func() { Expect(err).ToNot(HaveOccurred()) sql := rec.only() - Expect(sql).To(MatchRegexp(`(?i)expires_at\s+IS\s+NULL\s+OR\s+expires_at\s*>\s*now\(\)`)) + Expect(sql).To(MatchRegexp(`(?i)expires_at\s+IS\s+NOT\s+NULL\s+AND\s+expires_at\s*<=\s*now\(\)`)) // A bound parameter in the comparison position is a Go-side cutoff // wearing the same behaviour; gorm's logger explains binds into the // text, so a time literal here is exactly that defect. - Expect(sql).ToNot(MatchRegexp(`expires_at\s*>\s*['$]`)) + Expect(sql).ToNot(MatchRegexp(`expires_at\s*<=\s*['$]`)) }) It("compares expires_at against the database clock in PurgeExpired", func() { @@ -276,6 +359,41 @@ var _ = Describe("ResponseMetadataStore", func() { Expect(sql).ToNot(MatchRegexp(`expires_at\s*<=\s*['$]`)) }) + It("ages an untagged row out against the database clock too, never a Go-side cutoff", func() { + // The retention leg is subtracted from the SERVER's now() with + // make_interval, so the only bind is a number of seconds. A + // timestamp literal here would be this process's clock deciding + // which rows the whole deployment can still see. + _, err := store.PurgeExpired(ctx) + Expect(err).ToNot(HaveOccurred()) + + sql := rec.only() + Expect(sql).To(MatchRegexp(`(?i)created_at\s*<=\s*now\(\)\s*-\s*make_interval\(secs\s*=>`)) + Expect(sql).ToNot(MatchRegexp(`created_at\s*<=\s*['$]`)) + }) + + It("selects and deletes on ONE predicate, so a hydrate cannot resurrect what a sweep killed", func() { + // Two spellings of "this row is dead" drift, and the drift is + // silent both ways: a row invisible to every reader that nothing + // deletes, or a row a sweep removed that a hydrate had just served. + _, err := store.ListUnexpired(ctx) + Expect(err).ToNot(HaveOccurred()) + listSQL := rec.only() + + rec.reset() + _, err = store.PurgeExpired(ctx) + Expect(err).ToNot(HaveOccurred()) + purgeSQL := rec.only() + + predicate := regexp.MustCompile(`(?is)\(\(expires_at.*?make_interval\(secs\s*=>[^)]*\)\)\)`) + listPredicate := predicate.FindString(listSQL) + Expect(listPredicate).ToNot(BeEmpty(), "ListUnexpired must carry the shared predicate") + Expect(purgeSQL).To(ContainSubstring(listPredicate), + "PurgeExpired must delete exactly what ListUnexpired refuses to return") + Expect(listSQL).To(ContainSubstring("NOT "+listPredicate), + "and ListUnexpired must be its negation rather than a second spelling") + }) + It("issues one statement per call, so neither reads the clock into Go first", func() { _, err := store.ListUnexpired(ctx) Expect(err).ToNot(HaveOccurred()) diff --git a/core/services/finetune/service.go b/core/services/finetune/service.go index 3e2431df2..8e3978a00 100644 --- a/core/services/finetune/service.go +++ b/core/services/finetune/service.go @@ -40,19 +40,24 @@ type FineTuneService struct { mu sync.Mutex // jobs is the cross-replica job store: an in-memory map kept consistent across - // replicas via NATS, optionally read-through to PostgreSQL in distributed mode. + // replicas over the deployment's fan-out carrier, optionally read-through to + // PostgreSQL in distributed mode. jobs *syncstate.SyncedMap[string, *schema.FineTuneJob] } // NewFineTuneService creates a new FineTuneService. In distributed mode pass the -// shared NATS client and PostgreSQL store so jobs stay consistent across -// replicas; pass nil for both in standalone mode, where the disk Loader hydrates -// the map and there is nothing to broadcast. +// deployment's broadcast carrier and PostgreSQL store so jobs stay consistent +// across replicas; pass nil for both in standalone mode, where the disk Loader +// hydrates the map and there is nothing to broadcast. +// +// bus is messaging.Broadcaster and not the NATS client: this state.*.delta +// family travels on whatever the deployment's fan-out carrier is, and in +// distributed mode that is PostgreSQL LISTEN/NOTIFY. func NewFineTuneService( appConfig *config.ApplicationConfig, modelLoader *model.ModelLoader, configLoader *config.ModelConfigLoader, - nats messaging.MessagingClient, + bus messaging.Broadcaster, store *distributed.FineTuneStore, ) *FineTuneService { s := &FineTuneService{ @@ -72,7 +77,7 @@ func NewFineTuneService( s.jobs = syncstate.New(syncstate.Config[string, *schema.FineTuneJob]{ Name: "finetune.jobs", Key: func(j *schema.FineTuneJob) string { return j.ID }, - Nats: nats, + Bus: bus, Store: syncStore, Loader: s.loadJobsFromDisk, // ignored when Store is set (distributed mode) }) diff --git a/core/services/finetune/service_test.go b/core/services/finetune/service_test.go index dc7c53290..5c40daabb 100644 --- a/core/services/finetune/service_test.go +++ b/core/services/finetune/service_test.go @@ -14,13 +14,14 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/distributed" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" ) // newTestService builds a standalone FineTuneService wired to the given bus. The // model/config loaders are nil because the read/sync paths under test never touch // them; the data dir is a throwaway temp dir so the disk Loader finds nothing. -func newTestService(bus *testutil.FakeBus) *FineTuneService { +func newTestService(bus messaging.Broadcaster) *FineTuneService { appConfig := &config.ApplicationConfig{ Context: context.Background(), DataPath: GinkgoT().TempDir(), @@ -77,6 +78,35 @@ var _ = Describe("FineTuneService", func() { Expect(err).To(HaveOccurred(), "a delete on A must remove the job from B") }) + It("publishes on the carrier it was handed, and on the fine-tune family's own subject", func() { + // S1 in the wiring table: this service takes a + // messaging.Broadcaster and must broadcast on THAT and on nothing + // else. The subject is asserted by name rather than inferred from + // delivery, because a service that published on the wrong family + // would still reach a peer of its own kind here while colliding + // with quantization in production. + job := &schema.FineTuneJob{ID: "job-subject", UserID: "user-1", Status: "queued", CreatedAt: "2026-06-27T10:00:00Z"} + Expect(a.jobs.Set(ctx, job)).To(Succeed()) + + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(1)) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("quant.jobs"))).To(Equal(0)) + }) + + It("broadcasts nothing at all when it is built with no carrier", func() { + // Standalone is the other half of the same wiring rule: a nil + // carrier must be a strict no-op rather than a panic on the first + // job a single-binary user starts. + solo := newTestService(nil) + DeferCleanup(func() { Expect(solo.Close()).To(Succeed()) }) + + Expect(solo.jobs.Set(ctx, &schema.FineTuneJob{ID: "solo", UserID: "u", CreatedAt: "2026-06-27T10:00:00Z"})).To(Succeed()) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(0)) + + got, err := solo.GetJob("u", "solo") + Expect(err).ToNot(HaveOccurred(), "a carrier-less service must still serve its own reads") + Expect(got.ID).To(Equal("solo")) + }) + It("propagates a status update from A to B", func() { job := &schema.FineTuneJob{ID: "job-3", UserID: "user-1", Status: "training", CreatedAt: "2026-06-27T10:00:00Z"} Expect(a.jobs.Set(ctx, job)).To(Succeed()) diff --git a/core/services/messaging/client.go b/core/services/messaging/client.go index e01c7d9ca..7450d8b17 100644 --- a/core/services/messaging/client.go +++ b/core/services/messaging/client.go @@ -263,19 +263,6 @@ func (c *Client) QueueSubscribeReply(subject, queue string, handler func(data [] }) } -// SubscribeJSON creates a subscription that automatically unmarshals JSON messages. -// Invalid JSON messages are logged and skipped. -func SubscribeJSON[T any](c MessagingClient, subject string, handler func(T)) (Subscription, error) { - return c.Subscribe(subject, func(data []byte) { - var evt T - if err := json.Unmarshal(data, &evt); err != nil { - xlog.Warn("Failed to unmarshal NATS message", "subject", subject, "error", err) - return - } - handler(evt) - }) -} - // QueueSubscribeJSON creates a queue subscription that automatically unmarshals JSON messages. // Invalid JSON messages are logged and skipped. func QueueSubscribeJSON[T any](c MessagingClient, subject, queue string, handler func(T)) (Subscription, error) { diff --git a/core/services/messaging/json.go b/core/services/messaging/json.go new file mode 100644 index 000000000..ccc27b531 --- /dev/null +++ b/core/services/messaging/json.go @@ -0,0 +1,26 @@ +package messaging + +import ( + "encoding/json" + + "github.com/mudler/xlog" +) + +// SubscribeJSON creates a subscription that automatically unmarshals JSON +// messages. Invalid JSON messages are logged and skipped. +// +// Its parameter is Broadcaster and not MessagingClient, so a call site can be +// moved onto the PostgreSQL carrier without also being rewritten. It lives in +// its own file because client.go is deleted once the request/reply and queue +// halves are retired, and a generic helper with this many call sites must not +// go with it. +func SubscribeJSON[T any](c Broadcaster, subject string, handler func(T)) (Subscription, error) { + return c.Subscribe(subject, func(data []byte) { + var evt T + if err := json.Unmarshal(data, &evt); err != nil { + xlog.Warn("Failed to unmarshal a broadcast message", "subject", subject, "error", err) + return + } + handler(evt) + }) +} diff --git a/core/services/quantization/service.go b/core/services/quantization/service.go index cd9cbcead..c51bd0021 100644 --- a/core/services/quantization/service.go +++ b/core/services/quantization/service.go @@ -40,19 +40,24 @@ type QuantizationService struct { mu sync.Mutex // jobs is the cross-replica job store: an in-memory map kept consistent across - // replicas via NATS, optionally read-through to PostgreSQL in distributed mode. + // replicas over the deployment's fan-out carrier, optionally read-through to + // PostgreSQL in distributed mode. jobs *syncstate.SyncedMap[string, *schema.QuantizationJob] } // NewQuantizationService creates a new QuantizationService. In distributed mode -// pass the shared NATS client and PostgreSQL store so jobs stay consistent across -// replicas; pass nil for both in standalone mode, where the disk Loader hydrates -// the map and there is nothing to broadcast. +// pass the deployment's broadcast carrier and PostgreSQL store so jobs stay +// consistent across replicas; pass nil for both in standalone mode, where the +// disk Loader hydrates the map and there is nothing to broadcast. +// +// bus is messaging.Broadcaster and not the NATS client: this state.*.delta +// family travels on whatever the deployment's fan-out carrier is, and in +// distributed mode that is PostgreSQL LISTEN/NOTIFY. func NewQuantizationService( appConfig *config.ApplicationConfig, modelLoader *model.ModelLoader, configLoader *config.ModelConfigLoader, - nats messaging.MessagingClient, + bus messaging.Broadcaster, store *distributed.QuantStore, ) *QuantizationService { s := &QuantizationService{ @@ -72,7 +77,7 @@ func NewQuantizationService( s.jobs = syncstate.New(syncstate.Config[string, *schema.QuantizationJob]{ Name: "quant.jobs", Key: func(j *schema.QuantizationJob) string { return j.ID }, - Nats: nats, + Bus: bus, Store: syncStore, Loader: s.loadJobsFromDisk, // ignored when Store is set (distributed mode) }) diff --git a/core/services/quantization/service_test.go b/core/services/quantization/service_test.go index 665728614..3c24989e2 100644 --- a/core/services/quantization/service_test.go +++ b/core/services/quantization/service_test.go @@ -15,6 +15,7 @@ import ( "github.com/mudler/LocalAI/core/config" "github.com/mudler/LocalAI/core/schema" "github.com/mudler/LocalAI/core/services/distributed" + "github.com/mudler/LocalAI/core/services/messaging" "github.com/mudler/LocalAI/core/services/testutil" ) @@ -22,7 +23,7 @@ import ( // The model/config loaders are nil because the read/sync paths under test never // touch them; the data dir is a throwaway temp dir so the disk Loader finds // nothing. -func newTestService(bus *testutil.FakeBus) *QuantizationService { +func newTestService(bus messaging.Broadcaster) *QuantizationService { appConfig := &config.ApplicationConfig{ Context: context.Background(), DataPath: GinkgoT().TempDir(), @@ -79,6 +80,29 @@ var _ = Describe("QuantizationService", func() { Expect(err).To(HaveOccurred(), "a delete on A must remove the job from B") }) + It("publishes on the carrier it was handed, and on the quantization family's own subject", func() { + // S2 in the wiring table. The negative half is the one that matters: + // finetune and quant land on the SAME LISTEN channel in distributed + // mode, so the subject is all that separates them. + job := &schema.QuantizationJob{ID: "job-subject", UserID: "user-1", Status: "queued", CreatedAt: "2026-06-27T10:00:00Z"} + Expect(a.jobs.Set(ctx, job)).To(Succeed()) + + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("quant.jobs"))).To(Equal(1)) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("finetune.jobs"))).To(Equal(0)) + }) + + It("broadcasts nothing at all when it is built with no carrier", func() { + solo := newTestService(nil) + DeferCleanup(func() { Expect(solo.Close()).To(Succeed()) }) + + Expect(solo.jobs.Set(ctx, &schema.QuantizationJob{ID: "solo", UserID: "u", CreatedAt: "2026-06-27T10:00:00Z"})).To(Succeed()) + Expect(bus.PublishCount(messaging.SubjectSyncStateDelta("quant.jobs"))).To(Equal(0)) + + got, err := solo.GetJob("u", "solo") + Expect(err).ToNot(HaveOccurred(), "a carrier-less service must still serve its own reads") + Expect(got.ID).To(Equal("solo")) + }) + It("propagates a status update from A to B", func() { job := &schema.QuantizationJob{ID: "job-3", UserID: "user-1", Status: "quantizing", CreatedAt: "2026-06-27T10:00:00Z"} Expect(a.jobs.Set(ctx, job)).To(Succeed()) diff --git a/core/services/syncstate/syncstate.go b/core/services/syncstate/syncstate.go index 8e70138d2..7373c2f05 100644 --- a/core/services/syncstate/syncstate.go +++ b/core/services/syncstate/syncstate.go @@ -5,9 +5,18 @@ // that is surfaced to the HTTP/UI API; without cross-replica sync a poll that // lands on a replica which did not originate a change sees stale or missing data. // SyncedMap collapses the three legs each feature otherwise hand-wires - an -// in-memory map, a NATS broadcast/apply path, and optional durable read-through - -// into one well-tested component so cross-replica consistency is a configuration -// choice rather than a bespoke re-implementation. +// in-memory map, a broadcast/apply path over the deployment's fan-out carrier, +// and optional durable read-through - into one well-tested component so +// cross-replica consistency is a configuration choice rather than a bespoke +// re-implementation. +// +// The carrier is messaging.Broadcaster and nothing narrower, so a deployment +// can carry these deltas on PostgreSQL LISTEN/NOTIFY. That is not a detail of +// the transport: LISTEN/NOTIFY is at most once to CONNECTED listeners and never +// replays, so a map whose only convergence path were the deltas would answer +// from state it can never repair. Store (or Loader) is what hydrate, the +// reconnect callback and the reconcile ticker read, and it is the reason a +// dropped delta is a gap that closes rather than a value that was never set. package syncstate import ( @@ -36,9 +45,20 @@ type Store[K comparable, V any] interface { // Config configures a SyncedMap. type Config[K comparable, V any] struct { - Name string // subject namespace, e.g. "finetune.jobs" - Key func(V) K // extract the key from a value - Nats messaging.MessagingClient // nil => standalone: in-memory only, no broadcast/subscribe + Name string // subject namespace, e.g. "finetune.jobs" + Key func(V) K // extract the key from a value + + // Bus is the fan-out carrier. nil => standalone: in-memory only, no + // broadcast and no subscribe. + // + // It is messaging.Broadcaster rather than messaging.MessagingClient because + // this component only ever publishes and subscribes: request/reply and + // queue groups are not part of what a replicated map needs, and demanding + // them would rule out every carrier that does not have them. The field is + // not called Nats because a field of that name holding a PostgreSQL carrier + // is a comment that claims more than the code does. + Bus messaging.Broadcaster + Store Store[K, V] // optional read-through persistence Loader func(ctx context.Context) ([]V, error) // source when there is no Store (e.g. disk reload) OnApply func(op string, k K, v V) // optional hook after an applied change (e.g. ShutdownModel) @@ -144,21 +164,24 @@ func (m *SyncedMap[K, V]) Start(ctx context.Context) error { // goroutines, so it cannot be cancelled or deferred within this scope. m.lifeCtx, m.cancel = context.WithCancel(context.Background()) // #nosec G118 -- cancel is invoked in Close() - if m.cfg.Nats != nil { + if m.cfg.Bus != nil { for _, filter := range m.subscribeFilters() { - sub, err := messaging.SubscribeJSON(m.cfg.Nats, filter, m.apply) + sub, err := messaging.SubscribeJSON(m.cfg.Bus, filter, m.apply) if err != nil { return err } m.subs = append(m.subs, sub) } - // nats.go transparently resubscribes on reconnect, but it cannot know we - // kept derived in-memory state that may have drifted while the link was - // down, so re-hydrate from the durable source. Detected via an optional - // interface so MessagingClient itself stays minimal; standalone/test - // clients without the method simply fall back to the reconcile ticker. - if r, ok := m.cfg.Nats.(interface{ OnReconnect(func()) }); ok { + // A carrier that reconnects restores its own registrations, but it + // cannot know we kept derived in-memory state that drifted while the + // link was down: every delta published in that window was delivered to + // the replicas that were connected and to nobody else, and neither + // carrier replays. Re-hydrating from the durable source is what turns + // that gap into a delay instead of a permanently wrong map. Detected + // via an optional interface so Broadcaster itself stays minimal; + // carriers without the method fall back to the reconcile ticker. + if r, ok := m.cfg.Bus.(interface{ OnReconnect(func()) }); ok { r.OnReconnect(func() { if err := m.hydrate(m.lifeCtx); err != nil { xlog.Warn("syncstate: reconnect re-hydrate failed", "name", m.cfg.Name, "error", err) @@ -261,12 +284,12 @@ func (m *SyncedMap[K, V]) Snapshot() map[K]V { return out } -// publish broadcasts a delta. Standalone (nil Nats) is a strict no-op. +// publish broadcasts a delta. Standalone (nil Bus) is a strict no-op. func (m *SyncedMap[K, V]) publish(op string, k K, v V) { - if m.cfg.Nats == nil { + if m.cfg.Bus == nil { return } - if err := m.cfg.Nats.Publish(m.publishSubject(), delta[K, V]{Op: op, Key: k, Value: v}); err != nil { + if err := m.cfg.Bus.Publish(m.publishSubject(), delta[K, V]{Op: op, Key: k, Value: v}); err != nil { xlog.Warn("syncstate: failed to broadcast delta", "name", m.cfg.Name, "op", op, "error", err) } } diff --git a/core/services/syncstate/syncstate_pg_test.go b/core/services/syncstate/syncstate_pg_test.go new file mode 100644 index 000000000..4e23b6286 --- /dev/null +++ b/core/services/syncstate/syncstate_pg_test.go @@ -0,0 +1,293 @@ +// SPDX-License-Identifier: MIT + +package syncstate_test + +import ( + "context" + "strings" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "gorm.io/gorm" + + "github.com/mudler/LocalAI/core/services/distributed" + "github.com/mudler/LocalAI/core/services/pgbus" + "github.com/mudler/LocalAI/core/services/syncstate" + "github.com/mudler/LocalAI/core/services/testutil" +) + +// ftAdapter is the same bridge the fine-tune service builds, spelled here so +// this file exercises the component over the REAL durable source rather than +// over a map that cannot fail the way a table fails. +type ftAdapter struct{ s *distributed.FineTuneStore } + +func (a ftAdapter) List(_ context.Context) ([]*distributed.FineTuneJobRecord, error) { + recs, err := a.s.ListAll() + if err != nil { + return nil, err + } + out := make([]*distributed.FineTuneJobRecord, len(recs)) + for i := range recs { + r := recs[i] + out[i] = &r + } + return out, nil +} + +func (a ftAdapter) Upsert(_ context.Context, r *distributed.FineTuneJobRecord) error { + return a.s.Upsert(r) +} + +func (a ftAdapter) Delete(_ context.Context, k string) error { return a.s.Delete(k) } + +// This file is the carriage proof, and it is on a real PostgreSQL for a reason +// the doubles cannot cover. +// +// A FakeBus delivers synchronously, in process, with no payload limit and no +// connection to lose. The three things this component's correctness actually +// rests on in distributed mode are exactly the three it cannot express: a +// notification bigger than PostgreSQL's 8000-byte payload cap has to spill to a +// row and come back byte identical; two maps in different families share ONE +// LISTEN channel, so separation is a filter decision and not a channel +// decision; and a listener whose session is terminated has to come back and +// re-hydrate, because every delta published while it was gone reached the +// replicas that were connected and nobody else. +// +// That last one is the invariant. A delta that was dropped or never delivered +// must not be able to read as a state that was never set: the notification says +// only that something changed, and the table is what says what it changed to. +var _ = Describe("SyncedMap on the PostgreSQL broadcast carrier", func() { + var ( + db *gorm.DB + dsn string + ) + + BeforeEach(func() { + db, dsn = testutil.SetupTestDBWithDSN() + Expect(pgbus.Migrate(context.Background(), db)).To(Succeed()) + }) + + // newBus is one replica's end of the carrier: its own pinned LISTEN + // connection, so a spec can drop one replica's session without touching the + // peer it is asserting against. + newBus := func() *pgbus.Bus { + GinkgoHelper() + b, err := pgbus.New(context.Background(), pgbus.Config{DSN: dsn, DB: db}) + Expect(err).ToNot(HaveOccurred()) + DeferCleanup(b.Close) + return b + } + + // terminate drops exactly this carrier's LISTEN session, matched on the + // application name it reports in pg_stat_activity, so the peer bus in the + // same spec stays up. + terminate := func(b *pgbus.Bus) { + GinkgoHelper() + var killed int64 + Expect(db.Raw( + "SELECT count(pg_terminate_backend(pid)) FROM pg_stat_activity WHERE application_name = ?", + b.ApplicationName(), + ).Scan(&killed).Error).To(Succeed()) + Expect(killed).To(BeNumerically("==", 1), "expected exactly one listener session to drop") + } + + Describe("two replicas over one shared store", func() { + var ( + ftStore *distributed.FineTuneStore + busA *pgbus.Bus + busB *pgbus.Bus + a, b *syncstate.SyncedMap[string, *distributed.FineTuneJobRecord] + ) + + newMap := func(bus *pgbus.Bus) *syncstate.SyncedMap[string, *distributed.FineTuneJobRecord] { + GinkgoHelper() + m := syncstate.New(syncstate.Config[string, *distributed.FineTuneJobRecord]{ + Name: "finetune.jobs", + Key: func(r *distributed.FineTuneJobRecord) string { return r.ID }, + // The carrier goes in as messaging.Broadcaster. Before this + // change the field was typed as the NATS client, so this line + // did not compile and the durable re-hydrate leg below had no + // consumer that could reach it. + Bus: bus, + Store: ftAdapter{s: ftStore}, + }) + Expect(m.Start(context.Background())).To(Succeed()) + DeferCleanup(func() { Expect(m.Close()).To(Succeed()) }) + return m + } + + rec := func(id, status, message string) *distributed.FineTuneJobRecord { + return &distributed.FineTuneJobRecord{ + ID: id, UserID: "u1", Model: "m", Backend: "bk", + TrainingType: "lora", TrainingMethod: "sft", Status: status, Message: message, + } + } + + BeforeEach(func() { + var err error + ftStore, err = distributed.NewFineTuneStore(db) + Expect(err).ToNot(HaveOccurred()) + + busA, busB = newBus(), newBus() + a, b = newMap(busA), newMap(busB) + }) + + It("carries a Set from one replica to the other", func() { + Expect(a.Set(context.Background(), rec("job-1", "queued", ""))).To(Succeed()) + + Eventually(func() string { + if r, ok := b.Get("job-1"); ok { + return r.Status + } + return "" + }, 30*time.Second, 50*time.Millisecond).Should(Equal("queued")) + }) + + It("carries a Delete from one replica to the other", func() { + Expect(a.Set(context.Background(), rec("job-2", "queued", ""))).To(Succeed()) + Eventually(func() bool { _, ok := b.Get("job-2"); return ok }, 30*time.Second, 50*time.Millisecond). + Should(BeTrue()) + + Expect(a.Delete(context.Background(), "job-2")).To(Succeed()) + Eventually(func() bool { _, ok := b.Get("job-2"); return ok }, 30*time.Second, 50*time.Millisecond). + Should(BeFalse(), "a delete that does not carry leaves a job the peer can never clear") + }) + + It("carries a delta too large for one notification, byte identical", func() { + // Comfortably over the carrier's 8000-byte notification cap, so + // this delta can only arrive through the spill row. A truncated or + // re-encoded payload here would surface as a job whose message the + // peer renders differently from the replica that owns it, with + // nothing failing. + big := strings.Repeat("training log line; ", 1200) + Expect(len(big)).To(BeNumerically(">", 8000)) + + Expect(a.Set(context.Background(), rec("job-3", "training", big))).To(Succeed()) + + Eventually(func() int { + if r, ok := b.Get("job-3"); ok { + return len(r.Message) + } + return 0 + }, 30*time.Second, 50*time.Millisecond).Should(Equal(len(big))) + + got, ok := b.Get("job-3") + Expect(ok).To(BeTrue()) + Expect(got.Message).To(Equal(big), "the spilled payload must come back exactly as it was published") + }) + + It("re-hydrates what it missed once its dropped listener reconnects", func() { + // The whole reason the durable store exists. The row is written + // straight to the shared table and never broadcast, which is what a + // peer's delta amounts to for a replica whose session was gone when + // it was published: at most once, no replay, no gap signal. If the + // reconnect callback is not registered, or the map converges on + // deltas alone, this replica answers as though the job had never + // been created - for good. + // + // The row is written BEFORE the session is dropped, and never + // broadcast, so this spec cannot pass by racing a reconnect against + // a NOTIFY: there is no NOTIFY, and the only path into b is a + // re-hydrate that the reconnect triggers. + Expect(ftStore.Upsert(rec("job-4", "completed", ""))).To(Succeed()) + _, present := b.Get("job-4") + Expect(present).To(BeFalse(), "a row written straight to the table reaches no map until something re-reads it") + + terminate(busB) + Eventually(busB.IsConnected, 60*time.Second, 100*time.Millisecond).Should(BeTrue()) + + Eventually(func() string { + if r, ok := b.Get("job-4"); ok { + return r.Status + } + return "" + }, 60*time.Second, 100*time.Millisecond). + Should(Equal("completed"), "a reconnected replica must re-read the durable source, not wait for a delta that is gone") + }) + + It("keeps carrying deltas after the reconnect", func() { + // The half the re-hydrate does not cover: a carrier that came back + // deaf re-hydrates once and then looks exactly like a deployment + // where nobody is publishing. + terminate(busB) + Eventually(busB.IsConnected, 60*time.Second, 100*time.Millisecond).Should(BeTrue()) + + Eventually(func() bool { + // Retried because a NOTIFY issued between the drop and the + // re-LISTEN is genuinely lost: this carrier is at most once. + Expect(a.Set(context.Background(), rec("job-5", "queued", ""))).To(Succeed()) + _, ok := b.Get("job-5") + return ok + }, 60*time.Second, 200*time.Millisecond).Should(BeTrue()) + }) + }) + + Describe("two families sharing one LISTEN channel", func() { + // Every state.* subject maps onto the single localai_state channel, so + // a subscriber hears its peers' families too and separation is decided + // by the filter alone. A carrier that delivered to every subscriber on + // the channel would put quantization jobs into the fine-tune map, and + // the fake bus cannot express the hazard because it has no channels. + newNamed := func(bus *pgbus.Bus, name string) *syncstate.SyncedMap[string, *job] { + GinkgoHelper() + m := syncstate.New(syncstate.Config[string, *job]{Name: name, Key: jobKey, Bus: bus}) + Expect(m.Start(context.Background())).To(Succeed()) + DeferCleanup(func() { Expect(m.Close()).To(Succeed()) }) + return m + } + + It("delivers a quantization delta to the quantization map and not to the fine-tune map", func() { + busA, busB := newBus(), newBus() + + quantA := newNamed(busA, "quant.jobs") + quantB := newNamed(busB, "quant.jobs") + ftB := newNamed(busB, "finetune.jobs") + + Expect(quantA.Set(context.Background(), &job{ID: "q-1", Status: "running"})).To(Succeed()) + + // The positive leg first, so the negative one below is asserted + // AFTER delivery has demonstrably happened rather than before it + // could have. + Eventually(func() bool { _, ok := quantB.Get("q-1"); return ok }, 30*time.Second, 50*time.Millisecond). + Should(BeTrue()) + + Expect(ftB.List()).To(BeEmpty(), "a fine-tune map must not apply a quantization delta off the shared channel") + }) + }) + + Describe("per-tenant subjects on the real carrier", func() { + // Task 10 pinned this against a fake bus. Repeated here because this is + // the commit where those subjects actually become LISTEN traffic, and + // because both the tenant subject and the cluster-wide one land on the + // same channel: if tenancy were a channel decision it would be right by + // accident, and it is not. + newTenant := func(bus *pgbus.Bus, tenant string) *syncstate.SyncedMap[string, *job] { + GinkgoHelper() + m := syncstate.New(syncstate.Config[string, *job]{ + Name: "agent.tasks", Key: jobKey, Bus: bus, PerTenant: true, Tenant: tenant, + }) + Expect(m.Start(context.Background())).To(Succeed()) + DeferCleanup(func() { Expect(m.Close()).To(Succeed()) }) + return m + } + + It("reaches the same tenant and the cluster-wide view, and no other tenant", func() { + busA, busB := newBus(), newBus() + + u1A := newTenant(busA, "u1") + u1B := newTenant(busB, "u1") + u2B := newTenant(busB, "u2") + clusterB := newTenant(busB, "") + + Expect(u1A.Set(context.Background(), &job{ID: "t-1", Status: "running"})).To(Succeed()) + + Eventually(func() bool { _, ok := u1B.Get("t-1"); return ok }, 30*time.Second, 50*time.Millisecond). + Should(BeTrue(), "the same tenant on another replica must see it") + Eventually(func() bool { _, ok := clusterB.Get("t-1"); return ok }, 30*time.Second, 50*time.Millisecond). + Should(BeTrue(), "the cluster-wide view hydrates across tenants, so it must apply across tenants") + + Expect(u2B.List()).To(BeEmpty(), "another tenant must never see this tenant's task") + }) + }) +}) diff --git a/core/services/syncstate/syncstate_test.go b/core/services/syncstate/syncstate_test.go index 47116a1b4..d5339135c 100644 --- a/core/services/syncstate/syncstate_test.go +++ b/core/services/syncstate/syncstate_test.go @@ -23,6 +23,12 @@ func jobKey(j *job) string { return j.ID } const stateName = "test.jobs" +// The retype this whole change rests on. Config.Bus is messaging.Broadcaster, +// so anything a spec or a deployment hands it must satisfy that and nothing +// wider; a fake that quietly needed MessagingClient would mean the component +// still could not be handed the PostgreSQL carrier. +var _ messaging.Broadcaster = (*testutil.FakeBus)(nil) + func deltaSubject() string { return messaging.SubjectSyncStateDelta(stateName) } // fakeStore is an in-memory Store that records call counts so specs can assert @@ -96,8 +102,8 @@ var _ = Describe("SyncedMap", func() { BeforeEach(func() { bus = testutil.NewFakeBus() - a = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) - b = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) + a = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) + b = syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) Expect(a.Start(ctx)).To(Succeed()) Expect(b.Start(ctx)).To(Succeed()) }) @@ -157,8 +163,8 @@ var _ = Describe("SyncedMap", func() { Describe("echo-loop guard", func() { It("applies its own broadcast once and does not re-publish", func() { bus := testutil.NewFakeBus() - a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) - b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) + a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) + b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) Expect(a.Start(ctx)).To(Succeed()) Expect(b.Start(ctx)).To(Succeed()) defer func() { @@ -183,8 +189,8 @@ var _ = Describe("SyncedMap", func() { bus := testutil.NewFakeBus() storeA := newFakeStore() storeB := newFakeStore() - a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: storeA}) - b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: storeB}) + a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus, Store: storeA}) + b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus, Store: storeB}) Expect(a.Start(ctx)).To(Succeed()) Expect(b.Start(ctx)).To(Succeed()) defer func() { @@ -215,9 +221,9 @@ var _ = Describe("SyncedMap", func() { ops []string keys []string ) - a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) + a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) b := syncstate.New(syncstate.Config[string, *job]{ - Name: stateName, Key: jobKey, Nats: bus, + Name: stateName, Key: jobKey, Bus: bus, OnApply: func(op string, k string, _ *job) { mu.Lock() ops = append(ops, op) @@ -242,7 +248,32 @@ var _ = Describe("SyncedMap", func() { }) }) - Describe("standalone (nil Nats)", func() { + Describe("standalone (nil Bus)", func() { + It("registers nothing and broadcasts nothing, and still serves reads from hydrate", func() { + // The bus exists in this spec and is deliberately NOT handed over. + // Zero publishes and zero subscribers is the assertion: "nothing was + // delivered" cannot tell a strict no-op apart from a subscription + // nobody happened to publish to, and a component that reached for a + // carrier of its own would pass that weaker check. + bus := testutil.NewFakeBus() + store := newFakeStore(&job{ID: "seeded", Status: "completed"}) + + m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Store: store}) + Expect(m.Start(ctx)).To(Succeed()) + defer func() { Expect(m.Close()).To(Succeed()) }() + + Expect(bus.Subscribers()).To(Equal(0), "a standalone map must register no subscription anywhere") + + got, ok := m.Get("seeded") + Expect(ok).To(BeTrue(), "hydrate must serve reads with no carrier at all") + Expect(got.Status).To(Equal("completed")) + Expect(m.List()).To(HaveLen(1)) + + Expect(m.Set(ctx, &job{ID: "local", Status: "running"})).To(Succeed()) + Expect(m.Delete(ctx, "seeded")).To(Succeed()) + Expect(bus.PublishCount(deltaSubject())).To(Equal(0), "a standalone map must not broadcast") + }) + It("works in-memory with no panic and nothing to broadcast", func() { m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey}) Expect(m.Start(ctx)).To(Succeed()) @@ -268,7 +299,7 @@ var _ = Describe("SyncedMap", func() { It("re-reads the source when the messaging client reconnects", func() { bus := testutil.NewFakeBus() store := newFakeStore(&job{ID: "init", Status: "running"}) - m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus, Store: store}) + m := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus, Store: store}) Expect(m.Start(ctx)).To(Succeed()) defer func() { Expect(m.Close()).To(Succeed()) }() @@ -305,7 +336,7 @@ var _ = Describe("SyncedMap per-tenant subjects", func() { m := syncstate.New(syncstate.Config[string, *job]{ Name: stateName, Key: jobKey, - Nats: bus, + Bus: bus, PerTenant: true, Tenant: tenant, }) @@ -437,8 +468,8 @@ var _ = Describe("SyncedMap per-tenant subjects", func() { // finetune.jobs, quantization and the responses store are unscoped // adopters. Pin that this change moved none of them. bus := testutil.NewFakeBus() - a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) - b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Nats: bus}) + a := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) + b := syncstate.New(syncstate.Config[string, *job]{Name: stateName, Key: jobKey, Bus: bus}) Expect(a.Start(ctx)).To(Succeed()) Expect(b.Start(ctx)).To(Succeed()) tenant := newTenantMap(bus, "u1") diff --git a/core/services/testutil/fakebus.go b/core/services/testutil/fakebus.go index d178fc3ee..1607a5e9c 100644 --- a/core/services/testutil/fakebus.go +++ b/core/services/testutil/fakebus.go @@ -43,6 +43,15 @@ type fakeBusSub struct { handler func([]byte) } +// The two interfaces this double stands in for. Broadcaster is asserted here +// and not left to the first adopter: every cross-replica map now takes a +// messaging.Broadcaster, and a fake that drifted out of that interface would +// break each adopter's suite in turn rather than the package that owns it. +var ( + _ messaging.MessagingClient = (*FakeBus)(nil) + _ messaging.Broadcaster = (*FakeBus)(nil) +) + // NewFakeBus returns a ready-to-use in-memory bus. func NewFakeBus() *FakeBus { return &FakeBus{publishCounts: map[string]int{}} @@ -74,6 +83,17 @@ func (b *FakeBus) PublishCount(subject string) int { return b.publishCounts[subject] } +// Subscribers reports how many live subscriptions this bus is carrying. +// +// It exists so a spec can assert the NEGATIVE: a component configured +// standalone must register nothing at all, and "nothing was delivered" cannot +// tell that apart from "a subscription exists and nobody published". +func (b *FakeBus) Subscribers() int { + b.mu.Lock() + defer b.mu.Unlock() + return len(b.subs) +} + type fakeBusSubscription struct { bus *FakeBus subRef fakeBusSub diff --git a/docs/content/features/distributed-mode.md b/docs/content/features/distributed-mode.md index b6c880b3f..73da03a47 100644 --- a/docs/content/features/distributed-mode.md +++ b/docs/content/features/distributed-mode.md @@ -38,7 +38,7 @@ Each model gets its own gRPC backend process, so a single worker can serve multi ## Prerequisites - **PostgreSQL** (with pgvector extension recommended for RAG) - used for node registry, job store, auth, and shared state - - Each frontend replica holds **one extra PostgreSQL session** beyond its connection pool, pinned for the life of the process, and creates a `bus_messages` table. Both belong to the broadcast carrier that is replacing NATS for cross-replica fan-out. Size `max_connections` for one additional session per frontend replica. + - Each frontend replica holds **one extra PostgreSQL session** beyond its connection pool, pinned for the life of the process, and creates a `bus_messages` table. Both belong to the broadcast carrier that is replacing NATS for cross-replica fan-out; it already carries the four `state.*.delta` families (see [Cross-replica in-memory state](#cross-replica-in-memory-state)). Size `max_connections` for one additional session per frontend replica. - That session reports an `application_name` of `localai_pgbus_`, so `SELECT count(*) FROM pg_stat_activity WHERE application_name LIKE 'localai_pgbus_%'` counts the replicas currently listening. If the carrier loses its session it redials and re-registers on its own; a broadcast published while it was down is not replayed, which is why nothing that must survive a gap is carried by a broadcast alone. - `bus_messages` holds only broadcasts too large for a PostgreSQL notification, and every replica retires rows older than ten minutes. The table is a spill buffer, not a log: it is not a place to read past events from. - **NATS** server - used for agent-worker coordination and the frontend's own cross-replica events. **Serve-backend workers do not connect to it at all**: every verb they take, and file staging with it, is an HTTP route on the worker's tunnel. Set no `LOCALAI_NATS_URL` on a `local-ai worker`. The frontend and any `local-ai agent-worker` still need one. @@ -121,6 +121,25 @@ The peer link is served at `/api/cluster/peer` and authenticates with `LOCALAI_R **The peer link has no per-replica credential yet.** It checks the shared registration token and takes the replica id in `?id=` on trust. Anything already holding that token - every worker holds it - can therefore open a peer link, relay through it to every worker tunnel a replica owns, by declaring another replica's id displace that replica's inbound link, and hold sessions open against the per-session receive window, which the peer-link code sizes at roughly 31 GiB of unread data per session and which on this route is also a memory budget an attacker can point at one replica. Treat `LOCALAI_REGISTRATION_TOKEN` as a cluster-wide secret with the blast radius of the whole fleet: give it its own value per deployment, do not reuse it elsewhere, and keep `/api/cluster/peer` on a network only your replicas and workers can reach. Per-replica credentials for this route are planned. {{% /notice %}} +### Cross-replica in-memory state + +Several features keep state in a frontend's process memory and surface it over the API: fine-tune jobs, quantization jobs, agent tasks and Open Responses metadata. A round-robin load balancer sends a follow-up request to any replica, so each of those maps is kept current on every replica by a broadcast. + +**Those four families travel on PostgreSQL, not on NATS.** Each mutation is a `NOTIFY` on the database the deployment already runs, and each replica holds one `LISTEN` session for it. There is nothing to configure: the carrier uses the same database URL as `--auth-database-url` / `LOCALAI_AUTH_DATABASE_URL`. + +| Map | Subject | +|-----|---------| +| Fine-tune jobs | `state.finetune-jobs.delta` | +| Quantization jobs | `state.quant-jobs.delta` | +| Agent tasks | `state.agent-tasks.delta` and `state.agent-tasks..delta` | +| Open Responses metadata | `state.responses-metadata.delta` | + +Two carrier details are visible to an operator. + +**The 8000-byte notification cap.** PostgreSQL refuses a `pg_notify` payload of 8000 bytes or more, and that limit is measured against the whole encoded notification, not just the value being replicated. A broadcast that does not fit is written to the `bus_messages` table and the notification carries the row id instead; the receiving replica reads the row and delivers the original bytes. This is an ordinary path and not an error: a fine-tune job carrying a long training message spills every time. Rows are retired ten minutes after they are written, by every replica, so `bus_messages` is a spill buffer and never a log of past events. + +**A broadcast is at most once, and is never replayed.** A `NOTIFY` reaches the sessions that are listening when it is issued and nobody else. A replica whose session was down in that window never receives the change, and no error is reported anywhere. That is why every one of these maps is backed by a durable table: the broadcast says only that something changed, and the table says what it changed to. A replica re-reads its table when its listener reconnects, so a missed delta is a delay and never a value that reads as though it had never been set. + ### Open Responses across replicas A response created by `POST /v1/responses` is held by the replica that served the request. A round-robin load balancer sends the follow-up poll, the `previous_response_id` chain and the cancel to any replica, so that metadata is replicated to every frontend and is also written to a `response_metadata` table in PostgreSQL. @@ -135,20 +154,26 @@ What crosses replicas and what does not: | Cancellation | The request is, the `CancelFunc` is not | The cancel is forwarded to the owning replica, which holds the function that stops generation | | Stream resume buffer (`starting_after`) | No | It is the full token log; replicating it would put every generated token on the bus. A resume request that lands on the wrong replica is refused with an explicit error, never with a silently truncated event list | -**Retention.** Rows carry the expiry of the response they describe, and each replica sweeps expired rows every five minutes. The expiry comes from the Open Responses store TTL, which is **`0` (no expiration) by default**: +**Retention.** Each replica sweeps dead rows out of `response_metadata` every five minutes, and a row is dead when either of two things is true. -```yaml -environment: - LOCALAI_OPEN_RESPONSES_STORE_TTL: "1h" -``` +- It carries the expiry of the response it describes, and that expiry has passed. The expiry comes from the Open Responses store TTL, which is `0` (no expiration) by default: -Leave it at `0` in distributed mode and nothing ever expires: `response_metadata` grows for the life of the deployment, and every replica that restarts re-hydrates every response the cluster has ever created. Set a TTL that matches how long clients are allowed to poll for a response. + ```yaml + environment: + LOCALAI_OPEN_RESPONSES_STORE_TTL: "1h" + ``` -These rows carry the request body and the generated output, not just identifiers. They live in the same database as the rest of the cluster state, and the TTL above is the only thing that removes them. +- It carries no expiry, because the TTL is `0`, and it is more than **24 hours** old. + +The 24-hour floor is the table's own bound and it is independent of the TTL. A TTL of `0` is a reasonable answer for the in-memory map it governs, which dies with the process; a table has no such bound, so without a floor `response_metadata` would grow for the life of the deployment and every restarting replica would re-hydrate every response the cluster had ever created. + +The floor never overrides a TTL you set. A row that names an expiry is judged on that expiry alone, longer or shorter than 24 hours. What the floor bounds is only how long a response stays resolvable **on a replica that did not create it**: the owning replica keeps it in memory for exactly as long as the TTL says. Set a TTL that matches how long clients are allowed to poll for a response. + +These rows carry the request body and the generated output, not just identifiers. They live in the same database as the rest of the cluster state. ### Agent tasks are scoped to their tenant -Every frontend replica keeps agent task definitions in memory so that `GET /api/agent/tasks` answers from any replica. That in-memory copy is kept current by a broadcast on the cluster bus, and the broadcast carries the owning user in the subject: +Every frontend replica keeps agent task definitions in memory so that `GET /api/agent/tasks` answers from any replica. That in-memory copy is kept current by a broadcast on the PostgreSQL carrier described above, and the broadcast carries the owning user in the subject: | Map | Subject it publishes on | Subjects it applies | |-----|------------------------|---------------------| diff --git a/tests/e2e/distributed/syncstate_distributed_test.go b/tests/e2e/distributed/syncstate_distributed_test.go index acd2797e6..b0af894e7 100644 --- a/tests/e2e/distributed/syncstate_distributed_test.go +++ b/tests/e2e/distributed/syncstate_distributed_test.go @@ -4,7 +4,7 @@ import ( "context" "github.com/mudler/LocalAI/core/services/distributed" - "github.com/mudler/LocalAI/core/services/messaging" + "github.com/mudler/LocalAI/core/services/pgbus" "github.com/mudler/LocalAI/core/services/syncstate" . "github.com/onsi/ginkgo/v2" @@ -42,48 +42,59 @@ func (a ftSyncStore) Delete(_ context.Context, k string) error { return a.s.Dele // This suite is the real-infrastructure counterpart to the fake-bus unit tests: // two SyncedMap instances stand in for two LocalAI frontend replicas, each with -// its OWN NATS connection to a shared NATS server and a SHARED PostgreSQL store - -// the exact distributed-mode invariant (single shared DB, per-replica process -// state). It proves the delta path works over the wire and that a late-joining -// replica recovers via store hydrate (the at-most-once gap a fake bus cannot -// exercise). -var _ = Describe("SyncedMap two-replica sync over real NATS", Label("Distributed"), func() { +// its OWN pinned LISTEN connection to the shared PostgreSQL and a SHARED store +// on the same database - the exact distributed-mode invariant (single shared +// DB, per-replica process state). It proves the delta path works over the wire +// and that a late-joining replica recovers via store hydrate (the at-most-once +// gap a fake bus cannot exercise). +// +// There is no NATS here any more. This family's deltas ride the deployment's +// PostgreSQL broadcast carrier, so a suite that still proved them over a broker +// would be proving something the product no longer does. +var _ = Describe("SyncedMap two-replica sync over the PostgreSQL carrier", Label("Distributed"), func() { var ( infra *TestInfra + db *gorm.DB ftStore *distributed.FineTuneStore ) BeforeEach(func() { infra = SetupInfra("localai_syncstate_dist_test") - db, err := gorm.Open(pgdriver.Open(infra.PGURL), &gorm.Config{ + var err error + db, err = gorm.Open(pgdriver.Open(infra.PGURL), &gorm.Config{ Logger: logger.Default.LogMode(logger.Silent), }) Expect(err).ToNot(HaveOccurred()) + Expect(pgbus.Migrate(infra.Ctx, db)).To(Succeed()) + ftStore, err = distributed.NewFineTuneStore(db) Expect(err).ToNot(HaveOccurred()) }) - // newReplica builds an independent "replica": its own NATS client to the - // shared server plus a SyncedMap over the shared store, started (hydrate + - // subscribe) and cleaned up automatically. + // newReplica builds an independent "replica": its own carrier, with its own + // pinned LISTEN connection, plus a SyncedMap over the shared store, started + // (hydrate + subscribe) and cleaned up automatically. + // + // No flush is needed where NATS needed one: pgbus.Subscribe returns only + // after the server has acknowledged the LISTEN, so a subscription that has + // been registered cannot miss a NOTIFY issued afterwards. newReplica := func() *syncstate.SyncedMap[string, *distributed.FineTuneJobRecord] { GinkgoHelper() - nc, err := messaging.New(infra.NatsURL) + bus, err := pgbus.New(infra.Ctx, pgbus.Config{DSN: infra.PGURL, DB: db}) Expect(err).ToNot(HaveOccurred()) sm := syncstate.New(syncstate.Config[string, *distributed.FineTuneJobRecord]{ Name: "finetune.jobs", Key: func(r *distributed.FineTuneJobRecord) string { return r.ID }, - Nats: nc, + Bus: bus, Store: ftSyncStore{s: ftStore}, }) Expect(sm.Start(infra.Ctx)).To(Succeed()) - FlushNATS(nc) // ensure the subscription is registered server-side before any publish DeferCleanup(func() { _ = sm.Close() - nc.Close() + bus.Close() }) return sm } @@ -102,7 +113,7 @@ var _ = Describe("SyncedMap two-replica sync over real NATS", Label("Distributed Expect(a.Set(infra.Ctx, rec("job-1", "queued"))).To(Succeed()) Eventually(func() bool { _, ok := b.Get("job-1"); return ok }, "10s", "50ms"). - Should(BeTrue(), "replica B must observe the job created on A via NATS") + Should(BeTrue(), "replica B must observe the job created on A over the carrier") got, ok := b.Get("job-1") Expect(ok).To(BeTrue())